Updating upload size and speed must be done separately

src/BtPieceMessage.cc

@@ -99,8 +99,8 @@ void BtPieceMessage::doReceivedAction()
   }
   auto slot = getBtMessageDispatcher()->getOutstandingRequest(index_, begin_,
                                                               blockLength_);
-  getPeer()->updateDownloadLength(blockLength_);
-  downloadContext_->updateDownloadLength(blockLength_);
+  getPeer()->updateDownload(blockLength_);
+  downloadContext_->updateDownload(blockLength_);
   if (slot) {
     getPeer()->snubbing(false);
     std::shared_ptr<Piece> piece = getPieceStorage()->getPiece(index_);
@@ -175,8 +175,9 @@ size_t BtPieceMessage::getMessageHeaderLength()
 
 namespace {
 struct PieceSendUpdate : public ProgressUpdate {
-  PieceSendUpdate(std::shared_ptr<Peer> peer, size_t headerLength)
-      : peer(std::move(peer)), headerLength(headerLength)
+  PieceSendUpdate(DownloadContext* dctx, std::shared_ptr<Peer> peer,
+                  size_t headerLength)
+      : dctx(dctx), peer(std::move(peer)), headerLength(headerLength)
   {
   }
   virtual void update(size_t length, bool complete) CXX11_OVERRIDE
@@ -187,7 +188,9 @@ struct PieceSendUpdate : public ProgressUpdate {
       length -= m;
     }
     peer->updateUploadLength(length);
+    dctx->updateUploadLength(length);
   }
+  DownloadContext* dctx;
   std::shared_ptr<Peer> peer;
   size_t headerLength;
 };
@@ -216,12 +219,13 @@ void BtPieceMessage::pushPieceData(int64_t offset, int32_t length) const
   r = getPieceStorage()->getDiskAdaptor()->readData(
       buf.get() + MESSAGE_HEADER_LENGTH, length, offset);
   if (r == length) {
+    const auto& peer = getPeer();
     getPeerConnection()->pushBytes(
         buf.release(), length + MESSAGE_HEADER_LENGTH,
-        make_unique<PieceSendUpdate>(getPeer(), MESSAGE_HEADER_LENGTH));
-    // To avoid upload rate overflow, we update the length here at
-    // once.
-    downloadContext_->updateUploadLength(length);
+        make_unique<PieceSendUpdate>(downloadContext_, peer,
+                                     MESSAGE_HEADER_LENGTH));
+    peer->updateUploadSpeed(length);
+    downloadContext_->updateUploadSpeed(length);
   }
   else {
     throw DL_ABORT_EX(EX_DATA_READ);

src/DownloadCommand.cc

@@ -194,8 +194,8 @@ bool DownloadCommand::executeInternal()
       bufSize = streamFilter_->getBytesProcessed();
     }
     getSocketRecvBuffer()->drain(bufSize);
-    peerStat_->updateDownloadLength(bufSize);
-    getDownloadContext()->updateDownloadLength(bufSize);
+    peerStat_->updateDownload(bufSize);
+    getDownloadContext()->updateDownload(bufSize);
   }
   bool segmentPartComplete = false;
   // Note that GrowSegment::complete() always returns false.

src/DownloadContext.cc

@@ -288,19 +288,28 @@ void DownloadContext::setSignature(std::unique_ptr<Signature> signature)
   signature_ = std::move(signature);
 }
 
