2010-03-28 Tatsuhiro Tsujikawa <t-tujikawa@users.sourceforge.net>

	Added --always-resume and --max-resume-failure-tries option.  If
	--always-resume=false is given, when all given URIs does not
	support resume or aria2 encounters N URIs which does not support
	resume
        (N is the value specified using --max-resume-failure-tries
	option), aria2 download file from scratch.  The default behavior
	is --always-resume=true, which means if all URIs do not support
	resume, download fails. I think this is OK because user normally
	don't like to see that partially downloaded file is
	overwritten(this is particularly true if file size is big).  This
	option is useful when aria2 is used as download backend and
	graceful falling back to overwritten behavior is preferable.
	Added exit status value 8, which means download failed because
	server did not support resume.
	* src/AbstractCommand.cc
	* src/DefaultPieceStorage.cc
	* src/DownloadCommand.cc
	* src/DownloadResultCode.h
	* src/FileEntry.h
	* src/FtpNegotiationCommand.cc
	* src/HttpResponse.cc
	* src/HttpResponseCommand.cc
	* src/OptionHandlerFactory.cc
	* src/RequestGroup.cc
	* src/RequestGroup.h
	* src/SegmentMan.cc
	* src/SegmentMan.h
	* src/prefs.cc
	* src/prefs.h
	* src/usage_text.h
	* test/DefaultPieceStorageTest.cc
	* test/SegmentManTest.cc

ChangeLog

@@ -1,3 +1,38 @@
+2010-03-28  Tatsuhiro Tsujikawa  <t-tujikawa@users.sourceforge.net>
+
+	Added --always-resume and --max-resume-failure-tries option.  If
+	--always-resume=false is given, when all given URIs does not
+	support resume or aria2 encounters N URIs which does not support
+	resume
+        (N is the value specified using --max-resume-failure-tries
+	option), aria2 download file from scratch.  The default behavior
+	is --always-resume=true, which means if all URIs do not support
+	resume, download fails. I think this is OK because user normally
+	don't like to see that partially downloaded file is
+	overwritten(this is particularly true if file size is big).  This
+	option is useful when aria2 is used as download backend and
+	graceful falling back to overwritten behavior is preferable.
+	Added exit status value 8, which means download failed because
+	server did not support resume.
+	* src/AbstractCommand.cc
+	* src/DefaultPieceStorage.cc
+	* src/DownloadCommand.cc
+	* src/DownloadResultCode.h
+	* src/FileEntry.h
+	* src/FtpNegotiationCommand.cc
+	* src/HttpResponse.cc
+	* src/HttpResponseCommand.cc
+	* src/OptionHandlerFactory.cc
+	* src/RequestGroup.cc
+	* src/RequestGroup.h
+	* src/SegmentMan.cc
+	* src/SegmentMan.h
+	* src/prefs.cc
+	* src/prefs.h
+	* src/usage_text.h
+	* test/DefaultPieceStorageTest.cc
+	* test/SegmentManTest.cc
+
 2010-03-25  Tatsuhiro Tsujikawa  <t-tujikawa@users.sourceforge.net>
 
 	Added --remove-control-file option to -i list options.

src/AbstractCommand.cc

