Merge branch 'master' into random-webseeding

1 files changed, 53 insertions, 33 deletions

diff --git a/src/DHTInteractionCommand.cc b/src/DHTInteractionCommand.cc
index fb1e35b0..ca6554c5 100644
--- a/src/DHTInteractionCommand.cc
+++ b/src/DHTInteractionCommand.cc
@@ -54,35 +54,40 @@
 #include "UDPTrackerRequest.h"
 #include "fmt.h"
 #include "wallclock.h"
+#include "TrackerWatcherCommand.h"
 
 namespace aria2 {
 
 // TODO This name of this command is misleading, because now it also
 // handles UDP trackers as well as DHT.
 DHTInteractionCommand::DHTInteractionCommand(cuid_t cuid, DownloadEngine* e)
-  : Command{cuid},
-    e_{e},
-    dispatcher_{nullptr},
-    receiver_{nullptr},
-    taskQueue_{nullptr}
-{}
+    : Command{cuid},
+      e_{e},
+      dispatcher_{nullptr},
+      receiver_{nullptr},
+      taskQueue_{nullptr}
+{
+  setStatusRealtime();
+}
 
 DHTInteractionCommand::~DHTInteractionCommand()
 {
   disableReadCheckSocket(readCheckSocket_);
 }
 
-void DHTInteractionCommand::setReadCheckSocket(const std::shared_ptr<SocketCore>& socket)
+void DHTInteractionCommand::setReadCheckSocket(
+    const std::shared_ptr<SocketCore>& socket)
 {
   readCheckSocket_ = socket;
-  if(socket) {
+  if (socket) {
     e_->addSocketForReadCheck(socket, this);
   }
 }
 
-void DHTInteractionCommand::disableReadCheckSocket(const std::shared_ptr<SocketCore>& socket)
+void DHTInteractionCommand::disableReadCheckSocket(
+    const std::shared_ptr<SocketCore>& socket)
 {
-  if(socket) {
+  if (socket) {
     e_->deleteSocketForReadCheck(socket, this);
   }
 }
@@ -91,11 +96,14 @@ bool DHTInteractionCommand::execute()
 {
   // We need to keep this command alive while TrackerWatcherCommand
   // needs this.
-  if(e_->getRequestGroupMan()->downloadFinished() ||
-     (e_->isHaltRequested() && udpTrackerClient_->getNumWatchers() == 0)) {
+  if (e_->getRequestGroupMan()->downloadFinished() ||
+      (e_->isHaltRequested() && udpTrackerClient_->getNumWatchers() == 0)) {
+    A2_LOG_DEBUG("DHTInteractionCommand exiting");
     return true;
-  } else if(e_->isForceHaltRequested()) {
+  }
+  else if (e_->isForceHaltRequested()) {
     udpTrackerClient_->failAll();
+    A2_LOG_DEBUG("DHTInteractionCommand exiting");
     return true;
   }
 
@@ -105,51 +113,63 @@ bool DHTInteractionCommand::execute()
   uint16_t remotePort;
   std::array<unsigned char, 64_k> data;
   try {
-    while(1) {
+    while (1) {
       ssize_t length = connection_->receiveMessage(data.data(), data.size(),
                                                    remoteAddr, remotePort);
-      if(length <= 0) {
+      if (length <= 0) {
         break;
       }
-      if(data[0] == 'd') {
+      if (data[0] == 'd') {
         // udp tracker response does not start with 'd', so assume
         // this message belongs to DHT. nothrow.
         receiver_->receiveMessage(remoteAddr, remotePort, data.data(), length);
-      } else {
+      }
+      else {
         // this may be udp tracker response. nothrow.
-        udpTrackerClient_->receiveReply(data.data(), length, remoteAddr,
-                                        remotePort, global::wallclock());
+        std::shared_ptr<UDPTrackerRequest> req;
+        if (udpTrackerClient_->receiveReply(req, data.data(), length,
+                                            remoteAddr, remotePort,
+                                            global::wallclock()) == 0) {
+          if (req->action == UDPT_ACT_ANNOUNCE) {
+            auto c = static_cast<TrackerWatcherCommand*>(req->user_data);
+            if (c) {
+              c->setStatus(Command::STATUS_ONESHOT_REALTIME);
+              e_->setNoWait(true);
+            }
+          }
+        }
       }
     }
-  } catch(RecoverableException& e) {
+  }
+  catch (RecoverableException& e) {
     A2_LOG_INFO_EX("Exception thrown while receiving UDP message.", e);
   }
   receiver_->handleTimeout();
   udpTrackerClient_->handleTimeout(global::wallclock());
   dispatcher_->sendMessages();
-  while(!udpTrackerClient_->getPendingRequests().empty()) {
+  while (!udpTrackerClient_->getPendingRequests().empty()) {
     // no throw
-    ssize_t length = udpTrackerClient_->createRequest(data.data(), data.size(),
-                                                      remoteAddr, remotePort,
-                                                      global::wallclock());
-    if(length == -1) {
+    ssize_t length = udpTrackerClient_->createRequest(
+        data.data(), data.size(), remoteAddr, remotePort, global::wallclock());
+    if (length == -1) {
       break;
     }
     try {
       // throw
       connection_->sendMessage(data.data(), length, remoteAddr, remotePort);
       udpTrackerClient_->requestSent(global::wallclock());
-    } catch(RecoverableException& e) {
+    }
+    catch (RecoverableException& e) {
       A2_LOG_INFO_EX("Exception thrown while sending UDP tracker request.", e);
       udpTrackerClient_->requestFail(UDPT_ERR_NETWORK);
     }
   }
-  e_->addCommand(std::unique_ptr<Command>(this));
+  e_->addRoutineCommand(std::unique_ptr<Command>(this));
   return false;
 }
 
-void DHTInteractionCommand::setMessageDispatcher
-(DHTMessageDispatcher* dispatcher)
+void DHTInteractionCommand::setMessageDispatcher(
+    DHTMessageDispatcher* dispatcher)
 {
   dispatcher_ = dispatcher;
 }
@@ -164,14 +184,14 @@ void DHTInteractionCommand::setTaskQueue(DHTTaskQueue* taskQueue)
   taskQueue_ = taskQueue;
 }
 
-void DHTInteractionCommand::setConnection
-(std::unique_ptr<DHTConnection> connection)
+void DHTInteractionCommand::setConnection(
+    std::unique_ptr<DHTConnection> connection)
 {
   connection_ = std::move(connection);
 }
 
-void DHTInteractionCommand::setUDPTrackerClient
-(const std::shared_ptr<UDPTrackerClient>& udpTrackerClient)
+void DHTInteractionCommand::setUDPTrackerClient(
+    const std::shared_ptr<UDPTrackerClient>& udpTrackerClient)
 {
   udpTrackerClient_ = udpTrackerClient;
 }