-void DownloadContext::updateDownloadLength(size_t bytes)
+void DownloadContext::updateDownload(size_t bytes)
 {
-  netStat_.updateDownloadLength(bytes);
+  netStat_.updateDownload(bytes);
   RequestGroupMan* rgman = ownerRequestGroup_->getRequestGroupMan();
   if (rgman) {
-    rgman->getNetStat().updateDownloadLength(bytes);
+    rgman->getNetStat().updateDownload(bytes);
+  }
+}
+
+void DownloadContext::updateUploadSpeed(size_t bytes)
+{
+  netStat_.updateUploadSpeed(bytes);
+  auto rgman = ownerRequestGroup_->getRequestGroupMan();
+  if (rgman) {
+    rgman->getNetStat().updateUploadSpeed(bytes);
   }
 }
 
 void DownloadContext::updateUploadLength(size_t bytes)
 {
   netStat_.updateUploadLength(bytes);
-  RequestGroupMan* rgman = ownerRequestGroup_->getRequestGroupMan();
+  auto rgman = ownerRequestGroup_->getRequestGroupMan();
   if (rgman) {
     rgman->getNetStat().updateUploadLength(bytes);
   }

src/DownloadContext.h

@@ -225,11 +225,12 @@ public:
 
   // This method also updates global download length held by
   // RequestGroupMan via getOwnerRequestGroup().
-  void updateDownloadLength(size_t bytes);
+  void updateDownload(size_t bytes);
 
   // This method also updates global upload length held by
   // RequestGroupMan via getOwnerRequestGroup().
   void updateUploadLength(size_t bytes);
+  void updateUploadSpeed(size_t bytes);
 };
 
 } // namespace aria2

src/NetStat.cc

@@ -68,18 +68,25 @@ int NetStat::calculateAvgUploadSpeed()
   return avgUploadSpeed_ = uploadSpeed_.calculateAvgSpeed();
 }
 
-void NetStat::updateDownloadLength(size_t bytes)
+void NetStat::updateDownload(size_t bytes)
 {
   downloadSpeed_.update(bytes);
   sessionDownloadLength_ += bytes;
 }
 
-void NetStat::updateUploadLength(size_t bytes)
+void NetStat::updateUpload(size_t bytes)
 {
   uploadSpeed_.update(bytes);
   sessionUploadLength_ += bytes;
 }
 
+void NetStat::updateUploadSpeed(size_t bytes) { uploadSpeed_.update(bytes); }
+
+void NetStat::updateUploadLength(size_t bytes)
+{
+  sessionUploadLength_ += bytes;
+}
+
 int NetStat::getMaxDownloadSpeed() const
 {
   return downloadSpeed_.getMaxSpeed();

src/NetStat.h

@@ -67,7 +67,11 @@ public:
 
   int calculateAvgUploadSpeed();
 
-  void updateDownloadLength(size_t bytes);
+  void updateDownload(size_t bytes);
+
+  void updateUpload(size_t bytes);
+
+  void updateUploadSpeed(size_t bytes);
 
   void updateUploadLength(size_t bytes);
 

src/Peer.cc

@@ -174,16 +174,22 @@ void Peer::snubbing(bool b)
   res_->snubbing(b);
 }
 
+void Peer::updateUploadSpeed(int32_t bytes)
+{
+  assert(res_);
+  res_->updateUploadSpeed(bytes);
+}
+
 void Peer::updateUploadLength(int32_t bytes)
 {
   assert(res_);
   res_->updateUploadLength(bytes);
 }
 
-void Peer::updateDownloadLength(int32_t bytes)
+void Peer::updateDownload(int32_t bytes)
 {
   assert(res_);
-  res_->updateDownloadLength(bytes);
+  res_->updateDownload(bytes);
 }
 
 void Peer::updateSeeder()

src/Peer.h

@@ -171,9 +171,11 @@ public:
 
   void snubbing(bool b);
 
+  void updateUploadSpeed(int32_t bytes);
+
   void updateUploadLength(int32_t bytes);
 