@@ -148,6 +148,15 @@ bool AbstractCommand::execute() {
       if(!_requestGroup->getPieceStorage().isNull()) {
         _segments.clear();
         _requestGroup->getSegmentMan()->getInFlightSegment(_segments, cuid);
+        if(!req.isNull() && _segments.empty()) {
+          // This command previously has assigned segments, but it is
+          // canceled. So discard current request chain.
+          if(logger->debug()) {
+            logger->debug("CUID#%s - It seems previously assigned segments are"
+                          "canceled. Restart.", util::itos(cuid).c_str());
+          }
+          return prepareForRetry(0);
+        }
         if(req.isNull() || req->getMaxPipelinedRequest() == 1 ||
            _requestGroup->getDownloadContext()->getFileEntries().size() == 1) {
           if(_segments.empty()) {
@@ -215,6 +224,9 @@ bool AbstractCommand::execute() {
                     util::itos(cuid).c_str(), req->getUri().c_str());
       _fileEntry->addURIResult(req->getUri(), err.getCode());
       _requestGroup->setLastUriResult(req->getUri(), err.getCode());
+      if(err.getCode() == downloadresultcode::CANNOT_RESUME) {
+        _requestGroup->increaseResumeFailureCount();
+      }
     }
     onAbort();
     tryReserved();
@@ -233,7 +245,6 @@ bool AbstractCommand::execute() {
     const unsigned int maxTries = getOption()->getAsInt(PREF_MAX_TRIES);
     bool isAbort = maxTries != 0 && req->getTryCount() >= maxTries;
     if(isAbort) {
-      onAbort();
       if(logger->info()) {
         logger->info(MSG_MAX_TRY, util::itos(cuid).c_str(), req->getTryCount());
       }
@@ -241,6 +252,10 @@ bool AbstractCommand::execute() {
                     req->getUri().c_str());
       _fileEntry->addURIResult(req->getUri(), err.getCode());
       _requestGroup->setLastUriResult(req->getUri(), err.getCode());
+      if(err.getCode() == downloadresultcode::CANNOT_RESUME) {
+        _requestGroup->increaseResumeFailureCount();
+      }
+      onAbort();
       tryReserved();
       return true;
     } else {
@@ -323,7 +338,52 @@ void AbstractCommand::onAbort() {
     logger->debug("CUID#%s - Aborting download", util::itos(cuid).c_str());
   }
   if(!_requestGroup->getPieceStorage().isNull()) {
-    _requestGroup->getSegmentMan()->cancelSegment(cuid);
+    SharedHandle<SegmentMan> segmentMan = _requestGroup->getSegmentMan();
+    segmentMan->cancelSegment(cuid);
+    // Don't do following process if BitTorrent is involved or files
+    // in DownloadContext is more than 1. The latter condition is
+    // limitation of current implementation.
+    if(!getOption()->getAsBool(PREF_ALWAYS_RESUME) &&
+       !_fileEntry.isNull() &&
+       segmentMan->calculateSessionDownloadLength() == 0 &&
+       !_requestGroup->p2pInvolved() &&
+       _requestGroup->getDownloadContext()->getFileEntries().size() == 1) {
+      const int maxTries = getOption()->getAsInt(PREF_MAX_RESUME_FAILURE_TRIES);
+      if((maxTries > 0 && _requestGroup->getResumeFailureCount() >= maxTries)||
+         _fileEntry->emptyRequestUri()) {
+        // Local file exists, but given servers(or at least contacted
+        // ones) doesn't support resume. Let's restart download from
+        // scratch.
+        logger->notice("CUID#%s - Failed to resume download."
+                       " Download from scratch.",
+                       util::itos(cuid).c_str());
+        if(logger->debug()) {
+          logger->debug("CUID#%s - Gathering URIs that has CANNOT_RESUME error",
+                        util::itos(cuid).c_str());
+        }
+        // Set PREF_ALWAYS_RESUME to V_TRUE to avoid repeating this
+        // process.
+        getOption()->put(PREF_ALWAYS_RESUME, V_TRUE);
+        std::deque<URIResult> res;
+        _fileEntry->extractURIResult(res, downloadresultcode::CANNOT_RESUME);
+        if(!res.empty()) {
+          segmentMan->cancelAllSegments();
+          segmentMan->eraseSegmentWrittenLengthMemo();
+          _requestGroup->getPieceStorage()->markPiecesDone(0);
+          std::vector<std::string> uris;
+          uris.reserve(res.size());
+          std::transform(res.begin(), res.end(), std::back_inserter(uris),
+                         std::mem_fun_ref(&URIResult::getURI));
+          if(logger->debug()) {
+            logger->debug("CUID#%s - %lu URIs found.",
+                          util::itos(cuid).c_str(),
+                          static_cast<unsigned long int>(uris.size()));
+          }
+          _fileEntry->addUris(uris.begin(), uris.end());
+          segmentMan->recognizeSegmentFor(_fileEntry);
+        }
+      }
+    }
   }
 }
 

src/DefaultPieceStorage.cc

@@ -596,6 +596,10 @@ void DefaultPieceStorage::markPiecesDone(uint64_t length)
 {
   if(length == bitfieldMan->getTotalLength()) {
     bitfieldMan->setAllBit();
+  } else if(length == 0) {
+    // TODO this would go to markAllPiecesUndone()
+    bitfieldMan->clearAllBit();
+    usedPieces.clear();
   } else {
     size_t numPiece = length/bitfieldMan->getBlockLength();
     if(numPiece > 0) {

src/DownloadCommand.cc

@@ -304,14 +304,15 @@ bool DownloadCommand::prepareForNextSegment() {
       if(!tempSegment->complete()) {
         return prepareForRetry(0);
       }
+      SharedHandle<SegmentMan> segmentMan = _requestGroup->getSegmentMan();
       SharedHandle<Segment> nextSegment =
-        _requestGroup->getSegmentMan()->getSegment(cuid,
-                                                   tempSegment->getIndex()+1);
-      if(!nextSegment.isNull() && nextSegment->getWrittenLength() == 0) {
+        segmentMan->getCleanSegmentIfOwnerIsIdle
+        (cuid, tempSegment->getIndex()+1);
+      if(nextSegment.isNull()) {
+        return prepareForRetry(0);
+      } else {
         e->commands.push_back(this);
         return false;
-      } else {
-        return prepareForRetry(0);
       }
     } else {
       return prepareForRetry(0);

src/DownloadResultCode.h

@@ -50,6 +50,7 @@ enum RESULT {
   TOO_SLOW_DOWNLOAD_SPEED = 5,
   NETWORK_PROBLEM = 6,
   IN_PROGRESS = 7,
+  CANNOT_RESUME = 8,
 };
 
 } // namespace downloadresultcode

src/FileEntry.h

@@ -269,6 +269,11 @@ public:
   }
 
   bool removeUri(const std::string& uri);
+
+  bool emptyRequestUri() const
+  {
+    return _uris.empty() && _inFlightRequests.empty() && _requestPool.empty();
+  }
 };
 
 // Returns the first FileEntry which isRequested() method returns

src/FtpNegotiationCommand.cc

@@ -394,7 +394,10 @@ bool FtpNegotiationCommand::onFileSizeDetermined(uint64_t totalLength)
       poolConnection();
       return false;
     }
-
+    // We have to make sure that command that has Request object must
+    // have segment after PieceStorage is initialized. See
+    // AbstractCommand::execute()
+    _requestGroup->getSegmentMan()->getSegment(cuid, 0);
     return true;
   } else {
     _requestGroup->adjustFilename
@@ -424,6 +427,10 @@ bool FtpNegotiationCommand::onFileSizeDetermined(uint64_t totalLength)
       return false;
     }
     _requestGroup->loadAndOpenFile(infoFile);
+    // We have to make sure that command that has Request object must
+    // have segment after PieceStorage is initialized. See
+    // AbstractCommand::execute()
+    _requestGroup->getSegmentMan()->getSegment(cuid, 0);
 
     prepareForNextAction(this);
 
@@ -574,7 +581,8 @@ bool FtpNegotiationCommand::recvRest(const SharedHandle<Segment>& segment) {
   // then throw exception here.
   if(status != 350) {
     if(!segment.isNull() && segment->getPositionToWrite() != 0) {
-      throw DL_ABORT_EX("FTP server doesn't support resuming.");
+      throw DL_ABORT_EX2("FTP server doesn't support resuming.",
+                         downloadresultcode::CANNOT_RESUME);
     }
   }
   sequence = SEQ_SEND_RETR;

src/HttpResponse.cc

@@ -79,7 +79,7 @@ void HttpResponse::validateResponse() const
     // compare the received range against the requested range
     RangeHandle responseRange = httpHeader->getRange();
     if(!httpRequest->isRangeSatisfied(responseRange)) {
-      throw DL_ABORT_EX
+      throw DL_ABORT_EX2
         (StringFormat
          (EX_INVALID_RANGE_HEADER,
           util::itos(httpRequest->getStartByte(), true).c_str(),
@@ -87,7 +87,8 @@ void HttpResponse::validateResponse() const
           util::uitos(httpRequest->getEntityLength(), true).c_str(),
           util::itos(responseRange->getStartByte(), true).c_str(),
           util::itos(responseRange->getEndByte(), true).c_str(),
-          util::uitos(responseRange->getEntityLength(), true).c_str()).str());
+          util::uitos(responseRange->getEntityLength(), true).c_str()).str(),
+         downloadresultcode::CANNOT_RESUME);
     }
   }
 }

src/HttpResponseCommand.cc

@@ -245,6 +245,9 @@ bool HttpResponseCommand::handleDefaultEncoding
     _requestGroup->loadAndOpenFile(infoFile);
     File file(_requestGroup->getFirstFilePath());
 
+    // We have to make sure that command that has Request object must
+    // have segment after PieceStorage is initialized. See
+    // AbstractCommand::execute()
     SharedHandle<Segment> segment =
       _requestGroup->getSegmentMan()->getSegment(cuid, 0);
     // pipelining requires implicit range specified. But the request for
@@ -348,6 +351,11 @@ bool HttpResponseCommand::handleOtherEncoding
     poolConnection();
     return true;
   }
