Added PieceStorage::onDownloadIncomplete() virtual function.

In DefaultPieceStorage::onDownloadIncomplete(), we call
StreamPieceSelector::onBitfieldInit().
Added GeomStreamPieceSelectorTest.

13 files changed, 45 insertions, 4 deletions

diff --git a/src/BtCheckIntegrityEntry.cc b/src/BtCheckIntegrityEntry.cc
index a51f17ca..f7428bff 100644
--- a/src/BtCheckIntegrityEntry.cc
+++ b/src/BtCheckIntegrityEntry.cc
@@ -51,11 +51,12 @@ BtCheckIntegrityEntry::~BtCheckIntegrityEntry() {}
 void BtCheckIntegrityEntry::onDownloadIncomplete
 (std::vector<Command*>& commands, DownloadEngine* e)
 {
+  const SharedHandle<PieceStorage>& ps = getRequestGroup()->getPieceStorage();
+  ps->onDownloadIncomplete();
   if(getRequestGroup()->getOption()->getAsBool(PREF_HASH_CHECK_ONLY)) {
     return;
   }
-  const SharedHandle<DiskAdaptor>& diskAdaptor =
-    getRequestGroup()->getPieceStorage()->getDiskAdaptor();
+  const SharedHandle<DiskAdaptor>& diskAdaptor = ps->getDiskAdaptor();
   if(diskAdaptor->isReadOnlyEnabled()) {
     // Now reopen DiskAdaptor with read only disabled.
     diskAdaptor->closeFile();
diff --git a/src/DefaultPieceStorage.cc b/src/DefaultPieceStorage.cc
index 47a43480..7268bcf8 100644
--- a/src/DefaultPieceStorage.cc
+++ b/src/DefaultPieceStorage.cc
@@ -810,4 +810,9 @@ size_t DefaultPieceStorage::getNextUsedIndex(size_t index)
   return bitfieldMan_->countBlock();
 }
 
+void DefaultPieceStorage::onDownloadIncomplete()
+{
+  streamPieceSelector_->onBitfieldInit();
+}
+
 } // namespace aria2
diff --git a/src/DefaultPieceStorage.h b/src/DefaultPieceStorage.h
index 18d664de..30f0ecf9 100644
--- a/src/DefaultPieceStorage.h
+++ b/src/DefaultPieceStorage.h
@@ -272,6 +272,8 @@ public:
 
   virtual size_t getNextUsedIndex(size_t index);
 
+  virtual void onDownloadIncomplete();
+
   /**
    * This method is made private for test purpose only.
    */
diff --git a/src/DefaultStreamPieceSelector.cc b/src/DefaultStreamPieceSelector.cc
index 2a899a89..09fda560 100644
--- a/src/DefaultStreamPieceSelector.cc
+++ b/src/DefaultStreamPieceSelector.cc
@@ -54,4 +54,6 @@ bool DefaultStreamPieceSelector::select
     (index, minSplitSize, ignoreBitfield, length);
 }
 
+void DefaultStreamPieceSelector::onBitfieldInit() {}
+
 } // namespace aria2
diff --git a/src/DefaultStreamPieceSelector.h b/src/DefaultStreamPieceSelector.h
index 20f4f42c..33d72ac2 100644
--- a/src/DefaultStreamPieceSelector.h
+++ b/src/DefaultStreamPieceSelector.h
@@ -51,6 +51,8 @@ public:
    size_t minSplitSize,
    const unsigned char* ignoreBitfield,
    size_t length);
+
+  virtual void onBitfieldInit();
 private:
   BitfieldMan* bitfieldMan_;
 };
diff --git a/src/GeomStreamPieceSelector.cc b/src/GeomStreamPieceSelector.cc
index be1cd59c..e69b5ec7 100644
--- a/src/GeomStreamPieceSelector.cc
+++ b/src/GeomStreamPieceSelector.cc
@@ -41,7 +41,8 @@ GeomStreamPieceSelector::GeomStreamPieceSelector
 (BitfieldMan* bitfieldMan,
  double base)
   : bitfieldMan_(bitfieldMan),
-    base_(base)
+    base_(base),
+    offsetIndex_(0)
 {}
 
 GeomStreamPieceSelector::~GeomStreamPieceSelector() {}
@@ -53,7 +54,16 @@ bool GeomStreamPieceSelector::select
  size_t length)
 {
   return bitfieldMan_->getGeomMissingUnusedIndex
-    (index, minSplitSize, ignoreBitfield, length, base_, 0);
+    (index, minSplitSize, ignoreBitfield, length, base_, offsetIndex_);
+}
+
+void GeomStreamPieceSelector::onBitfieldInit()
+{
+  size_t index;
+  bool r = bitfieldMan_->getFirstMissingIndex(index);
+  if(r) {
+    offsetIndex_ = index;
+  }
 }
 
 } // namespace aria2
