Flush cached data on checkout/cancel Segment

This change prevents the cached data from overlap because of BT peers.

src/DownloadCommand.cc

@@ -126,7 +126,7 @@ void flushWrDiskCacheEntry(WrDiskCache* wrDiskCache,
                            const SharedHandle<Segment>& segment)
 {
   const SharedHandle<Piece>& piece = segment->getPiece();
-  if(piece && piece->getWrDiskCacheEntry()) {
+  if(piece->getWrDiskCacheEntry()) {
     piece->flushWrCache(wrDiskCache);
     if(piece->getWrDiskCacheEntry()->getError() !=
        WrDiskCacheEntry::CACHE_ERR_SUCCESS) {

src/GrowSegment.cc

@@ -62,6 +62,7 @@ std::string GrowSegment::getDigest()
 void GrowSegment::clear(WrDiskCache* diskCache)
 {
   writtenLength_ = 0;
+  // cache won't be used in this object.
   piece_->clearAllBlock(0);
 }
 

src/Piece.cc

@@ -319,6 +319,9 @@ void Piece::removeUser(cuid_t cuid)
 void Piece::initWrCache(WrDiskCache* diskCache,
                         const SharedHandle<DiskAdaptor>& diskAdaptor)
 {
+  if(!diskCache) {
+    return;
+  }
   assert(wrCache_ == 0);
   wrCache_ = new WrDiskCacheEntry(diskAdaptor);
   bool rv = diskCache->add(wrCache_);
@@ -327,6 +330,9 @@ void Piece::initWrCache(WrDiskCache* diskCache,
 
 void Piece::flushWrCache(WrDiskCache* diskCache)
 {
+  if(!diskCache) {
+    return;
+  }
   assert(wrCache_);
   ssize_t size = static_cast<ssize_t>(wrCache_->getSize());
   diskCache->update(wrCache_, -size);
@@ -335,6 +341,9 @@ void Piece::flushWrCache(WrDiskCache* diskCache)
 
 void Piece::clearWrCache(WrDiskCache* diskCache)
 {
+  if(!diskCache) {
+    return;
+  }
   assert(wrCache_);
   ssize_t size = static_cast<ssize_t>(wrCache_->getSize());
   diskCache->update(wrCache_, -size);
@@ -344,8 +353,11 @@ void Piece::clearWrCache(WrDiskCache* diskCache)
 void Piece::updateWrCache(WrDiskCache* diskCache, unsigned char* data,
                           size_t offset, size_t len, int64_t goff)
 {
-  A2_LOG_DEBUG(fmt("updateWrCache entry=%p", wrCache_));
+  if(!diskCache) {
+    return;
+  }
   assert(wrCache_);
+  A2_LOG_DEBUG(fmt("updateWrCache entry=%p", wrCache_));
   WrDiskCacheEntry::DataCell* cell = new WrDiskCacheEntry::DataCell();
   cell->goff = goff;
   cell->data = data;
@@ -360,7 +372,7 @@ void Piece::updateWrCache(WrDiskCache* diskCache, unsigned char* data,
 
 void Piece::releaseWrCache(WrDiskCache* diskCache)
 {
-  if(wrCache_) {
+  if(diskCache && wrCache_) {
     diskCache->remove(wrCache_);
     delete wrCache_;
     wrCache_ = 0;

src/SegmentMan.cc

@@ -53,6 +53,8 @@
 #include "FileEntry.h"
 #include "wallclock.h"
 #include "fmt.h"
+#include "WrDiskCacheEntry.h"
+#include "DownloadFailureException.h"
 
 namespace aria2 {
 
@@ -112,6 +114,22 @@ void SegmentMan::setDownloadContext
   downloadContext_ = downloadContext;
 }
 
+namespace {
+void flushWrDiskCache(WrDiskCache* wrDiskCache,
+                      const SharedHandle<Piece>& piece)
+{
+  piece->flushWrCache(wrDiskCache);
+  if(piece->getWrDiskCacheEntry()->getError() !=
+     WrDiskCacheEntry::CACHE_ERR_SUCCESS) {
+    piece->clearAllBlock(wrDiskCache);
+    throw DOWNLOAD_FAILURE_EXCEPTION2
+      (fmt("Write disk cache flush failure index=%lu",
+           static_cast<unsigned long>(piece->getIndex())),
+       piece->getWrDiskCacheEntry()->getErrorCode());
+  }
+}
+} // namespace
+
 SharedHandle<Segment> SegmentMan::checkoutSegment
 (cuid_t cuid, const SharedHandle<Piece>& piece)
 {
@@ -121,6 +139,15 @@ SharedHandle<Segment> SegmentMan::checkoutSegment
   A2_LOG_DEBUG(fmt("Attach segment#%lu to CUID#%" PRId64 ".",
                    static_cast<unsigned long>(piece->getIndex()),
                    cuid));
+
+  if(piece->getWrDiskCacheEntry()) {
+    // Flush cached data here, because the cached data may be overlapped
+    // if BT peers are involved.
+    A2_LOG_DEBUG(fmt("Flushing cached data, size=%lu",
+                     piece->getWrDiskCacheEntry()->getSize()));
+    flushWrDiskCache(pieceStorage_->getWrDiskCache(), piece);
+  }
+
   piece->setUsedBySegment(true);
   SharedHandle<Segment> segment;
   if(piece->getLength() == 0) {
@@ -136,6 +163,7 @@ SharedHandle<Segment> SegmentMan::checkoutSegment
                    segment->getLength(),
                    segment->getSegmentLength(),
                    segment->getWrittenLength()));
+
   if(piece->getLength() > 0) {
     std::map<size_t, int32_t>::iterator positr =
       segmentWrittenLengthMemo_.find(segment->getIndex());
@@ -254,12 +282,24 @@ void SegmentMan::cancelSegmentInternal
 {
   A2_LOG_DEBUG(fmt("Canceling segment#%lu",
                    static_cast<unsigned long>(segment->getIndex())));
-  segment->getPiece()->setUsedBySegment(false);
+  const SharedHandle<Piece>& piece = segment->getPiece();
-  pieceStorage_->cancelPiece(segment->getPiece(), cuid);
+  piece->setUsedBySegment(false);
+  pieceStorage_->cancelPiece(piece, cuid);
   segmentWrittenLengthMemo_[segment->getIndex()] = segment->getWrittenLength();
   A2_LOG_DEBUG(fmt("Memorized segment index=%lu, writtenLength=%d",
                    static_cast<unsigned long>(segment->getIndex()),
                    segment->getWrittenLength()));
+  // TODO In PieceStorage::cancelPiece(), WrDiskCacheEntry may be
+  // released.
+  if(piece->getWrDiskCacheEntry()) {
+    // Flush cached data here, because the cached data may be overlapped
+    // if BT peers are involved.
+    A2_LOG_DEBUG(fmt("Flushing cached data, size=%lu",
+                     piece->getWrDiskCacheEntry()->getSize()));
+    flushWrDiskCache(pieceStorage_->getWrDiskCache(), piece);
+    // TODO Exception may cause some segments (pieces) are not
+    // canceled.
+  }
 }
 
 void SegmentMan::cancelSegment(cuid_t cuid) {