+  // We have to make sure that command that has Request object must
+  // have segment after PieceStorage is initialized. See
+  // AbstractCommand::execute()
+  _requestGroup->getSegmentMan()->getSegment(cuid, 0);
+
   e->commands.push_back
     (createHttpDownloadCommand(httpResponse,
                                getTransferEncodingDecoder(httpResponse),

src/OptionHandlerFactory.cc

@@ -65,7 +65,18 @@ OptionHandlers OptionHandlerFactory::createOptionHandlers()
                                     V_FALSE));
     op->addTag(TAG_ADVANCED);
     handlers.push_back(op);
-  } 
+  }
+  {
+    SharedHandle<OptionHandler> op(new BooleanOptionHandler
+                                   (PREF_ALWAYS_RESUME,
+                                    TEXT_ALWAYS_RESUME,
+                                    V_TRUE,
+                                    OptionHandler::OPT_ARG));
+    op->addTag(TAG_ADVANCED);
+    op->addTag(TAG_FTP);
+    op->addTag(TAG_HTTP);
+    handlers.push_back(op);
+  }
 #ifdef ENABLE_ASYNC_DNS
   {
     SharedHandle<OptionHandler> op(new BooleanOptionHandler
@@ -322,6 +333,17 @@ OptionHandlers OptionHandlerFactory::createOptionHandlers()
     op->addTag(TAG_HTTP);
     handlers.push_back(op);
   }
+  {
+    SharedHandle<OptionHandler> op(new NumberOptionHandler
+                                   (PREF_MAX_RESUME_FAILURE_TRIES,
+                                    TEXT_MAX_RESUME_FAILURE_TRIES,
+                                    "0",
+                                    0));
+    op->addTag(TAG_ADVANCED);
+    op->addTag(TAG_FTP);
+    op->addTag(TAG_HTTP);
+    handlers.push_back(op);    
+  }
   {
     SharedHandle<OptionHandler> op(new BooleanOptionHandler
                                    (PREF_NO_CONF,

src/RequestGroup.cc

@@ -131,6 +131,7 @@ RequestGroup::RequestGroup(const SharedHandle<Option>& option):
   _maxUploadSpeedLimit(option->getAsInt(PREF_MAX_UPLOAD_LIMIT)),
   _belongsToGID(0),
   _requestGroupMan(0),
+  _resumeFailureCount(0),
   _logger(LogFactory::getInstance())
 {
   _fileAllocationEnabled = _option->get(PREF_FILE_ALLOCATION) != V_NONE;
@@ -1173,6 +1174,15 @@ void RequestGroup::setDownloadContext
   }
 }
 
+bool RequestGroup::p2pInvolved() const
+{
+#ifdef ENABLE_BITTORRENT
+  return _downloadContext->hasAttribute(bittorrent::BITTORRENT);
+#else // !ENABLE_BITTORRENT
+  return false;
+#endif // !ENABLE_BITTORRENT
+}
+
 gid_t RequestGroup::newGID()
 {
   if(_gidCounter == INT64_MAX) {

src/RequestGroup.h

@@ -167,6 +167,8 @@ private:
 
   RequestGroupMan* _requestGroupMan;
 
+  int _resumeFailureCount;
+
   Logger* _logger;
 
   void validateFilename(const std::string& expectedFilename,
@@ -500,6 +502,18 @@ public:
     _requestGroupMan = requestGroupMan;
   }
 
+  int getResumeFailureCount() const
+  {
+    return _resumeFailureCount;
+  }
+
+  void increaseResumeFailureCount()
+  {
+    ++_resumeFailureCount;
+  }
+
+  bool p2pInvolved() const;
+
   static void resetGIDCounter() { _gidCounter = 0; }
 
   static gid_t newGID();

src/SegmentMan.cc

@@ -208,14 +208,46 @@ void SegmentMan::getSegment
 }
 
 SharedHandle<Segment> SegmentMan::getSegment(cuid_t cuid, size_t index) {
-  if(_downloadContext->getNumPieces() <= index) {
+  if(index > 0 && _downloadContext->getNumPieces() <= index) {
     return SharedHandle<Segment>();
   }
   return checkoutSegment(cuid, _pieceStorage->getMissingPiece(index));
 }
 
+SharedHandle<Segment> SegmentMan::getCleanSegmentIfOwnerIsIdle
+(cuid_t cuid, size_t index)
+{
+  if(index > 0 && _downloadContext->getNumPieces() <= index) {
+    return SharedHandle<Segment>();
+  }
+  for(SegmentEntries::const_iterator itr = usedSegmentEntries.begin(),
+        eoi = usedSegmentEntries.end(); itr != eoi; ++itr) {
+    const SharedHandle<SegmentEntry>& segmentEntry = *itr;
+    if(segmentEntry->segment->getIndex() == index) {
+      if(segmentEntry->cuid == cuid) {
+        return segmentEntry->segment;
+      }
+      if(segmentEntry->segment->getWrittenLength() > 0) {
+        return SharedHandle<Segment>();
+      }
+      cuid_t owner = segmentEntry->cuid;
+      SharedHandle<PeerStat> ps = getPeerStat(owner);
+      if(ps.isNull() || (!ps.isNull() && ps->getStatus() == PeerStat::IDLE)) {
+        cancelSegment(owner);
+        return getSegment(cuid, index);
+      } else {
+        return SharedHandle<Segment>();
+      }
+    }
+  }
+  return checkoutSegment(cuid, _pieceStorage->getMissingPiece(index));
+}
+
 void SegmentMan::cancelSegment(const SharedHandle<Segment>& segment)
 {
+  if(logger->debug()) {
+    logger->debug("Canceling segment#%d", segment->getIndex());
+  }
   _pieceStorage->cancelPiece(segment->getPiece());
   _segmentWrittenLengthMemo[segment->getIndex()] = segment->getWrittenLength();
   if(logger->debug()) {
@@ -253,6 +285,21 @@ void SegmentMan::cancelSegment
   }
 }
 
+void SegmentMan::cancelAllSegments()
+{
+  for(std::deque<SharedHandle<SegmentEntry> >::iterator itr =
+        usedSegmentEntries.begin(), eoi = usedSegmentEntries.end();
+      itr != eoi; ++itr) {
+    cancelSegment((*itr)->segment);
+  }
+  usedSegmentEntries.clear();
+}
+
+void SegmentMan::eraseSegmentWrittenLengthMemo()
+{
+  _segmentWrittenLengthMemo.clear();
+}
+
 class FindSegmentEntry {
 private:
   SharedHandle<Segment> _segment;
@@ -304,6 +351,18 @@ void SegmentMan::registerPeerStat(const SharedHandle<PeerStat>& peerStat)
   peerStats.push_back(peerStat);
 }
 
+SharedHandle<PeerStat> SegmentMan::getPeerStat(cuid_t cuid) const
+{
+  for(std::vector<SharedHandle<PeerStat> >::const_iterator i =
+        peerStats.begin(), eoi = peerStats.end(); i != eoi; ++i) {
+    if((*i)->getCuid() == cuid) {
+      return *i;
+    }
+  }
+  return SharedHandle<PeerStat>();
+}
+
+
 class PeerStatHostProtoEqual {
 private:
   const SharedHandle<PeerStat>& _peerStat;

src/SegmentMan.h

@@ -157,6 +157,15 @@ public:
    * whose isNull call is true.
    */
   SharedHandle<Segment> getSegment(cuid_t cuid, size_t index);
+
+  // Returns a segment whose index is index.  If the segment whose
+  // index is index is free, it is assigned to cuid and it is
+  // returned.  If it has already be assigned to another cuid, and if
+  // it is idle state and segment's written length is 0, then cancels
+  // the assignment and re-attach the segment to given cuid and the
+  // segment is returned. Otherwise returns null.
+  SharedHandle<Segment> getCleanSegmentIfOwnerIsIdle(cuid_t cuid, size_t index);
+
   /**
    * Updates download status.
    */
@@ -169,6 +178,10 @@ public:
 
   void cancelSegment(cuid_t cuid, const SharedHandle<Segment>& segment);
 
+  void cancelAllSegments();
+
+  void eraseSegmentWrittenLengthMemo();
+
   /**
    * Tells SegmentMan that the segment has been downloaded successfully.
    */
@@ -204,6 +217,8 @@ public:
     return peerStats;
   }
 
+  SharedHandle<PeerStat> getPeerStat(cuid_t cuid) const;
+
   // If there is slower PeerStat than given peerStat for the same
   // hostname and protocol in _fastestPeerStats, the former is
   // replaced with latter. If there are no PeerStat with same hostname

src/prefs.cc

@@ -180,6 +180,10 @@ const std::string PREF_DISABLE_IPV6("disable-ipv6");
 const std::string PREF_HUMAN_READABLE("human-readable");
 // value: true | false
 const std::string PREF_REMOVE_CONTROL_FILE("remove-control-file");
+// value: true | false
+const std::string PREF_ALWAYS_RESUME("always-resume");
+// value: 1*digit
+const std::string PREF_MAX_RESUME_FAILURE_TRIES("max-resume-failure-tries");
 
 /**
  * FTP related preferences

src/prefs.h

@@ -184,6 +184,10 @@ extern const std::string PREF_DISABLE_IPV6;
 extern const std::string PREF_HUMAN_READABLE;
 // value: true | false
 extern const std::string PREF_REMOVE_CONTROL_FILE;
+// value: true | false
+extern const std::string PREF_ALWAYS_RESUME;
+// value: 1*digit
+extern const std::string PREF_MAX_RESUME_FAILURE_TRIES;
 
 /**
  * FTP related preferences

src/usage_text.h

@@ -640,3 +640,18 @@
     "                              with --allow-overwrite=true, download always\n" \
     "                              starts from scratch. This will be useful for\n" \
     "                              users behind proxy server which disables resume.")
+#define TEXT_ALWAYS_RESUME                      \
+  _(" --always-resume[=true|false] Always resume download. If false is given, when\n" \
+    "                              all given URIs does not support resume or aria2\n" \
+    "                              encounters N URIs which does not support resume\n" \
+    "                              (N is the value specified using\n"   \
+    "                              --max-resume-failure-tries option), aria2\n" \
+    "                              download file from scratch.\n"       \
+    "                              See --max-resume-failure-tries option.")
+#define TEXT_MAX_RESUME_FAILURE_TRIES                                   \
+  _(" --max-resume-failure-tries=N When used with --always-resume=false, aria2\n" \
+    "                              download file from scratch when aria2 detects N\n" \
+    "                              number of URIs that does not support resume. If N\n" \
+    "                              is 0, aria2 download file from scratch when all\n" \
+    "                              given URIs do not support resume.\n" \
+    "                              See --always-resume option.")

test/DefaultPieceStorageTest.cc

@@ -267,6 +267,9 @@ void DefaultPieceStorageTest::testMarkPiecesDone()
   for(size_t i = 0; i < (totalLength+pieceLength-1)/pieceLength; ++i) {
     CPPUNIT_ASSERT(ps.hasPiece(i));
   }
+
+  ps.markPiecesDone(0);
+  CPPUNIT_ASSERT_EQUAL((uint64_t)0, ps.getCompletedLength());
 }
 
 void DefaultPieceStorageTest::testGetCompletedLength()

test/SegmentManTest.cc

@@ -18,18 +18,34 @@ class SegmentManTest:public CppUnit::TestFixture {
   CPPUNIT_TEST(testCompleteSegment);
   CPPUNIT_TEST(testGetSegment_sameFileEntry);
   CPPUNIT_TEST(testRegisterPeerStat);
+  CPPUNIT_TEST(testCancelAllSegments);
+  CPPUNIT_TEST(testGetPeerStat);
+  CPPUNIT_TEST(testGetCleanSegmentIfOwnerIsIdle);
   CPPUNIT_TEST_SUITE_END();
 private:
-
+  SharedHandle<Option> _option;
+  SharedHandle<DownloadContext> _dctx;
+  SharedHandle<DefaultPieceStorage> _pieceStorage;
+  SharedHandle<SegmentMan> _segmentMan;
 public:
-  void setUp() {
+  void setUp()
+  {
+    size_t pieceLength = 1024*1024;
+    uint64_t totalLength = 64*1024*1024;
+    _option.reset(new Option());
+    _dctx.reset
+      (new DownloadContext(pieceLength, totalLength, "aria2.tar.bz2"));
+    _pieceStorage.reset(new DefaultPieceStorage(_dctx, _option.get()));
+    _segmentMan.reset(new SegmentMan(_option.get(), _dctx, _pieceStorage));
   }
 
   void testNullBitfield();
   void testCompleteSegment();
-  void testGetPeerStat();
   void testGetSegment_sameFileEntry();
   void testRegisterPeerStat();
+  void testCancelAllSegments();
+  void testGetPeerStat();
+  void testGetCleanSegmentIfOwnerIsIdle();
 };
 
 
@@ -144,4 +160,49 @@ void SegmentManTest::testRegisterPeerStat()
   CPPUNIT_ASSERT_EQUAL((size_t)2, segman.getPeerStats().size());
 }
 
+void SegmentManTest::testCancelAllSegments()
+{
+  _segmentMan->getSegment(1, 0);
+  _segmentMan->getSegment(2, 1);
+  CPPUNIT_ASSERT(_segmentMan->getSegment(3, 0).isNull());
+  CPPUNIT_ASSERT(_segmentMan->getSegment(4, 1).isNull());
+  _segmentMan->cancelAllSegments();
+  CPPUNIT_ASSERT(!_segmentMan->getSegment(3, 0).isNull());
+  CPPUNIT_ASSERT(!_segmentMan->getSegment(4, 1).isNull());
+}
+
+void SegmentManTest::testGetPeerStat()
+{
+  SharedHandle<PeerStat> peerStat1(new PeerStat(1));
+  _segmentMan->registerPeerStat(peerStat1);
+  CPPUNIT_ASSERT_EQUAL((cuid_t)1, _segmentMan->getPeerStat(1)->getCuid());
+}
+
+void SegmentManTest::testGetCleanSegmentIfOwnerIsIdle()
+{
+  SharedHandle<Segment> seg1 = _segmentMan->getSegment(1, 0);
+  SharedHandle<Segment> seg2 = _segmentMan->getSegment(2, 1);
+  seg2->updateWrittenLength(100);
+  CPPUNIT_ASSERT(!_segmentMan->getCleanSegmentIfOwnerIsIdle(3, 0).isNull());
+  SharedHandle<PeerStat> peerStat3(new PeerStat(3));
+  _segmentMan->registerPeerStat(peerStat3);
+  CPPUNIT_ASSERT(!_segmentMan->getCleanSegmentIfOwnerIsIdle(4, 0).isNull());
+  SharedHandle<PeerStat> peerStat4(new PeerStat(4));
+  peerStat4->downloadStart();
+  _segmentMan->registerPeerStat(peerStat4);
+  // Owner PeerStat is not IDLE
+  CPPUNIT_ASSERT(_segmentMan->getCleanSegmentIfOwnerIsIdle(5, 0).isNull());
+  // Segment::updateWrittenLength != 0
+  CPPUNIT_ASSERT(_segmentMan->getCleanSegmentIfOwnerIsIdle(5, 1).isNull());
+
+  // Test with UnknownLengthPieceStorage
+  SharedHandle<DownloadContext> dctx(new DownloadContext(1024, 0, "aria2"));
+  SharedHandle<UnknownLengthPieceStorage> ps
+    (new UnknownLengthPieceStorage(dctx, _option.get()));
+  _segmentMan.reset(new SegmentMan(_option.get(), dctx, ps));
+  
+  CPPUNIT_ASSERT(!_segmentMan->getCleanSegmentIfOwnerIsIdle(1, 0).isNull());
+
+}
+
 } // namespace aria2