Don't sum in-flight piece which does not intersect filtered ranges

When calculating completed length, don't sum the completed length of
piece which does not intersect filtered ranges.

6 files changed, 81 insertions, 7 deletions

diff --git a/src/BitfieldMan.cc b/src/BitfieldMan.cc
index feb387c1..dab736c8 100644
--- a/src/BitfieldMan.cc
+++ b/src/BitfieldMan.cc
@@ -609,6 +609,14 @@ bool BitfieldMan::isAllFilterBitSet() const
   return testAllBitSet(filterBitfield_, bitfieldLength_, blocks_);
 }
 
+bool BitfieldMan::isFilterBitSet(size_t index) const {
+  if(filterBitfield_) {
+    return bitfield::test(filterBitfield_, blocks_, index);
+  } else {
+    return false;
+  }
+}
+
 bool BitfieldMan::isBitSet(size_t index) const
 {
   return bitfield::test(bitfield_, blocks_, index);
diff --git a/src/BitfieldMan.h b/src/BitfieldMan.h
index 818e6957..f82557d4 100644
--- a/src/BitfieldMan.h
+++ b/src/BitfieldMan.h
@@ -200,6 +200,9 @@ public:
   bool isAllBitSet() const;
 
   bool isAllFilterBitSet() const;
+  // Returns true if index bit is set in filterBitfield_.  If
+  // filterBitfield_ is NULL, returns false.
+  bool isFilterBitSet(size_t index) const;
 
   const unsigned char* getBitfield() const
   {
diff --git a/src/DefaultPieceStorage.cc b/src/DefaultPieceStorage.cc
index f1f13f2f..929359d3 100644
--- a/src/DefaultPieceStorage.cc
+++ b/src/DefaultPieceStorage.cc
@@ -537,14 +537,29 @@ int64_t DefaultPieceStorage::getCompletedLength()
 int64_t DefaultPieceStorage::getFilteredCompletedLength()
 {
   return bitfieldMan_->getFilteredCompletedLength()+
-    getInFlightPieceCompletedLength();
+    getInFlightPieceFilteredCompletedLength();
 }
 
-size_t DefaultPieceStorage::getInFlightPieceCompletedLength() const
+int64_t DefaultPieceStorage::getInFlightPieceCompletedLength() const
 {
-  return std::accumulate(usedPieces_.begin(), usedPieces_.end(),
-                         0, adopt2nd(std::plus<size_t>(),
-                                     mem_fun_sh(&Piece::getCompletedLength)));
+  int64_t len = 0;
+  for(UsedPieceSet::const_iterator i = usedPieces_.begin(),
+        eoi = usedPieces_.end(); i != eoi; ++i) {
+    len += (*i)->getCompletedLength();
+  }
+  return len;
+}
+
+int64_t DefaultPieceStorage::getInFlightPieceFilteredCompletedLength() const
+{
+  int64_t len = 0;
+  for(UsedPieceSet::const_iterator i = usedPieces_.begin(),
+        eoi = usedPieces_.end(); i != eoi; ++i) {
+    if(bitfieldMan_->isFilterBitSet((*i)->getIndex())) {
+      len += (*i)->getCompletedLength();
+    }
+  }
+  return len;
 }
 
 // not unittested
diff --git a/src/DefaultPieceStorage.h b/src/DefaultPieceStorage.h
index 7851f171..ec3af3d4 100644
--- a/src/DefaultPieceStorage.h
+++ b/src/DefaultPieceStorage.h
@@ -110,8 +110,11 @@ private:
   void deleteUsedPiece(const SharedHandle<Piece>& piece);
   SharedHandle<Piece> findUsedPiece(size_t index) const;
 
-  size_t getInFlightPieceCompletedLength() const;
-
+  // Returns the sum of completed length of in-flight pieces
+  int64_t getInFlightPieceCompletedLength() const;
+  // Returns the sum of completed length of in-flight pieces
+  // intersecting filter ranges.
+  int64_t getInFlightPieceFilteredCompletedLength() const;
 public:
   // Setting randomPieceStatsOrdering to true means a piece is chosen in
   // random when more than 2 pieces has the same rarity.
diff --git a/test/BitfieldManTest.cc b/test/BitfieldManTest.cc
index eb03fad8..bc7ebcc2 100644
--- a/test/BitfieldManTest.cc
+++ b/test/BitfieldManTest.cc
@@ -18,6 +18,7 @@ class BitfieldManTest:public CppUnit::TestFixture {
   CPPUNIT_TEST(testGetFirstMissingIndex);
   CPPUNIT_TEST(testIsAllBitSet);
   CPPUNIT_TEST(testFilter);
+  CPPUNIT_TEST(testIsFilterBitSet);
   CPPUNIT_TEST(testAddFilter_zeroLength);
   CPPUNIT_TEST(testAddNotFilter);
   CPPUNIT_TEST(testAddNotFilter_zeroLength);
@@ -51,6 +52,7 @@ public:
   
   void testIsAllBitSet();
   void testFilter();
+  void testIsFilterBitSet();
   void testAddFilter_zeroLength();
   void testAddNotFilter();
   void testAddNotFilter_zeroLength();
@@ -244,6 +246,17 @@ void BitfieldManTest::testFilter()
   CPPUNIT_ASSERT_EQUAL((int64_t)31ULL, btman2.getFilteredTotalLength());
 }
 
+void BitfieldManTest::testIsFilterBitSet()
+{
+  BitfieldMan btman(2, 32);
+  CPPUNIT_ASSERT(!btman.isFilterBitSet(0));
+  btman.addFilter(0, 2);
+  CPPUNIT_ASSERT(btman.isFilterBitSet(0));
+  CPPUNIT_ASSERT(!btman.isFilterBitSet(1));
+  btman.addFilter(2, 4);
+  CPPUNIT_ASSERT(btman.isFilterBitSet(1));
+}
+
 void BitfieldManTest::testAddFilter_zeroLength()
 {
   BitfieldMan bits(1024, 1024*1024);
diff --git a/test/DefaultPieceStorageTest.cc b/test/DefaultPieceStorageTest.cc
index e12c707c..b2ef7d20 100644
--- a/test/DefaultPieceStorageTest.cc
+++ b/test/DefaultPieceStorageTest.cc
@@ -37,6 +37,7 @@ class DefaultPieceStorageTest:public CppUnit::TestFixture {
   CPPUNIT_TEST(testCancelPiece);
   CPPUNIT_TEST(testMarkPiecesDone);
   CPPUNIT_TEST(testGetCompletedLength);
+  CPPUNIT_TEST(testGetFilteredCompletedLength);
   CPPUNIT_TEST(testGetNextUsedIndex);
   CPPUNIT_TEST_SUITE_END();
 private:
@@ -71,6 +72,7 @@ public:
   void testCancelPiece();
   void testMarkPiecesDone();
   void testGetCompletedLength();
+  void testGetFilteredCompletedLength();
   void testGetNextUsedIndex();
 };
 
@@ -345,6 +347,36 @@ void DefaultPieceStorageTest::testGetCompletedLength()
   CPPUNIT_ASSERT_EQUAL((int64_t)256*1024*1024, ps.getCompletedLength());
 }
 
+void DefaultPieceStorageTest::testGetFilteredCompletedLength()
+{
+  const size_t pieceLength = 1024*1024;
+  SharedHandle<DownloadContext> dctx(new DownloadContext());
+  dctx->setPieceLength(pieceLength);
+  SharedHandle<FileEntry> files[] = {
+    SharedHandle<FileEntry>(new FileEntry("foo", 2*pieceLength, 0)),
+    SharedHandle<FileEntry>(new FileEntry("bar", 4*pieceLength, 2*pieceLength))
+  };
+  files[1]->setRequested(false);
+  dctx->setFileEntries(&files[0], &files[2]);
+
+  DefaultPieceStorage ps(dctx, option_.get());
+  std::vector<SharedHandle<Piece> > inflightPieces(2);
+  inflightPieces[0] = SharedHandle<Piece>(new Piece(1, pieceLength));
+  inflightPieces[0]->completeBlock(0);
+  inflightPieces[1] = SharedHandle<Piece>(new Piece(2, pieceLength));
+  inflightPieces[1]->completeBlock(1);
+  inflightPieces[1]->completeBlock(2);
+
+  ps.addInFlightPiece(inflightPieces);
+  ps.setupFileFilter();
+
+  SharedHandle<Piece> piece = ps.getMissingPiece(0, 1);
+  ps.completePiece(piece);
+
+  CPPUNIT_ASSERT_EQUAL((int64_t)pieceLength+16*1024,
+                       ps.getFilteredCompletedLength());
+}
+
 void DefaultPieceStorageTest::testGetNextUsedIndex()
 {
   DefaultPieceStorage pss(dctx_, option_.get());