Factor the primary request out of the sim worker kernel

run/O2SimDevice.cxx

@@ -19,7 +19,6 @@
 #include <fairlogger/Logger.h>
 #include <DetectorsBase/Stack.h>
 #include <DetectorsBase/VMCSeederService.h>
-#include <SimulationDataFormat/PrimaryChunk.h>
 #include <TRandom.h>
 #include <SimConfig/SimConfig.h>
 #include <cstring>
@@ -130,7 +129,6 @@ bool O2SimDevice::isWorkAvailable(fair::mq::Channel& statuschannel, int workerID
 
 bool O2SimDevice::Kernel(int workerID, fair::mq::Channel& requestchannel, fair::mq::Channel& dataoutchannel, fair::mq::Channel* statuschannel)
 {
-  static int counter = 0;
   bool reproducibleSim = true;
   if (getenv("O2_DISABLE_REPRODUCIBLE_SIM")) {
     reproducibleSim = false;
@@ -158,9 +156,6 @@ bool O2SimDevice::Kernel(int workerID, fair::mq::Channel& requestchannel, fair::
     focus_on_part = std::atoi(p.second.c_str());
   }
 
-  fair::mq::MessagePtr request(requestchannel.NewSimpleMessage(PrimaryChunkRequest{workerID, -1, counter++})); // <-- don't need content; channel means -> give primaries
-  fair::mq::Parts reply;
-
   mVMCApp->setSimDataChannel(&dataoutchannel);
 
   // we log info with workerID prepended
@@ -170,92 +165,107 @@ bool O2SimDevice::Kernel(int workerID, fair::mq::Channel& requestchannel, fair::
     return str.str();
   };
 
+  auto fetched = fetchChunk(workerID, requestchannel);
+  if (fetched.status == ChunkStatus::Stop) {
+    return false;
+  }
+  if (fetched.status != ChunkStatus::Work) {
+    // EndOfWork deliberately keeps asking: the server answers the next request as idle, which is the Stop
+    return true;
+  }
+  auto& chunk = fetched.chunk;
+
+  auto info = chunk->mSubEventInfo;
+  LOG(info) << workerStr() << " Processing " << chunk->mParticles.size() << " primary particles "
+            << "for event " << info.eventID << "/" << info.maxEvents << " "
+            << "part " << info.part << "/" << info.nparts;
+
+  if (eventselection == nullptr || (focus_on_event == info.eventID && focus_on_part == info.part)) {
+    mVMCApp->setPrimaries(chunk->mParticles);
+  } else {
+    // nothing to transport here
+    mVMCApp->setPrimaries(std::vector<TParticle>{});
+    LOG(info) << workerStr() << " This chunk will be skipped";
+  }
+
+  mVMCApp->setSubEventInfo(&info);
+
+  if (reproducibleSim) {
+    LOG(info) << workerStr() << " Setting seed for this sub-event to " << chunk->mSubEventInfo.seed;
+    gRandom->SetSeed(chunk->mSubEventInfo.seed);
+    o2::base::VMCSeederService::instance().setSeed();
+  }
+
+  // Process one event
+  auto& conf = o2::conf::SimConfig::Instance();
+  if (strcmp(conf.getMCEngine().c_str(), "TGeant4") == 0 || strcmp(conf.getMCEngine().c_str(), "O2TrivialMCEngine") == 0) {
+    // this is preferred and necessary for Geant4
+    // since repeated "ProcessRun" might have significant overheads
+    mVMC->ProcessEvent();
+  } else {
+    // for Geant3 calling ProcessEvent is not enough
+    // as some hooks are not called
+    mVMC->ProcessRun(1);
+  }
+
+  FairSystemInfo sysinfo;
+  LOG(info) << workerStr() << " TIME-STAMP " << mTimer.RealTime() << "\t";
+  mTimer.Continue();
+  LOG(info) << workerStr() << " MEM-STAMP " << sysinfo.GetCurrentMemory() / (1024. * 1024) << " "
+            << sysinfo.GetMaxMemory() << " MB\n";
+  return true;
+}
+
+O2SimDevice::FetchedChunk O2SimDevice::fetchChunk(int workerID, fair::mq::Channel& requestchannel)
+{
+  static int counter = 0;
+  fair::mq::MessagePtr request(requestchannel.NewSimpleMessage(PrimaryChunkRequest{workerID, -1, counter++})); // <-- don't need content; channel means -> give primaries
+  fair::mq::Parts reply;
+
+  auto workerStr = [workerID]() {
+    std::stringstream str;
+    str << "[W" << workerID << "]";
+    return str.str();
+  };
+
   doLogInfo(workerID, "Requesting work chunk");
   int timeoutinMS = 2000;
   auto sendcode = requestchannel.Send(request, timeoutinMS);
-  if (sendcode > 0) {
-    doLogInfo(workerID, "Waiting for answer");
-    // asking for primary generation
-
-    auto code = requestchannel.Receive(reply);
-    if (code > 0) {
-      doLogInfo(workerID, "Primary chunk received");
-      auto rawmessage = std::move(reply.At(0));
-      auto header = *(o2::PrimaryChunkAnswer*)(rawmessage->GetData());
-      if (!header.payload_attached) {
-        doLogInfo(workerID, "No payload; Server in stage " + std::string(PrimStateToString[(int)header.serverstate]));
-        // if no payload attached we inspect the server state, to see what to do
-        if (header.serverstate == O2PrimaryServerState::Initializing || header.serverstate == O2PrimaryServerState::WaitingEvent) {
-          sleep(1); // back-off and retry
-          return true;
-        }
-        // we need to decide what to do when the server is idle ---> if this happens immediately after a new batch request it means that the server might just lag a bit behind
-        return false;
-      } else {
-        auto payload = std::move(reply.At(1));
-        // wrap incoming bytes as a TMessageWrapper which offers "adoption" of a buffer
-        auto message = new TMessageWrapper(payload->GetData(), payload->GetSize());
-        auto chunk = static_cast<o2::data::PrimaryChunk*>(message->ReadObjectAny(message->GetClass()));
-
-        bool goon = true;
-        // no particles and eventID == -1 --> indication for no more work
-        if (chunk->mParticles.size() == 0 && chunk->mSubEventInfo.eventID == -1) {
-          doLogInfo(workerID, "No particles in reply : quitting kernel");
-          goon = false;
-        }
-
-        if (goon) {
-
-          auto info = chunk->mSubEventInfo;
-          LOG(info) << workerStr() << " Processing " << chunk->mParticles.size() << " primary particles "
-                    << "for event " << info.eventID << "/" << info.maxEvents << " "
-                    << "part " << info.part << "/" << info.nparts;
-
-          if (eventselection == nullptr || (focus_on_event == info.eventID && focus_on_part == info.part)) {
-            mVMCApp->setPrimaries(chunk->mParticles);
-          } else {
-            // nothing to transport here
-            mVMCApp->setPrimaries(std::vector<TParticle>{});
-            LOG(info) << workerStr() << " This chunk will be skipped";
-          }
-
-          mVMCApp->setSubEventInfo(&info);
-
-          if (reproducibleSim) {
-            LOG(info) << workerStr() << " Setting seed for this sub-event to " << chunk->mSubEventInfo.seed;
-            gRandom->SetSeed(chunk->mSubEventInfo.seed);
-            o2::base::VMCSeederService::instance().setSeed();
-          }
-
-          // Process one event
-          auto& conf = o2::conf::SimConfig::Instance();
-          if (strcmp(conf.getMCEngine().c_str(), "TGeant4") == 0 || strcmp(conf.getMCEngine().c_str(), "O2TrivialMCEngine") == 0) {
-            // this is preferred and necessary for Geant4
-            // since repeated "ProcessRun" might have significant overheads
-            mVMC->ProcessEvent();
-          } else {
-            // for Geant3 calling ProcessEvent is not enough
-            // as some hooks are not called
-            mVMC->ProcessRun(1);
-          }
-
-          FairSystemInfo sysinfo;
-          LOG(info) << workerStr() << " TIME-STAMP " << mTimer.RealTime() << "\t";
-          mTimer.Continue();
-          LOG(info) << workerStr() << " MEM-STAMP " << sysinfo.GetCurrentMemory() / (1024. * 1024) << " "
-                    << sysinfo.GetMaxMemory() << " MB\n";
-        }
-        delete message;
-        delete chunk;
-      }
-    } else {
-      LOG(info) << workerStr() << " No primary answer received from server (within timeout). Return code " << code;
-    }
-  } else {
+  if (sendcode <= 0) {
     LOG(info) << workerStr() << " Requesting work from server not possible. Return code " << sendcode;
-    return false;
+    return {ChunkStatus::Stop, nullptr};
   }
-  return true;
+  doLogInfo(workerID, "Waiting for answer");
+  // asking for primary generation
+  auto code = requestchannel.Receive(reply);
+  if (code <= 0) {
+    LOG(info) << workerStr() << " No primary answer received from server (within timeout). Return code " << code;
+    return {ChunkStatus::Retry, nullptr};
+  }
+  doLogInfo(workerID, "Primary chunk received");
+  auto rawmessage = std::move(reply.At(0));
+  auto header = *(o2::PrimaryChunkAnswer*)(rawmessage->GetData());
+  if (!header.payload_attached) {
+    doLogInfo(workerID, "No payload; Server in stage " + std::string(PrimStateToString[(int)header.serverstate]));
+    // if no payload attached we inspect the server state, to see what to do
+    if (header.serverstate == O2PrimaryServerState::Initializing || header.serverstate == O2PrimaryServerState::WaitingEvent) {
+      sleep(1); // back-off and retry
+      return {ChunkStatus::Retry, nullptr};
+    }
+    // we need to decide what to do when the server is idle ---> if this happens immediately after a new batch request it means that the server might just lag a bit behind
+    return {ChunkStatus::Stop, nullptr};
+  }
+  auto payload = std::move(reply.At(1));
+  // wrap incoming bytes as a TMessageWrapper which offers "adoption" of a buffer
+  TMessageWrapper message(payload->GetData(), payload->GetSize());
+  std::unique_ptr<o2::data::PrimaryChunk> chunk(static_cast<o2::data::PrimaryChunk*>(message.ReadObjectAny(message.GetClass())));
+
+  // no particles and eventID == -1 --> indication for no more work
+  if (chunk->mParticles.size() == 0 && chunk->mSubEventInfo.eventID == -1) {
+    doLogInfo(workerID, "No particles in reply : quitting kernel");
+    return {ChunkStatus::EndOfWork, nullptr};
+  }
+  return {ChunkStatus::Work, std::move(chunk)};
 }
 
 bool O2SimDevice::ConditionalRun()

run/O2SimDevice.h

@@ -20,6 +20,7 @@
 #include <FairRunSim.h>
 #include <TStopwatch.h>
 #include "PrimaryServerState.h"
+#include <SimulationDataFormat/PrimaryChunk.h>
 
 class TVirtualMC;
 
@@ -60,6 +61,21 @@ class O2SimDevice final : public fair::mq::Device
 
   bool Kernel(int workerID, fair::mq::Channel& requestchannel, fair::mq::Channel& dataoutchannel, fair::mq::Channel* statuschannel = nullptr);
 
+  /// outcome of asking the primary server for work
+  enum class ChunkStatus {
+    Work,      ///< a chunk to transport
+    Retry,     ///< nothing yet (server busy or no answer); ask again
+    EndOfWork, ///< the server sent the end-of-work marker
+    Stop       ///< the request could not be sent, or the server has nothing more to give; stop asking
+  };
+  struct FetchedChunk {
+    ChunkStatus status;
+    std::unique_ptr<o2::data::PrimaryChunk> chunk; ///< set for ChunkStatus::Work
+  };
+
+  /// requests one chunk of primaries from the server on the given channel
+  static FetchedChunk fetchChunk(int workerID, fair::mq::Channel& requestchannel);
+
  protected:
   /// Overloads the ConditionalRun() method of fair::mq::Device
   bool ConditionalRun() final;