diff --git a/src/GeomStreamPieceSelector.h b/src/GeomStreamPieceSelector.h
index 760b0801..2c234742 100644
--- a/src/GeomStreamPieceSelector.h
+++ b/src/GeomStreamPieceSelector.h
@@ -51,9 +51,12 @@ public:
    size_t minSplitSize,
    const unsigned char* ignoreBitfield,
    size_t length);
+
+  virtual void onBitfieldInit();
 private:
   BitfieldMan* bitfieldMan_;
   double base_;
+  size_t offsetIndex_;
 };
 
 } // namespace aria2
diff --git a/src/InorderStreamPieceSelector.cc b/src/InorderStreamPieceSelector.cc
index edb6e94e..20171f73 100644
--- a/src/InorderStreamPieceSelector.cc
+++ b/src/InorderStreamPieceSelector.cc
@@ -54,4 +54,6 @@ bool InorderStreamPieceSelector::select
     (index, minSplitSize, ignoreBitfield, length);
 }
 
+void InorderStreamPieceSelector::onBitfieldInit() {}
+
 } // namespace aria2
diff --git a/src/InorderStreamPieceSelector.h b/src/InorderStreamPieceSelector.h
index a6d4b034..5d5481c8 100644
--- a/src/InorderStreamPieceSelector.h
+++ b/src/InorderStreamPieceSelector.h
@@ -51,6 +51,8 @@ public:
    size_t minSplitSize,
    const unsigned char* ignoreBitfield,
    size_t length);
+
+  virtual void onBitfieldInit();
 private:
   BitfieldMan* bitfieldMan_;
 };
diff --git a/src/PieceStorage.h b/src/PieceStorage.h
index e50495b7..b3d3e46c 100644
--- a/src/PieceStorage.h
+++ b/src/PieceStorage.h
@@ -283,6 +283,9 @@ public:
   // are not used and not completed. If all pieces after index+1 are
   // used or completed, returns the number of pieces.
   virtual size_t getNextUsedIndex(size_t index) = 0;
+
+  // Called when system detects download is not finished
+  virtual void onDownloadIncomplete() = 0;
 };
 
 typedef SharedHandle<PieceStorage> PieceStorageHandle;
diff --git a/src/StreamCheckIntegrityEntry.cc b/src/StreamCheckIntegrityEntry.cc
index 3de23736..aeaf5e8c 100644
--- a/src/StreamCheckIntegrityEntry.cc
+++ b/src/StreamCheckIntegrityEntry.cc
@@ -38,6 +38,7 @@
 #include "StreamFileAllocationEntry.h"
 #include "prefs.h"
 #include "Option.h"
+#include "PieceStorage.h"
 
 namespace aria2 {
 
@@ -51,6 +52,8 @@ StreamCheckIntegrityEntry::~StreamCheckIntegrityEntry() {}
 void StreamCheckIntegrityEntry::onDownloadIncomplete
 (std::vector<Command*>& commands, DownloadEngine* e)
 {
+  const SharedHandle<PieceStorage>& ps = getRequestGroup()->getPieceStorage();
+  ps->onDownloadIncomplete();
   if(getRequestGroup()->getOption()->getAsBool(PREF_HASH_CHECK_ONLY)) {
     return;
   }
diff --git a/src/StreamPieceSelector.h b/src/StreamPieceSelector.h
index a8f95513..55f9be96 100644
--- a/src/StreamPieceSelector.h
+++ b/src/StreamPieceSelector.h
@@ -60,6 +60,10 @@ public:
    size_t minSplitSize,
    const unsigned char* ignoreBitfield,
    size_t length) = 0;
+
+  // Called when initial bitfield was fixed. Optimize
+  // StreamPieceSelector to take advantages of the initial bitfield.
+  virtual void onBitfieldInit() = 0;
 };
 
 } // namespace aria2
diff --git a/src/UnknownLengthPieceStorage.h b/src/UnknownLengthPieceStorage.h
index d409b3e9..b6b340f6 100644
--- a/src/UnknownLengthPieceStorage.h
+++ b/src/UnknownLengthPieceStorage.h
@@ -292,6 +292,8 @@ public:
   virtual size_t getNextUsedIndex(size_t index) { return 0; }
 
   void setDiskWriterFactory(const SharedHandle<DiskWriterFactory>& diskWriterFactory);
+
+  virtual void onDownloadIncomplete() {}
 };
 
 typedef SharedHandle<UnknownLengthPieceStorage> UnknownLengthPieceStorageHandle;