sched: implement datadist controller rpcs

Author: ironMann

src/DataDistControl/DataDistControl.proto

@@ -25,7 +25,8 @@ message PartitionInfo {
 }
 
 enum PartitionState {
-  __IGNORE = 0;
+  IGNORE__ = 0;
+
   PARTITION_UNKNOWN = 1;
   // Partition not known to DataDistControl. Initialize by calling PartitionInitialize()
 
@@ -43,7 +44,10 @@ enum PartitionState {
   // All components configured, ready for dataflow commands
 
   PARTITION_TERMINATING = 6;
-  // Partition is terminating. EPN-FLP connections are cleanly terminated.
+  // Partition is terminating. EPN-FLP connections will be cleanly closed.
+
+  PARTITION_TERMINATED = 7;
+  // Partition is terminated. TfScheduler does not accept further requests
 }
 
 

src/TfScheduler/TfSchedulerConnManager.cxx

@@ -116,14 +116,9 @@ void TfSchedulerConnManager::disconnectTfBuilder(const TfBuilderConfigStatus &pT
   pResponse.set_status(0);
   const std::string &lTfBuilderId = pTfBuilderStatus.info().process_id();
 
-  std::scoped_lock lLock(mStfSenderClientsLock);
-
-  deleteTfBuilderRpcClient(lTfBuilderId);
-
-  if (!stfSendersReady()) {
-    IDDLOG("TfBuilder Connection error: StfSenders not ready.");
-    pResponse.set_status(ERROR_STF_SENDERS_NOT_READY);
-    return;
+  {
+    std::scoped_lock lLock(mStfSenderClientsLock);
+    deleteTfBuilderRpcClient(lTfBuilderId);
   }
 
   TfBuilderEndpoint lParam;
@@ -137,31 +132,70 @@ void TfSchedulerConnManager::disconnectTfBuilder(const TfBuilderConfigStatus &pT
       continue; // not connected
     }
 
-    lParam.set_tf_builder_id(lTfBuilderId);
-    lParam.set_endpoint(lSocketInfo.endpoint());
+    { // lock clients
+      std::scoped_lock lLock(mStfSenderClientsLock);
+
+      if (mStfSenderRpcClients.count(lStfSenderId) == 0) {
+        WDDLOG("disconnectTfBuilder: Unknown StfSender. stfs_id={}", lStfSenderId);
+        continue;
+      }
 
-    if (mStfSenderRpcClients.count(lStfSenderId) == 0) {
-      WDDLOG("disconnectTfBuilder: Unknown StfSender. stfs_id={}", lStfSenderId);
-      continue;
+      lParam.set_tf_builder_id(lTfBuilderId);
+      lParam.set_endpoint(lSocketInfo.endpoint());
+      StatusResponse lResponse;
+
+      auto &lRpcClient = mStfSenderRpcClients[lSocketInfo.peer_id()];
+      if(!lRpcClient->DisconnectTfBuilderRequest(lParam, lResponse).ok()) {
+        EDDLOG("StfSender Connection error: gRPC error. stfs_id={} tfb_id={}", lStfSenderId, lTfBuilderId);
+        pResponse.set_status(ERROR_GRPC_STF_SENDER);
+        continue;
+      }
+      // check StfSender status
+      if (lResponse.status() != 0) {
+        EDDLOG("TfBuilder Connection error. stfs_id={} tfb_id={} response={}", lStfSenderId, lTfBuilderId, lResponse.status());
+        pResponse.set_status(ERROR_STF_SENDER_CONNECTING);
+        continue;
+      }
     }
+  }
+}
 
-    auto &lRpcClient = mStfSenderRpcClients[lSocketInfo.peer_id()];
+// Partition RPC: keep sending until all TfBuilders are gone
+bool TfSchedulerConnManager::requestTfBuildersTerminate() {
+  std::vector<std::string> lFailedRpcsForDeletion;
 
-    StatusResponse lResponse;
-    if(!lRpcClient->DisconnectTfBuilderRequest(lParam, lResponse).ok()) {
-      EDDLOG("TfBuilder Connection error: gRPC error. stfs_id={} tfb_id={}", lStfSenderId, lTfBuilderId);
-      pResponse.set_status(ERROR_GRPC_STF_SENDER);
-      break;
+  std::scoped_lock lLock(mStfSenderClientsLock);
+
+  for (auto &lTfBuilder : mTfBuilderRpcClients) {
+    if (!lTfBuilder.second.mClient->TerminatePartition()) {
+      lFailedRpcsForDeletion.emplace_back(lTfBuilder.first);
     }
+  }
 
-    // check StfSender status
-    if (lResponse.status() != 0) {
-      EDDLOG("TfBuilder Connection error. stfs_id={} tfb_id={} response={}",
-        lStfSenderId, lTfBuilderId, lResponse.status());
-      pResponse.set_status(ERROR_STF_SENDER_CONNECTING);
-      break;
+  for (const auto &lId : lFailedRpcsForDeletion) {
+    deleteTfBuilderRpcClient(lId);
+  }
+
+  return mTfBuilderRpcClients.size() == 0;
+}
+
+// Partition RPC: notify all StfSenders and remove rpc clients
+bool TfSchedulerConnManager::requestStfSendersTerminate() {
+  std::vector<std::string> lFailedRpcsForDeletion;
+
+  std::scoped_lock lLock(mStfSenderClientsLock);
+
+  for (auto &lStfSender : mStfSenderRpcClients) {
+    if (!lStfSender.second->TerminatePartition()) {
+      lFailedRpcsForDeletion.emplace_back(lStfSender.first);
     }
   }
+
+  for (const auto &lId : lFailedRpcsForDeletion) {
+    deleteTfBuilderRpcClient(lId);
+  }
+
+  return mTfBuilderRpcClients.size() == 0;
 }
 
 
@@ -260,10 +294,10 @@ void TfSchedulerConnManager::StfSenderMonitoringThread()
 
       WDDLOG_RL(1000, "Waiting for StfSenders. ready={} total={}", lNumStfSenders, mPartitionInfo.mStfSenderIdList.size());
       lSleep = 250ms;
+    } else {
+      mStfSenderState = STF_SENDER_STATE_OK;
     }
 
-    mStfSenderState = STF_SENDER_STATE_OK;
-
     // wait for drop futures
     {
       {

src/TfScheduler/TfSchedulerConnManager.h

@@ -62,7 +62,7 @@ class TfSchedulerConnManager
     using namespace std::chrono_literals;
 
     while (!mStfSenderRpcClients.start()) {
-      std::this_thread::sleep_for(1s);
+      return false; // we'll be called back
     }
 
     mRunning = true;
@@ -98,6 +98,10 @@ class TfSchedulerConnManager
 
   void StfSenderMonitoringThread();
 
+  /// Partition RPCs
+  bool requestTfBuildersTerminate();
+  bool requestStfSendersTerminate();
+
   /// External requests by TfBuilders
   void connectTfBuilder(const TfBuilderConfigStatus &pTfBuilderStatus, TfBuilderConnectionResponse &pResponse /*out*/);
   void disconnectTfBuilder(const TfBuilderConfigStatus &pTfBuilderStatus, StatusResponse &pResponse /*out*/);
@@ -127,7 +131,13 @@ class TfSchedulerConnManager
     return mStfSenderRpcClients.checkStfSenderRpcConn(lStfSenderId);
   }
 
+  void deleteStfSenderRpcClient(const std::string &pId)
+  {
+    mStfSenderRpcClients.remove(pId);
+  }
+
   StfSenderState getStfSenderState() const { return mStfSenderState; }
+  std::size_t getStfSenderCount() const { return mStfSenderRpcClients.size(); }
 
 private:
   /// Partition information

src/TfScheduler/TfSchedulerDevice.cxx

@@ -42,95 +42,67 @@ TfSchedulerDevice::~TfSchedulerDevice()
 void TfSchedulerDevice::InitTask()
 {
   DataDistLogger::SetThreadName("tfs-main");
-
-  // Discovery
-  mDiscoveryConfig = std::make_shared<ConsulTfSchedulerService>(ProcessType::TfSchedulerService, Config::getEndpointOption(*GetConfig()));
-
-  auto &lStatus = mDiscoveryConfig->status();
-
-  lStatus.mutable_info()->set_type(TfSchedulerService);
-  lStatus.mutable_info()->set_process_id(Config::getIdOption(*GetConfig()));
-  lStatus.mutable_info()->set_ip_address(Config::getNetworkIfAddressOption(*GetConfig()));
-
-  mDiscoveryConfig->write();
-
-  // start the service thread
-  mServiceThread = create_thread_member("sched_service", &TfSchedulerDevice::TfSchedulerServiceThread, this);
 }
 
 void TfSchedulerDevice::PreRun()
 {
-
+  mDiscoveryConfig = std::make_unique<ConsulTfSchedulerService>(
+    ProcessType::TfSchedulerService,
+    Config::getEndpointOption(*GetConfig())
+  );
 }
 
-void TfSchedulerDevice::ResetTask()
+void TfSchedulerDevice::PostRun()
 {
-  // stop the scheduler service
-  if (mServiceThread.joinable()) {
-    mServiceThread.join();
-  }
+  IDDLOG("Stopping the TfScheduler and exiting. partition_id={}", mPartitionId);
 
-  // stop all instances
-  for (auto &lInstIt : mSchedulerInstances) {
-    lInstIt.second->stop();
+  // delete everything
+  if (mSchedInstance) {
+    mSchedInstance->stop();
   }
+  mSchedInstance.reset();
+  mPartitionId.clear();
 
-  mSchedulerInstances.clear();
-
+  throw "Intentional exit";
+}
 
-  DDDLOG("ResetTask() done.");
+void TfSchedulerDevice::ResetTask()
+{
 }
 
 bool TfSchedulerDevice::ConditionalRun()
 {
   // nothing to do here sleep for awhile
   std::this_thread::sleep_for(500ms);
 
-  // NOTE: Not using Run or ConditionalRun lets us put teardown in PostRun()
-  return true;
-}
-
-void TfSchedulerDevice::TfSchedulerServiceThread()
-{
-  // wait for the device to go into RUNNING state
-  WaitForRunningState();
-
-  while (IsRunningState()) {
-
+  if (!mSchedInstance) {
     // check for new requests
     PartitionRequest lNewPartitionRequest;
-    DDLOGF_RL(2000, DataDistSeverity::debug, "Checking for new partition creation requests.");
+    DDLOGF_RL(5000, DataDistSeverity::debug, "Checking for new partition creation requests.");
     if (mDiscoveryConfig->getNewPartitionRequest(lNewPartitionRequest)) {
       // new request
       IDDLOG("Request for starting a new partition. partition={}", lNewPartitionRequest.mPartitionId);
 
-      // check if we already have instance for the requested partition
-      if (mSchedulerInstances.count(lNewPartitionRequest.mPartitionId) == 0) {
-
-        // Create a new instance for the partition
-        auto [lNewInstIt, lEmplaced ] = mSchedulerInstances.emplace(
-          lNewPartitionRequest.mPartitionId,
-          std::make_unique<TfSchedulerInstanceHandler>(*this,
-            std::string("0"), // TODO: add multiple schedulers
-            lNewPartitionRequest
-          ) // value
-        );
-
-        if (lEmplaced) {
-          auto &lNewInstance = lNewInstIt->second;
-          lNewInstance->start();
-        }
-        IDDLOG("Created new scheduler instance. partition={}", lNewPartitionRequest.mPartitionId);
-        break; // Only service one partition per process
-      } else {
-        EDDLOG("Scheduler instance already exists. partition={}",lNewPartitionRequest.mPartitionId);
-      }
+      // Create a new instance for the partition
+      mPartitionId = lNewPartitionRequest.mPartitionId;
+      mSchedInstance = std::make_unique<TfSchedulerInstanceHandler>(*this,
+        std::string("0"), // TODO: add multiple schedulers
+        lNewPartitionRequest
+      );
+      mSchedInstance->start();
+
+      IDDLOG("Created new scheduler instance. partition={}", lNewPartitionRequest.mPartitionId);
     }
+  }
 
-    std::this_thread::sleep_for(1000ms);
+  if (mSchedInstance) {
+    if (mSchedInstance->isTerminated()) {
+      std::this_thread::sleep_for(2000ms);
+      return false; // -> PostRun() -> exit
+    }
   }
 
-  DDDLOG("Exiting TfSchedulerServiceThread.");
+  return true;
 }
 
 }

src/TfScheduler/TfSchedulerDevice.h

@@ -21,7 +21,6 @@
 #include <Utilities.h>
 
 #include <thread>
-#include <map>
 
 
 namespace o2
@@ -45,21 +44,14 @@ class TfSchedulerDevice : public DataDistDevice
 
  protected:
   void PreRun() final;
-  void PostRun() final { };
+  void PostRun() final;
   bool ConditionalRun() final;
 
-  /// Discovery configuration
-  std::shared_ptr<ConsulTfSchedulerService> mDiscoveryConfig;
-
-  /// Scheduler service thread
-  void TfSchedulerServiceThread();
-  std::thread mServiceThread;
-
   /// Scheduler Instances
-  // NOTE (unique_ptr): TfSchedulerInstance has threads that take *this.
-  //                    unique_ptr ensures that *this does not change due to container
-  std::map<std::string, std::unique_ptr<TfSchedulerInstanceHandler>> mSchedulerInstances;
+  std::string mPartitionId;
+  std::unique_ptr<TfSchedulerInstanceHandler> mSchedInstance;
 
+  std::unique_ptr<ConsulTfSchedulerService> mDiscoveryConfig;
 };
 }
 } /* namespace o2::DataDistribution */