2323#include < string>
2424#include < fairmq/Channel.h>
2525#include < mutex>
26+ #include < vector>
2627#include < fairlogger/Logger.h>
2728#include < fairmq/Parts.h>
2829#include < fairmq/TransportFactory.h>
@@ -274,36 +275,53 @@ bool waitForControlInput(int workerID)
274275// and send their hits to the merger, as the forked worker processes do otherwise.
275276std::atomic<bool > gServerDone {false };
276277std::atomic<int > gNextWorkerID {0 };
277- // the transport threads create their channels at the same moment; the shared transport factory is not made for that
278- std::mutex gChannelSetupMutex ;
279278
280- std::unique_ptr<fair::mq::Channel> makeWorkerChannel (char const * name, char const * type, std::shared_ptr<fair::mq::TransportFactory> const & factory, std::string const & address)
279+ // The channels of the transport threads, kept here so that they can be closed before the process exits
280+ struct MTChannels {
281+ std::shared_ptr<fair::mq::TransportFactory> factory;
282+ std::mutex setupMutex; // the threads create their channels at the same moment; the factory is not made for that
283+ // one request channel for all threads, used by one thread at a time: with one per thread, only one thread got work
284+ std::mutex requestMutex;
285+ std::unique_ptr<fair::mq::Channel> request;
286+ std::vector<std::unique_ptr<fair::mq::Channel>> data; // one per thread
287+ };
288+ MTChannels gMTChannels ;
289+
290+ fair::mq::Channel* makeWorkerChannel (std::unique_ptr<fair::mq::Channel>& owner, char const * name, char const * type, std::string const & address)
281291{
282- std::lock_guard<std::mutex> lock (gChannelSetupMutex );
283- auto channel = std::make_unique<fair::mq::Channel>(name, type, factory);
284- if (!channel ->Connect (address) || !channel ->Validate ()) {
292+ std::lock_guard<std::mutex> lock (gMTChannels . setupMutex );
293+ owner = std::make_unique<fair::mq::Channel>(name, type, gMTChannels . factory );
294+ if (!owner ->Connect (address) || !owner ->Validate ()) {
285295 LOG (fatal) << " Could not connect channel " << name << " of a transport thread to " << address;
286296 }
287- return channel;
297+ return owner.get ();
298+ }
299+
300+ // Closes the channels of the transport threads; with linger -1 this waits until the merger has taken all queued hits
301+ void closeMTChannels ()
302+ {
303+ auto & hooks = o2::base::MTWorkerHooks::instance ();
304+ hooks.fetchChunk = nullptr ;
305+ hooks.dataChannel = nullptr ;
306+ gMTChannels .request .reset ();
307+ gMTChannels .data .clear ();
308+ gMTChannels .factory .reset ();
288309}
289310
290311// The Geant4 worker threads already exist when this is called; they use these hooks from the first real ProcessRun.
291312void installMTWorkerHooks (std::string const & serveraddress, std::string const & mergeraddress)
292313{
293- auto factory = fair::mq::TransportFactory::CreateTransportFactory (" zeromq" );
314+ gMTChannels . factory = fair::mq::TransportFactory::CreateTransportFactory (" zeromq" );
294315 auto & hooks = o2::base::MTWorkerHooks::instance ();
295- hooks.fetchChunk = [factory, serveraddress]() -> std::unique_ptr<o2::data::PrimaryChunk> {
316+ hooks.fetchChunk = [serveraddress]() -> std::unique_ptr<o2::data::PrimaryChunk> {
296317 thread_local int workerID = gNextWorkerID ++;
297- // one request channel for all threads, used by one thread at a time
298- static std::mutex fetchMutex;
299- static std::unique_ptr<fair::mq::Channel> channel;
300- std::lock_guard<std::mutex> lock (fetchMutex);
301- if (!channel) {
302- channel = makeWorkerChannel (" primary-get" , " req" , factory, serveraddress);
318+ std::lock_guard<std::mutex> lock (gMTChannels .requestMutex );
319+ if (!gMTChannels .request ) {
320+ makeWorkerChannel (gMTChannels .request , " primary-get" , " req" , serveraddress);
303321 }
304322 using Status = o2::devices::O2SimDevice::ChunkStatus;
305323 while (true ) {
306- auto fetched = o2::devices::O2SimDevice::fetchChunk (workerID, *channel );
324+ auto fetched = o2::devices::O2SimDevice::fetchChunk (workerID, *gMTChannels . request );
307325 if (fetched.status == Status::Work) {
308326 return std::move (fetched.chunk );
309327 }
@@ -314,12 +332,16 @@ void installMTWorkerHooks(std::string const& serveraddress, std::string const& m
314332 // Retry and EndOfWork: ask again, as a worker process does; the server answers as idle once done
315333 }
316334 };
317- hooks.dataChannel = [factory, mergeraddress]() -> fair::mq::Channel* {
318- thread_local std::unique_ptr< fair::mq::Channel> channel;
335+ hooks.dataChannel = [mergeraddress]() -> fair::mq::Channel* {
336+ thread_local fair::mq::Channel* channel = nullptr ;
319337 if (!channel) {
320- channel = makeWorkerChannel (" simdata" , " push" , factory, mergeraddress);
338+ std::unique_ptr<fair::mq::Channel> owner;
339+ channel = makeWorkerChannel (owner, " simdata" , " push" , mergeraddress);
340+ channel->GetSocket ().SetLinger (-1 ); // keep queued hits until the merger has taken them
341+ std::lock_guard<std::mutex> lock (gMTChannels .setupMutex );
342+ gMTChannels .data .push_back (std::move (owner));
321343 }
322- return channel. get () ;
344+ return channel;
323345 };
324346}
325347
@@ -341,8 +363,7 @@ void installMTWorkerHooks(std::string const& serveraddress, std::string const& m
341363 }).detach ();
342364 o2::simpubsub::publishMessage (pushchannel, o2::simpubsub::simStatusString (" WORKER0" , " STATUS" , " SETUP COMPLETED" ));
343365
344- // a run ends when every thread has been told that there is no more work; ProcessRun is given a large
345- // but finite number of events and is repeated if the server still has work when it is exhausted
366+ // a run ends once every thread was told there is no more work; it is repeated if the events ran out first
346367 constexpr int kEventBlock = 100000 ;
347368 auto vmc = TVirtualMC::GetMC ();
348369 if (auto & beforeTransport = o2::base::MTWorkerHooks::instance ().beforeTransport ) {
@@ -368,7 +389,7 @@ void installMTWorkerHooks(std::string const& serveraddress, std::string const& m
368389 }
369390 }
370391 o2::SimSetup::shutdown ();
371- sleep ( 10 ); // give the merger time to receive the last messages
392+ closeMTChannels ();
372393 _exit (0 );
373394}
374395
0 commit comments