2010-06-23 Tatsuhiro Tsujikawa <t-tujikawa@users.sourceforge.net>

	Request ut_metadata in end-game mode(in other words, more
	aggresively)
	* src/HandshakeExtensionMessage.cc
	* src/Peer.cc
	* src/Peer.h
	* src/PeerSessionResource.cc
	* src/PeerSessionResource.h
	* src/UTMetadataRequestFactory.cc
	* test/UTMetadataRequestFactoryTest.cc

6 files changed, 31 insertions, 11 deletions

diff --git a/src/HandshakeExtensionMessage.cc b/src/HandshakeExtensionMessage.cc
index 5a743de2..20c2e621 100644
--- a/src/HandshakeExtensionMessage.cc
+++ b/src/HandshakeExtensionMessage.cc
@@ -123,6 +123,11 @@ void HandshakeExtensionMessage::doReceivedAction()
       if(metadataSize_ != attrs->metadataSize) {
         throw DL_ABORT_EX("Wrong metadata_size. Which one is correct!?");
       }
+      if(!peer_->isSeeder()) {
+        peer_->reconfigureSessionResource(dctx_->getPieceLength(),
+                                          dctx_->getTotalLength());
+        peer_->setAllBitfield();
+      }
     } else {
       attrs->metadataSize = metadataSize_;
       dctx_->getFirstFileEntry()->setLength(metadataSize_);
@@ -131,7 +136,9 @@ void HandshakeExtensionMessage::doReceivedAction()
       
       SharedHandle<PieceStorage> pieceStorage =
         dctx_->getOwnerRequestGroup()->getPieceStorage();
-      pieceStorage->setEndGamePieceNum(0);
+      peer_->reconfigureSessionResource(dctx_->getPieceLength(),
+                                        dctx_->getTotalLength());
+      peer_->setAllBitfield();
     }
   } else if(attrs->metadata.empty()) {
     throw DL_ABORT_EX("Peer didn't provide metadata_size."
diff --git a/src/Peer.cc b/src/Peer.cc
index 8955d4e3..84fa887e 100644
--- a/src/Peer.cc
+++ b/src/Peer.cc
@@ -80,6 +80,12 @@ void Peer::allocateSessionResource(size_t pieceLength, uint64_t totalLength)
   res_->getPeerStat().downloadStart();
 }
 
+void Peer::reconfigureSessionResource(size_t pieceLength, uint64_t totalLength)
+{
+  assert(res_);
+  res_->reconfigure(pieceLength, totalLength);
+}
+
 void Peer::releaseSessionResource()
 {
   delete res_;
diff --git a/src/Peer.h b/src/Peer.h
index f8dd168e..5ab040ad 100644
--- a/src/Peer.h
+++ b/src/Peer.h
@@ -159,6 +159,8 @@ public:
 
   void allocateSessionResource(size_t pieceLength, uint64_t totalLength);
 
+  void reconfigureSessionResource(size_t pieceLength, uint64_t totalLength);
+
   void releaseSessionResource();
 
   const Timer& getFirstContactTime() const
diff --git a/src/PeerSessionResource.cc b/src/PeerSessionResource.cc
index daae8bf7..f6a8dd78 100644
--- a/src/PeerSessionResource.cc
+++ b/src/PeerSessionResource.cc
@@ -273,4 +273,10 @@ size_t PeerSessionResource::countOutstandingUpload() const
   return dispatcher_->countOutstandingUpload();
 }
 
+void PeerSessionResource::reconfigure(size_t pieceLength, uint64_t totalLenth)
+{
+  delete bitfieldMan_;
+  bitfieldMan_ = new BitfieldMan(pieceLength, totalLenth);
+}
+
 } // namespace aria2
diff --git a/src/PeerSessionResource.h b/src/PeerSessionResource.h
index 4fe28d16..c20460bd 100644
--- a/src/PeerSessionResource.h
+++ b/src/PeerSessionResource.h
@@ -155,6 +155,8 @@ public:
 
   size_t getBitfieldLength() const;
 
+  void reconfigure(size_t index, uint64_t totalLength);
+
   bool hasPiece(size_t index) const;
 
   void markSeeder();
diff --git a/src/UTMetadataRequestFactory.cc b/src/UTMetadataRequestFactory.cc
index 8f075882..6ebc0e9c 100644
--- a/src/UTMetadataRequestFactory.cc
+++ b/src/UTMetadataRequestFactory.cc
@@ -52,26 +52,23 @@ void UTMetadataRequestFactory::create
 (std::vector<SharedHandle<BtMessage> >& msgs, size_t num,
  const SharedHandle<PieceStorage>& pieceStorage)
 {
-  for(size_t index = 0, numPieces = dctx_->getNumPieces();
-      index < numPieces && num; ++index) {
-    SharedHandle<Piece> p = pieceStorage->getMissingPiece(index);
+  while(num) {
+    SharedHandle<Piece> p = pieceStorage->getMissingPiece(peer_);
     if(p.isNull()) {
       if(logger_->debug()) {
-        logger_->debug
-          ("ut_metadata piece %lu is used or already acquired.",
-           static_cast<unsigned long>(index));
+        logger_->debug("No ut_metadata piece is available to download.");
       }
-      continue;
+      break;
     }
     --num;
     if(logger_->debug()) {
       logger_->debug("Creating ut_metadata request index=%lu",
-                     static_cast<unsigned long>(index));
+                     static_cast<unsigned long>(p->getIndex()));
     }
     SharedHandle<UTMetadataRequestExtensionMessage> m
       (new UTMetadataRequestExtensionMessage
        (peer_->getExtensionMessageID("ut_metadata")));
-    m->setIndex(index);
+    m->setIndex(p->getIndex());
     m->setDownloadContext(dctx_);
     m->setBtMessageDispatcher(dispatcher_);
     m->setBtMessageFactory(messageFactory_);
@@ -79,7 +76,7 @@ void UTMetadataRequestFactory::create
     
     SharedHandle<BtMessage> msg = messageFactory_->createBtExtendedMessage(m);
     msgs.push_back(msg);
-    tracker_->add(index);
+    tracker_->add(p->getIndex());
   }
 }