From 651966b138e7e5d92a7e2228a1dd23767bbc7bbe Mon Sep 17 00:00:00 2001
From: Christopher Jones <chrisdjones15@gmail.com>
Date: Tue, 20 Jul 2021 10:05:24 -0500
Subject: [PATCH] Allow global::OutputModule to use a StreamCache

---
 .../interface/global/OutputModuleBase.h       |  4 +--
 .../global/outputmoduleAbilityToImplementor.h | 36 +++++++++++++++++++
 FWCore/Framework/src/WorkerT.cc               | 20 +++++++++--
 3 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/FWCore/Framework/interface/global/OutputModuleBase.h b/FWCore/Framework/interface/global/OutputModuleBase.h
index 26e51f244545a..db4eb031c42f0 100644
--- a/FWCore/Framework/interface/global/OutputModuleBase.h
+++ b/FWCore/Framework/interface/global/OutputModuleBase.h
@@ -124,8 +124,8 @@ namespace edm {
       void doBeginJob();
       void doEndJob();
 
-      void doBeginStream(StreamID id);
-      void doEndStream(StreamID id);
+      void doBeginStream(StreamID id) { doBeginStream_(id); }
+      void doEndStream(StreamID id) { doEndStream_(id); }
 
       bool doEvent(EventTransitionInfo const&, ActivityRegistry*, ModuleCallingContext const*);
       //For now this is a placeholder
diff --git a/FWCore/Framework/interface/global/outputmoduleAbilityToImplementor.h b/FWCore/Framework/interface/global/outputmoduleAbilityToImplementor.h
index 216b73d1ce8f0..78d37717c66a0 100644
--- a/FWCore/Framework/interface/global/outputmoduleAbilityToImplementor.h
+++ b/FWCore/Framework/interface/global/outputmoduleAbilityToImplementor.h
@@ -48,6 +48,37 @@ namespace edm {
         virtual void respondToCloseInputFile(FileBlock const&) = 0;
       };
 
+      template <typename T, typename C>
+      class StreamCacheHolder : public virtual T {
+      public:
+        explicit StreamCacheHolder(edm::ParameterSet const& iPSet) : OutputModuleBase(iPSet) {}
+        StreamCacheHolder(StreamCacheHolder<T, C> const&) = delete;
+        StreamCacheHolder<T, C>& operator=(StreamCacheHolder<T, C> const&) = delete;
+        ~StreamCacheHolder() override {
+          for (auto c : caches_) {
+            delete c;
+          }
+        }
+
+      protected:
+        C* streamCache(edm::StreamID iID) const { return caches_[iID.value()]; }
+
+      private:
+        void preallocStreams(unsigned int iNStreams) final { caches_.resize(iNStreams, static_cast<C*>(nullptr)); }
+        void doBeginStream_(StreamID id) final { caches_[id.value()] = beginStream(id).release(); }
+        void doEndStream_(StreamID id) final {
+          endStream(id);
+          delete caches_[id.value()];
+          caches_[id.value()] = nullptr;
+        }
+
+        virtual std::unique_ptr<C> beginStream(edm::StreamID) const = 0;
+        virtual void endStream(edm::StreamID) const {}
+
+        //When threaded we will have a container for N items whre N is # of streams
+        std::vector<C*> caches_;
+      };
+
       template <typename T, typename C>
       class RunCacheHolder : public virtual T {
       public:
@@ -117,6 +148,11 @@ namespace edm {
         typedef edm::global::outputmodule::LuminosityBlockCacheHolder<edm::global::OutputModuleBase, C> Type;
       };
 
+      template <typename C>
+      struct AbilityToImplementor<edm::StreamCache<C>> {
+        typedef edm::global::outputmodule::StreamCacheHolder<edm::global::OutputModuleBase, C> Type;
+      };
+
     }  // namespace outputmodule
   }    // namespace global
 }  // namespace edm
diff --git a/FWCore/Framework/src/WorkerT.cc b/FWCore/Framework/src/WorkerT.cc
index 3978441cce32c..03e5e25f2a0f4 100644
--- a/FWCore/Framework/src/WorkerT.cc
+++ b/FWCore/Framework/src/WorkerT.cc
@@ -76,6 +76,16 @@ namespace edm {
       static bool constexpr value = true;
     };
 
+    template <typename T>
+    struct has_only_stream_transition_functions {
+      static bool constexpr value = false;
+    };
+
+    template <>
+    struct has_only_stream_transition_functions<edm::global::OutputModuleBase> {
+      static bool constexpr value = true;
+    };
+
     struct DoNothing {
       template <typename... T>
       inline void operator()(const T&...) {}
@@ -456,7 +466,10 @@ namespace edm {
 
   template <typename T>
   inline void WorkerT<T>::implBeginStream(StreamID id) {
-    std::conditional_t<workerimpl::has_stream_functions<T>::value, workerimpl::DoBeginStream<T>, workerimpl::DoNothing>
+    std::conditional_t<workerimpl::has_stream_functions<T>::value or
+                           workerimpl::has_only_stream_transition_functions<T>::value,
+                       workerimpl::DoBeginStream<T>,
+                       workerimpl::DoNothing>
         might_call;
     might_call(this, id);
   }
@@ -469,7 +482,10 @@ namespace edm {
 
   template <typename T>
   inline void WorkerT<T>::implEndStream(StreamID id) {
-    std::conditional_t<workerimpl::has_stream_functions<T>::value, workerimpl::DoEndStream<T>, workerimpl::DoNothing>
+    std::conditional_t<workerimpl::has_stream_functions<T>::value or
+                           workerimpl::has_only_stream_transition_functions<T>::value,
+                       workerimpl::DoEndStream<T>,
+                       workerimpl::DoNothing>
         might_call;
     might_call(this, id);
   }