-  void updateDownloadLength(int32_t bytes);
+  void updateDownload(int32_t bytes);
 
   /**
    * Returns the transfer rate from localhost to remote host.

src/PeerSessionResource.cc

@@ -196,6 +196,11 @@ int64_t PeerSessionResource::uploadLength() const
   return netStat_.getSessionUploadLength();
 }
 
+void PeerSessionResource::updateUploadSpeed(int32_t bytes)
+{
+  netStat_.updateUploadSpeed(bytes);
+}
+
 void PeerSessionResource::updateUploadLength(int32_t bytes)
 {
   netStat_.updateUploadLength(bytes);
@@ -206,9 +211,9 @@ int64_t PeerSessionResource::downloadLength() const
   return netStat_.getSessionDownloadLength();
 }
 
-void PeerSessionResource::updateDownloadLength(int32_t bytes)
+void PeerSessionResource::updateDownload(int32_t bytes)
 {
-  netStat_.updateDownloadLength(bytes);
+  netStat_.updateDownload(bytes);
   lastDownloadUpdate_ = global::wallclock();
 }
 

src/PeerSessionResource.h

@@ -182,11 +182,13 @@ public:
 
   int64_t uploadLength() const;
 
+  void updateUploadSpeed(int32_t bytes);
+
   void updateUploadLength(int32_t bytes);
 
   int64_t downloadLength() const;
 
-  void updateDownloadLength(int32_t bytes);
+  void updateDownload(int32_t bytes);
 
   const Timer& getLastDownloadUpdate() const { return lastDownloadUpdate_; }
 

src/PeerStat.cc

@@ -67,15 +67,9 @@ int PeerStat::calculateAvgUploadSpeed()
   return netStat_.calculateAvgUploadSpeed();
 }
 
-void PeerStat::updateDownloadLength(size_t bytes)
-{
-  netStat_.updateDownloadLength(bytes);
-}
+void PeerStat::updateDownload(size_t bytes) { netStat_.updateDownload(bytes); }
 
-void PeerStat::updateUploadLength(size_t bytes)
-{
-  netStat_.updateUploadLength(bytes);
-}
+void PeerStat::updateUpload(size_t bytes) { netStat_.updateUpload(bytes); }
 
 int PeerStat::getMaxDownloadSpeed() const
 {

src/PeerStat.h

@@ -68,9 +68,9 @@ public:
 
   int calculateAvgUploadSpeed();
 
-  void updateDownloadLength(size_t bytes);
+  void updateDownload(size_t bytes);
 
-  void updateUploadLength(size_t bytes);
+  void updateUpload(size_t bytes);
 
   int getMaxDownloadSpeed() const;
 

test/DefaultBtAnnounceTest.cc

@@ -67,8 +67,8 @@ public:
       torrentAttrs->infoHash.assign(std::begin(infoHash), std::end(infoHash));
       dctx_->setAttribute(CTX_ATTR_BT, std::move(torrentAttrs));
     }
-    dctx_->getNetStat().updateDownloadLength(pieceLength * 5);
-    dctx_->getNetStat().updateUploadLength(pieceLength * 6);
+    dctx_->getNetStat().updateDownload(pieceLength * 5);
+    dctx_->getNetStat().updateUpload(pieceLength * 6);
     bittorrent::setStaticPeerId(peerId);
 
     pieceStorage_.reset(new MockPieceStorage());

test/DefaultBtProgressInfoFileTest.cc

@@ -214,7 +214,7 @@ void DefaultBtProgressInfoFileTest::testSave()
   initializeMembers(1_k, 80_k);
 
   dctx_->setBasePath(A2_TEST_OUT_DIR "/save-temp");
-  dctx_->getNetStat().updateUploadLength(768);
+  dctx_->getNetStat().updateUpload(768);
   btRuntime_->setUploadLengthAtStartup(256);
   bitfield_->setAllBit();
   bitfield_->unsetBit(79);

test/PeerSessionResourceTest.cc

@@ -117,8 +117,8 @@ void PeerSessionResourceTest::testUpdateDownloadLength()
   PeerSessionResource res(1_k, 1_m);
 
   CPPUNIT_ASSERT_EQUAL((int64_t)0LL, res.downloadLength());
-  res.updateDownloadLength(100);
-  res.updateDownloadLength(200);
+  res.updateDownload(100);
+  res.updateDownload(200);
   CPPUNIT_ASSERT_EQUAL((int64_t)300LL, res.downloadLength());
 }