Improve peer bitfield and piece stat update

On HAVE message reception and PIECE message transmission, update peer
bitfield and piece stat only if current peer bitfield indicates it
does not have the piece.

src/BtHaveMessage.cc

@@ -55,10 +55,13 @@ void BtHaveMessage::doReceivedAction()
   if(isMetadataGetMode()) {
     return;
   }
-  getPeer()->updateBitfield(getIndex(), 1);
-  getPieceStorage()->addPieceStats(getIndex());
-  if(getPeer()->isSeeder() && getPieceStorage()->downloadFinished()) {
-    throw DL_ABORT_EX(MSG_GOOD_BYE_SEEDER);
+  size_t index = getIndex();
+  if(!getPeer()->hasPiece(index)) {
+    getPeer()->updateBitfield(index, 1);
+    getPieceStorage()->addPieceStats(index);
+    if(getPeer()->isSeeder() && getPieceStorage()->downloadFinished()) {
+      throw DL_ABORT_EX(MSG_GOOD_BYE_SEEDER);
+    }
   }
 }
 

src/BtPieceMessage.cc

@@ -179,8 +179,10 @@ size_t BtPieceMessage::getMessageHeaderLength()
 
 namespace {
 struct PieceSendUpdate : public ProgressUpdate {
-  PieceSendUpdate(const SharedHandle<Peer>& peer, size_t headerLength)
-    : peer(peer), headerLength(headerLength) {}
+  PieceSendUpdate(const SharedHandle<Peer>& peer, size_t headerLength,
+                  size_t index, const SharedHandle<PieceStorage>& pieceStorage)
+    : peer(peer), headerLength(headerLength), index(index),
+      pieceStorage(pieceStorage) {}
   virtual void update(size_t length, bool complete)
   {
     if(headerLength > 0) {
@@ -189,9 +191,17 @@ struct PieceSendUpdate : public ProgressUpdate {
       length -= m;
     }
     peer->updateUploadLength(length);
+    if(complete && !peer->hasPiece(index)) {
+      // Update peer's bitfield because peer may not send HAVE message
+      // to us.
+      peer->updateBitfield(index, 1);
+      pieceStorage->addPieceStats(index);
+    }
   }
   SharedHandle<Peer> peer;
   size_t headerLength;
+  size_t index;
+  SharedHandle<PieceStorage> pieceStorage;
 };
 } // namespace
 
@@ -224,7 +234,9 @@ void BtPieceMessage::pushPieceData(int64_t offset, int32_t length) const
     buf.reset(0);
     getPeerConnection()->pushBytes(dbuf, length+MESSAGE_HEADER_LENGTH,
                                    new PieceSendUpdate(getPeer(),
-                                                       MESSAGE_HEADER_LENGTH));
+                                                       MESSAGE_HEADER_LENGTH,
+                                                       index_,
+                                                       getPieceStorage()));
     // To avoid upload rate overflow, we update the length here at
     // once.
     downloadContext_->updateUploadLength(length);

test/BtHaveMessageTest.cc

@@ -115,6 +115,7 @@ void BtHaveMessageTest::testDoReceivedAction_goodByeSeeder()
   msg.doReceivedAction();
 
   pieceStorage->setDownloadFinished(true);
+  peer->updateBitfield(1, 0);
   try {
     msg.doReceivedAction();
     CPPUNIT_FAIL("exception must be thrown.");