Merge pull request #47029 from makortel/waitingTaskWithArenaHolder

Use `WaitingTaskHolder` to signal `doneWaiting()` instead of `WaitingTaskWithArenaHolder` in framework

FWCore/Concurrency/interface/WaitingTaskWithArenaHolder.h

@@ -40,9 +40,8 @@ namespace edm {
     // eventually intend for the task to be spawned.
     explicit WaitingTaskWithArenaHolder(oneapi::tbb::task_group&, WaitingTask* iTask);
 
-    // Takes ownership of the underlying task and uses the current
-    // arena.
-    explicit WaitingTaskWithArenaHolder(WaitingTaskHolder&& iTask);
+    // Captures the current arena.
+    explicit WaitingTaskWithArenaHolder(WaitingTaskHolder iTask);
 
     ~WaitingTaskWithArenaHolder();
 

FWCore/Concurrency/src/WaitingTaskWithArenaHolder.cc

@@ -25,7 +25,7 @@ namespace edm {
     m_task->increment_ref_count();
   }
 
-  WaitingTaskWithArenaHolder::WaitingTaskWithArenaHolder(WaitingTaskHolder&& iTask)
+  WaitingTaskWithArenaHolder::WaitingTaskWithArenaHolder(WaitingTaskHolder iTask)
       : m_task(iTask.release_no_decrement()),
         m_group(iTask.group()),
         m_arena(std::make_shared<oneapi::tbb::task_arena>(oneapi::tbb::task_arena::attach())) {}

FWCore/Framework/interface/CallbackExternalWork.h

@@ -114,10 +114,9 @@ namespace edm {
               WaitingTaskHolder produceTask =
                   Base::makeProduceTask(group, token, record, es, emitPostPrefetchingSignal, std::move(produceFunctor));
 
-              WaitingTaskWithArenaHolder waitingTaskWithArenaHolder =
-                  makeExceptionHandlerTask(std::move(produceTask), group);
+              WaitingTaskHolder waitingTaskHolder = makeExceptionHandlerTask(std::move(produceTask), group);
 
-              return makeAcquireTask(std::move(waitingTaskWithArenaHolder), group, token, record, es);
+              return makeAcquireTask(std::move(waitingTaskHolder), group, token, record, es);
             },
             std::move(iTask),
             iRecord,
@@ -134,15 +133,15 @@ namespace edm {
                            const TDecorator& iDec = TDecorator())
           : Base(iProd, std::move(iProduceFunc), iID, iDec), acquireFunction_(std::move(iAcquireFunc)) {}
 
-      WaitingTaskHolder makeAcquireTask(WaitingTaskWithArenaHolder waitingTaskWithArenaHolder,
+      WaitingTaskHolder makeAcquireTask(WaitingTaskHolder waitingTaskHolder,
                                         oneapi::tbb::task_group* group,
                                         ServiceWeakToken const& serviceToken,
                                         EventSetupRecordImpl const* record,
                                         EventSetupImpl const* eventSetupImpl) {
         return WaitingTaskHolder(
             *group,
             make_waiting_task(
-                [this, holder = std::move(waitingTaskWithArenaHolder), group, serviceToken, record, eventSetupImpl](
+                [this, holder = std::move(waitingTaskHolder), group, serviceToken, record, eventSetupImpl](
                     std::exception_ptr const* iException) mutable {
                   std::exception_ptr excptr;
                   if (iException) {
@@ -191,7 +190,7 @@ namespace edm {
                               ESModuleCallingContext const& context_;
                             };
                             EndGuard guard(record, context);
-                            acquireCache_ = (*acquireFunction_)(rec, holder);
+                            acquireCache_ = (*acquireFunction_)(rec, WaitingTaskWithArenaHolder(holder));
                           });
                         } catch (cms::Exception& iException) {
                           iException.addContext("Running acquire");
@@ -202,25 +201,24 @@ namespace edm {
                 }));
       }
 
-      WaitingTaskWithArenaHolder makeExceptionHandlerTask(WaitingTaskHolder produceTask,
-                                                          oneapi::tbb::task_group* group) {
-        return WaitingTaskWithArenaHolder(*group,
-                                          make_waiting_task([this, produceTask = std::move(produceTask)](
-                                                                std::exception_ptr const* iException) mutable {
-                                            std::exception_ptr excptr;
-                                            if (iException) {
-                                              excptr = *iException;
-                                            }
-                                            if (excptr) {
-                                              try {
-                                                convertException::wrap([excptr]() { std::rethrow_exception(excptr); });
-                                              } catch (cms::Exception& exception) {
-                                                exception.addContext("Running acquire and external work");
-                                                edm::exceptionContext(exception, Base::callingContext());
-                                                produceTask.doneWaiting(std::current_exception());
-                                              }
-                                            }
-                                          }));
+      WaitingTaskHolder makeExceptionHandlerTask(WaitingTaskHolder produceTask, oneapi::tbb::task_group* group) {
+        return WaitingTaskHolder(*group,
+                                 make_waiting_task([this, produceTask = std::move(produceTask)](
+                                                       std::exception_ptr const* iException) mutable {
+                                   std::exception_ptr excptr;
+                                   if (iException) {
+                                     excptr = *iException;
+                                   }
+                                   if (excptr) {
+                                     try {
+                                       convertException::wrap([excptr]() { std::rethrow_exception(excptr); });
+                                     } catch (cms::Exception& exception) {
+                                       exception.addContext("Running acquire and external work");
+                                       edm::exceptionContext(exception, Base::callingContext());
+                                       produceTask.doneWaiting(std::current_exception());
+                                     }
+                                   }
+                                 }));
       }
 
       std::shared_ptr<TAcquireFunc> acquireFunction_;

FWCore/Framework/interface/StreamSchedule.h

@@ -90,6 +90,8 @@
 #include "FWCore/Utilities/interface/propagate_const.h"
 #include "FWCore/Utilities/interface/thread_safety_macros.h"
 
+#include "oneapi/tbb/task_arena.h"
+
 #include <exception>
 #include <map>
 #include <memory>

FWCore/Framework/interface/global/EDFilterBase.h

@@ -37,7 +37,6 @@ namespace edm {
   class StreamID;
   class ActivityRegistry;
   class ThinnedAssociationsHelper;
-  class WaitingTaskWithArenaHolder;
   class EventForTransformer;
   class ServiceWeakToken;
 
@@ -75,10 +74,7 @@ namespace edm {
 
     private:
       bool doEvent(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*);
-      void doAcquire(EventTransitionInfo const&,
-                     ActivityRegistry*,
-                     ModuleCallingContext const*,
-                     WaitingTaskWithArenaHolder&);
+      void doAcquire(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*, WaitingTaskHolder&&);
       void doTransformAsync(WaitingTaskHolder iTask,
                             size_t iTransformIndex,
                             EventPrincipal const& iEvent,
@@ -169,7 +165,7 @@ namespace edm {
       virtual bool hasAcquire() const noexcept { return false; }
       bool hasAccumulator() const noexcept { return false; }
 
-      virtual void doAcquire_(StreamID, Event const&, edm::EventSetup const&, WaitingTaskWithArenaHolder&);
+      virtual void doAcquire_(StreamID, Event const&, edm::EventSetup const&, WaitingTaskHolder&&);
 
       void setModuleDescription(ModuleDescription const& md) { moduleDescription_ = md; }
       ModuleDescription moduleDescription_;

FWCore/Framework/interface/global/EDProducerBase.h

@@ -38,7 +38,6 @@ namespace edm {
   class GlobalSchedule;
   class ActivityRegistry;
   class ThinnedAssociationsHelper;
-  class WaitingTaskWithArenaHolder;
   class EventForTransformer;
   class ServiceWeakToken;
 
@@ -78,10 +77,7 @@ namespace edm {
 
     private:
       bool doEvent(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*);
-      void doAcquire(EventTransitionInfo const&,
-                     ActivityRegistry*,
-                     ModuleCallingContext const*,
-                     WaitingTaskWithArenaHolder&);
+      void doAcquire(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*, WaitingTaskHolder&&);
       void doTransformAsync(WaitingTaskHolder iTask,
                             size_t iTransformIndex,
                             EventPrincipal const& iEvent,
@@ -173,7 +169,7 @@ namespace edm {
 
       virtual bool hasAcquire() const noexcept { return false; }
 
-      virtual void doAcquire_(StreamID, Event const&, edm::EventSetup const&, WaitingTaskWithArenaHolder&);
+      virtual void doAcquire_(StreamID, Event const&, edm::EventSetup const&, WaitingTaskHolder&&);
 
       void setModuleDescription(ModuleDescription const& md) { moduleDescription_ = md; }
       ModuleDescription moduleDescription_;

FWCore/Framework/interface/global/OutputModuleBase.h

@@ -59,10 +59,7 @@ namespace edm {
       void doEndStream(StreamID id) { doEndStream_(id); }
 
       bool doEvent(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*);
-      void doAcquire(EventTransitionInfo const&,
-                     ActivityRegistry*,
-                     ModuleCallingContext const*,
-                     WaitingTaskWithArenaHolder&);
+      void doAcquire(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*, WaitingTaskHolder&&);
       //For now this is a placeholder
       /*virtual*/ void preActionBeforeRunEventAsync(WaitingTaskHolder iTask,
                                                     ModuleCallingContext const& iModuleCallingContext,
@@ -86,7 +83,7 @@ namespace edm {
       virtual void doEndRunSummary_(RunForOutput const&, EventSetup const&) {}
       virtual void doBeginLuminosityBlockSummary_(LuminosityBlockForOutput const&, EventSetup const&) {}
       virtual void doEndLuminosityBlockSummary_(LuminosityBlockForOutput const&, EventSetup const&) {}
-      virtual void doAcquire_(StreamID, EventForOutput const&, WaitingTaskWithArenaHolder&) {}
+      virtual void doAcquire_(StreamID, EventForOutput const&, WaitingTaskHolder&&) {}
 
       virtual bool hasAcquire() const noexcept { return false; }
     };

FWCore/Framework/interface/global/implementors.h

@@ -48,9 +48,9 @@
 // forward declarations
 namespace edm {
 
-  class WaitingTaskWithArenaHolder;
   class ServiceWeakToken;
   class ActivityRegistry;
+  class WaitingTaskWithArenaHolder;
 
   namespace global {
     namespace impl {
@@ -436,7 +436,7 @@ namespace edm {
       private:
         bool hasAcquire() const noexcept override { return true; }
 
-        void doAcquire_(StreamID, Event const&, edm::EventSetup const&, WaitingTaskWithArenaHolder&) final;
+        void doAcquire_(StreamID, Event const&, edm::EventSetup const&, WaitingTaskHolder&&) final;
 
         virtual void acquire(StreamID, Event const&, edm::EventSetup const&, WaitingTaskWithArenaHolder) const = 0;
       };

FWCore/Framework/interface/global/outputmoduleAbilityToImplementor.h

@@ -141,8 +141,8 @@ namespace edm {
       private:
         bool hasAcquire() const noexcept override { return true; }
 
-        void doAcquire_(StreamID id, EventForOutput const& event, WaitingTaskWithArenaHolder& holder) final {
-          acquire(id, event, holder);
+        void doAcquire_(StreamID id, EventForOutput const& event, WaitingTaskHolder&& holder) final {
+          acquire(id, event, WaitingTaskWithArenaHolder(std::move(holder)));
         }
 
         virtual void acquire(StreamID, EventForOutput const&, WaitingTaskWithArenaHolder) const = 0;

FWCore/Framework/interface/maker/Worker.h

@@ -32,7 +32,6 @@ the worker is reset().
 #include "FWCore/Framework/interface/ProductResolverIndexAndSkipBit.h"
 #include "FWCore/Concurrency/interface/WaitingTask.h"
 #include "FWCore/Concurrency/interface/WaitingTaskHolder.h"
-#include "FWCore/Concurrency/interface/WaitingTaskWithArenaHolder.h"
 #include "FWCore/Concurrency/interface/WaitingTaskList.h"
 #include "FWCore/MessageLogger/interface/MessageLogger.h"
 #include "FWCore/ServiceRegistry/interface/ActivityRegistry.h"
@@ -267,9 +266,7 @@ namespace edm {
     virtual void itemsToGetForSelection(std::vector<ProductResolverIndexAndSkipBit>&) const = 0;
     virtual bool implNeedToRunSelection() const noexcept = 0;
 
-    virtual void implDoAcquire(EventTransitionInfo const&,
-                               ModuleCallingContext const*,
-                               WaitingTaskWithArenaHolder&) = 0;
+    virtual void implDoAcquire(EventTransitionInfo const&, ModuleCallingContext const*, WaitingTaskHolder&&) = 0;
 
     virtual void implDoTransformAsync(WaitingTaskHolder,
                                       size_t iTransformIndex,
@@ -394,12 +391,14 @@ namespace edm {
                                                    ParentContext const&,
                                                    typename T::Context const*) noexcept;
 
-    void runAcquire(EventTransitionInfo const&, ParentContext const&, WaitingTaskWithArenaHolder&);
+    // runAcquire() must take a copy of WaitingTaskHolder
+    // see comment in runAcquireAfterAsyncPrefetch() definition
+    void runAcquire(EventTransitionInfo const&, ParentContext const&, WaitingTaskHolder);
 
     void runAcquireAfterAsyncPrefetch(std::exception_ptr,
                                       EventTransitionInfo const&,
                                       ParentContext const&,
-                                      WaitingTaskWithArenaHolder) noexcept;
+                                      WaitingTaskHolder) noexcept;
 
     std::exception_ptr handleExternalWorkException(std::exception_ptr iEPtr,
                                                    ParentContext const& parentContext) noexcept;
@@ -519,7 +518,7 @@ namespace edm {
                   typename T::TransitionInfoType const&,
                   ServiceToken const&,
                   ParentContext const&,
-                  WaitingTaskWithArenaHolder) noexcept {}
+                  WaitingTaskHolder) noexcept {}
       void execute() final {}
     };
 
@@ -530,7 +529,7 @@ namespace edm {
                   EventTransitionInfo const& eventTransitionInfo,
                   ServiceToken const& token,
                   ParentContext const& parentContext,
-                  WaitingTaskWithArenaHolder holder) noexcept
+                  WaitingTaskHolder holder) noexcept
           : m_worker(worker),
             m_eventTransitionInfo(eventTransitionInfo),
             m_parentContext(parentContext),
@@ -545,7 +544,7 @@ namespace edm {
         // to hold the exception_ptr
         std::exception_ptr temp_excptr;
         auto excptr = exceptionPtr();
-        // Caught exception is passed to Worker::runModuleAfterAsyncPrefetch(), which propagates it via WaitingTaskWithArenaHolder
+        // Caught exception is passed to Worker::runModuleAfterAsyncPrefetch(), which propagates it via WaitingTaskHolder
         CMS_SA_ALLOW try {
           //pre was called in prefetchAsync
           m_worker->emitPostModuleEventPrefetchingSignal();
@@ -563,12 +562,12 @@ namespace edm {
                         info = m_eventTransitionInfo,
                         parentContext = m_parentContext,
                         serviceToken = m_serviceToken,
-                        holder = m_holder]() {
+                        holder = std::move(m_holder)]() {
                          //Need to make the services available
                          ServiceRegistry::Operate operateRunAcquire(serviceToken.lock());
 
                          std::exception_ptr ptr;
-                         worker->runAcquireAfterAsyncPrefetch(ptr, info, parentContext, holder);
+                         worker->runAcquireAfterAsyncPrefetch(ptr, info, parentContext, std::move(holder));
                        });
             return;
           }
@@ -581,7 +580,7 @@ namespace edm {
       Worker* m_worker;
       EventTransitionInfo m_eventTransitionInfo;
       ParentContext const m_parentContext;
-      WaitingTaskWithArenaHolder m_holder;
+      WaitingTaskHolder m_holder;
       ServiceWeakToken m_serviceToken;
     };
 
@@ -1127,7 +1126,7 @@ namespace edm {
             auto* group = task.group();
             moduleTask = make_waiting_task(
                 [this, weakToken, transitionInfo, parentContext, ownRunTask, group](std::exception_ptr const* iExcept) {
-                  WaitingTaskWithArenaHolder runTaskHolder(
+                  WaitingTaskHolder runTaskHolder(
                       *group, new HandleExternalWorkExceptionTask(this, group, ownRunTask->release(), parentContext));
                   AcquireTask<T> t(this, transitionInfo, weakToken.lock(), parentContext, runTaskHolder);
                   t.execute();
@@ -1154,7 +1153,7 @@ namespace edm {
         auto group = task.group();
         if constexpr (T::isEvent_) {
           if (hasAcquire()) {
-            WaitingTaskWithArenaHolder runTaskHolder(
+            WaitingTaskHolder runTaskHolder(
                 *group, new HandleExternalWorkExceptionTask(this, group, moduleTask, parentContext));
             moduleTask = new AcquireTask<T>(this, transitionInfo, token, parentContext, std::move(runTaskHolder));
           }

FWCore/Framework/interface/maker/WorkerT.h

@@ -27,7 +27,6 @@ namespace edm {
   class ModuleProcessName;
   class ProductResolverIndexAndSkipBit;
   class ThinnedAssociationsHelper;
-  class WaitingTaskWithArenaHolder;
 
   template <typename T>
   class WorkerT : public Worker {
@@ -90,7 +89,7 @@ namespace edm {
     void itemsToGetForSelection(std::vector<ProductResolverIndexAndSkipBit>&) const final;
     bool implNeedToRunSelection() const noexcept final;
 
-    void implDoAcquire(EventTransitionInfo const&, ModuleCallingContext const*, WaitingTaskWithArenaHolder&) final;
+    void implDoAcquire(EventTransitionInfo const&, ModuleCallingContext const*, WaitingTaskHolder&&) final;
 
     size_t transformIndex(edm::BranchDescription const&) const noexcept final;
     void implDoTransformAsync(WaitingTaskHolder,

FWCore/Framework/interface/stream/EDFilter.h

@@ -28,8 +28,6 @@
 
 namespace edm {
 
-  class WaitingTaskWithArenaHolder;
-
   namespace stream {
 
     template <typename... T>
@@ -66,8 +64,8 @@ namespace edm {
       bool hasAbilityToProduceInEndLumis() const final { return HasAbilityToProduceInEndLumis<T...>::value; }
 
     private:
-      void doAcquire_(Event const& ev, EventSetup const& es, WaitingTaskWithArenaHolder& holder) final {
-        doAcquireIfNeeded(this, ev, es, holder);
+      void doAcquire_(Event const& ev, EventSetup const& es, WaitingTaskHolder&& holder) final {
+        doAcquireIfNeeded(this, ev, es, std::move(holder));
       }
     };
 

FWCore/Framework/interface/stream/EDFilterAdaptorBase.h

@@ -36,7 +36,6 @@ namespace edm {
 
   class ModuleCallingContext;
   class ActivityRegistry;
-  class WaitingTaskWithArenaHolder;
 
   namespace maker {
     template <typename T>
@@ -70,10 +69,7 @@ namespace edm {
     private:
       bool doEvent(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*);
 
-      void doAcquire(EventTransitionInfo const&,
-                     ActivityRegistry*,
-                     ModuleCallingContext const*,
-                     WaitingTaskWithArenaHolder&);
+      void doAcquire(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*, WaitingTaskHolder&&);
 
       //For now this is a placeholder
       /*virtual*/ void preActionBeforeRunEventAsync(WaitingTaskHolder,