Pack Piece message header and body into one chunk

2 files changed, 18 insertions, 12 deletions

diff --git a/src/BtPieceMessage.cc b/src/BtPieceMessage.cc
index b0362d72..04025f48 100644
--- a/src/BtPieceMessage.cc
+++ b/src/BtPieceMessage.cc
@@ -157,7 +157,7 @@ void BtPieceMessage::doReceivedAction()
 
 size_t BtPieceMessage::MESSAGE_HEADER_LENGTH = 13;
 
-unsigned char* BtPieceMessage::createMessageHeader()
+void BtPieceMessage::createMessageHeader(unsigned char* msgHeader) const
 {
   /**
    * len --- 9+blockLength, 4bytes
@@ -166,12 +166,10 @@ unsigned char* BtPieceMessage::createMessageHeader()
    * begin --- begin, 4bytes
    * total: 13bytes
    */
-  unsigned char* msgHeader = new unsigned char[MESSAGE_HEADER_LENGTH];
   bittorrent::createPeerMessageString(msgHeader, MESSAGE_HEADER_LENGTH,
                                       9+blockLength_, ID);
   bittorrent::setIntParam(&msgHeader[5], index_);
   bittorrent::setIntParam(&msgHeader[9], begin_);
-  return msgHeader;
 }
 
 size_t BtPieceMessage::getMessageHeaderLength()
@@ -181,13 +179,19 @@ size_t BtPieceMessage::getMessageHeaderLength()
 
 namespace {
 struct PieceSendUpdate : public ProgressUpdate {
-  PieceSendUpdate(const SharedHandle<Peer>& peer)
-    : peer(peer) {}
+  PieceSendUpdate(const SharedHandle<Peer>& peer, size_t headerLength)
+    : peer(peer), headerLength(headerLength) {}
   virtual void update(size_t length, bool complete)
   {
+    if(headerLength > 0) {
+      size_t m = std::min(headerLength, length);
+      headerLength -= m;
+      length -= m;
+    }
     peer->updateUploadLength(length);
   }
   SharedHandle<Peer> peer;
+  size_t headerLength;
 };
 } // namespace
 
@@ -201,8 +205,6 @@ void BtPieceMessage::send()
                   getPeer()->getIPAddress().c_str(),
                   getPeer()->getPort(),
                   toString().c_str()));
-  getPeerConnection()->pushBytes(createMessageHeader(),
-                                 getMessageHeaderLength());
   int64_t pieceDataOffset =
     static_cast<int64_t>(index_)*downloadContext_->getPieceLength()+begin_;
   pushPieceData(pieceDataOffset, blockLength_);
@@ -211,14 +213,18 @@ void BtPieceMessage::send()
 void BtPieceMessage::pushPieceData(int64_t offset, int32_t length) const
 {
   assert(length <= 16*1024);
-  array_ptr<unsigned char> buf(new unsigned char[length]);
+  array_ptr<unsigned char> buf
+    (new unsigned char[length+MESSAGE_HEADER_LENGTH]);
+  createMessageHeader(buf);
   ssize_t r;
-  r = getPieceStorage()->getDiskAdaptor()->readData(buf, length, offset);
+  r = getPieceStorage()->getDiskAdaptor()->readData(buf+MESSAGE_HEADER_LENGTH,
+                                                    length, offset);
   if(r == length) {
     unsigned char* dbuf = buf;
     buf.reset(0);
-    getPeerConnection()->pushBytes(dbuf, length,
-                                   new PieceSendUpdate(getPeer()));
+    getPeerConnection()->pushBytes(dbuf, length+MESSAGE_HEADER_LENGTH,
+                                   new PieceSendUpdate(getPeer(),
+                                                       MESSAGE_HEADER_LENGTH));
     // To avoid upload rate overflow, we update the length here at
     // once.
     downloadContext_->updateUploadLength(length);
diff --git a/src/BtPieceMessage.h b/src/BtPieceMessage.h
index 26e01ea5..09fef7dc 100644
--- a/src/BtPieceMessage.h
+++ b/src/BtPieceMessage.h
@@ -96,7 +96,7 @@ public:
 
   virtual void doReceivedAction();
 
-  unsigned char* createMessageHeader();
+  void createMessageHeader(unsigned char* msgHeader) const;
 
   size_t getMessageHeaderLength();