Flush out commQueue before stopping listener (#830)

* Flush out commQueue before stopping listener

Co-authored-by: Adit Ranadive <[email protected]>
Signed-off-by: Micha Dery <[email protected]>

Author: 2dm

src/core/agent_data.h

@@ -91,13 +91,14 @@ class nixlAgentData {
                            std::hash<std::string>, strEqual>     remoteSections;
 
         // State/methods for listener thread
-        nixlMDStreamListener               *listener;
-        std::map<nixl_socket_peer_t, int>  remoteSockets;
-        std::thread                        commThread;
-        std::vector<nixl_comm_req_t>       commQueue;
-        std::mutex                         commLock;
-        bool                               commThreadStop;
-        bool                               useEtcd;
+        nixlMDStreamListener *listener;
+        std::map<nixl_socket_peer_t, int> remoteSockets;
+        std::thread commThread;
+        std::vector<nixl_comm_req_t> commQueue;
+        std::mutex commLock;
+        std::atomic<bool> commThreadStop;
+        std::atomic<bool> agentShutdown;
+        bool useEtcd;
         std::unique_ptr<nixlTelemetry> telemetry_;
         std::exception_ptr commThreadException_;
 

src/core/nixl_agent.cpp

@@ -194,12 +194,18 @@ nixlAgent::nixlAgent(const std::string &name, const nixlAgentConfig &cfg) :
 
     if (data->useEtcd || cfg.useListenThread) {
         data->commThreadStop = false;
+        data->agentShutdown = false;
         data->commThread = std::thread(&nixlAgentData::commWorker, data.get(), std::ref(*this));
     }
 }
 
 nixlAgent::~nixlAgent() {
     if (data && (data->useEtcd || data->config.useListenThread)) {
+        data->agentShutdown = true;
+        while (!data->commQueue.empty()) {
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
+
         data->commThreadStop = true;
         if(data->commThread.joinable()) data->commThread.join();
 

src/core/nixl_listener.cpp

@@ -676,6 +676,10 @@ nixlAgentData::commWorkerInternal(nixlAgent *myAgent) {
 }
 
 void nixlAgentData::enqueueCommWork(nixl_comm_req_t request){
+    if (agentShutdown) {
+        NIXL_WARN << "Agent shutting down, unable to accept new requests";
+        return;
+    }
     std::lock_guard<std::mutex> lock(commLock);
     commQueue.push_back(std::move(request));
 }