From 9eea95945568ecc01221fb827ff211d98207bf79 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Wed, 12 Jun 2013 23:12:01 +0900
Subject: [PATCH] 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      | 11 +++++++----
 src/BtPieceMessage.cc     | 18 +++++++++++++++---
 test/BtHaveMessageTest.cc |  1 +
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/BtHaveMessage.cc b/src/BtHaveMessage.cc
index 50b40f39..92ae8f8c 100644
--- a/src/BtHaveMessage.cc
+++ b/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);
+    }
   }
 }
 
diff --git a/src/BtPieceMessage.cc b/src/BtPieceMessage.cc
index 04025f48..12ffa2f0 100644
--- a/src/BtPieceMessage.cc
+++ b/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);
diff --git a/test/BtHaveMessageTest.cc b/test/BtHaveMessageTest.cc
index a0cd5a77..5c1c7e2a 100644
--- a/test/BtHaveMessageTest.cc
+++ b/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.");