Add MessageContext message discarding (#15662)

Author: autumn-mck

Framework/Core/include/Framework/MessageContext.h

@@ -53,6 +53,12 @@ struct Output;
 class MessageContext
 {
  public:
+  enum class DispatchState {
+    NotDispatched,
+    Dispatched,
+    Discarded,
+  };
+
   constexpr static ServiceKind service_kind = ServiceKind::Stream;
 
   // so far we are only using one instance per named channel
@@ -466,6 +472,11 @@ class MessageContext
   /// discarded.
   void clear();
 
+  /// Discard pending output messages without asserting that they were sent. This
+  /// is intended for exception teardown paths where normal post-processing will
+  /// not run.
+  void discard();
+
   FairMQDeviceProxy& proxy()
   {
     return mProxy;
@@ -490,16 +501,16 @@ class MessageContext
   o2::header::DataHeader* findMessageHeader(const Output& spec);
   o2::header::Stack* findMessageHeaderStack(const Output& spec);
   [[nodiscard]] int countDeviceOutputs(bool excludeDPLOrigin = false) const;
-  void fakeDispatch() { mDidDispatch = true; }
-  bool didDispatch() { return mDidDispatch; }
+  void fakeDispatch() { mDispatchState = DispatchState::Dispatched; }
+  [[nodiscard]] DispatchState dispatchState() const { return mDispatchState; }
   o2::framework::DataProcessingHeader* findMessageDataProcessingHeader(const Output& spec);
   std::pair<o2::header::DataHeader*, o2::framework::DataProcessingHeader*> findMessageHeaders(const Output& spec);
 
  private:
   FairMQDeviceProxy& mProxy;
   Messages mMessages;
   Messages mScheduledMessages;
-  bool mDidDispatch = false;
+  DispatchState mDispatchState = DispatchState::NotDispatched;
   DispatchControl mDispatchControl;
   /// Cached messages, in case we want to reuse them.
   std::unordered_map<int64_t, std::unique_ptr<fair::mq::Message>> mMessageCache;

Framework/Core/src/CommonServices.cxx

@@ -185,11 +185,17 @@ o2::framework::ServiceSpec CommonServices::streamContextSpec()
       auto& routes = processingContext.services().get<DeviceSpec const>().outputs;
       auto& timeslice = processingContext.services().get<TimingInfo>().timeslice;
       auto& messageContext = processingContext.services().get<MessageContext>();
+      auto dispatchState = messageContext.dispatchState();
+      O2_SIGNPOST_ID_FROM_POINTER(cid, stream_context, service);
+      // Do not report discarded messages as missing outputs.
+      if (dispatchState == MessageContext::DispatchState::Discarded) {
+        O2_SIGNPOST_EVENT_EMIT_ERROR(stream_context, cid, "postProcessingCallbacks", "Output messages discarded.");
+        return;
+      }
       // Check if we never created any data for this timeslice
-      // if we did not, but we still have didDispatched set to true
+      // if we did not, but messages were dispatched,
       // it means it was created out of band.
       bool userDidCreate = false;
-      O2_SIGNPOST_ID_FROM_POINTER(cid, stream_context, service);
       for (size_t ri = 0; ri < routes.size(); ++ri) {
         if (stream->routeCreated[ri] == true && stream->routeDPLCreated[ri] == false) {
           userDidCreate = true;
@@ -198,14 +204,14 @@ o2::framework::ServiceSpec CommonServices::streamContextSpec()
       }
       O2_SIGNPOST_EVENT_EMIT(stream_context, cid, "postProcessingCallbacks", "userDidCreate == %d && didDispatch == %d",
                              userDidCreate,
-                             messageContext.didDispatch());
-      if (userDidCreate == false && messageContext.didDispatch() == true) {
+                             dispatchState == MessageContext::DispatchState::Dispatched);
+      if (userDidCreate == false && dispatchState == MessageContext::DispatchState::Dispatched) {
         O2_SIGNPOST_EVENT_EMIT(stream_context, cid, "postProcessingCallbacks", "Data created out of band userDidCreate == %d && messageContext.didDispatch == %d",
                                userDidCreate,
-                               messageContext.didDispatch());
+                               dispatchState == MessageContext::DispatchState::Dispatched);
         return;
       }
-      if (userDidCreate == false && messageContext.didDispatch() == false) {
+      if (userDidCreate == false && dispatchState == MessageContext::DispatchState::NotDispatched) {
         O2_SIGNPOST_ID_FROM_POINTER(cid, stream_context, service);
         O2_SIGNPOST_EVENT_EMIT(stream_context, cid, "postProcessingCallbacks", "No data created.");
         return;

Framework/Core/src/MessageContext.cxx

@@ -84,7 +84,7 @@ int MessageContext::countDeviceOutputs(bool excludeDPLOrigin) const
 {
   // If we dispatched some messages before the end of the callback
   // we need to account for them as well.
-  int noutputs = mDidDispatch ? 1 : 0;
+  int noutputs = mDispatchState == DispatchState::Dispatched ? 1 : 0;
   constexpr o2::header::DataOrigin DataOriginDPL{"DPL"};
   for (auto it = mMessages.rbegin(); it != mMessages.rend(); ++it) {
     if (!excludeDPLOrigin || (*it)->header()->dataOrigin != DataOriginDPL) {
@@ -103,7 +103,16 @@ void MessageContext::clear()
 {
   // Verify that everything has been sent on clear.
   assert(std::all_of(mMessages.begin(), mMessages.end(), [](auto& m) { return m->empty(); }));
-  mDidDispatch = false;
+  assert(mScheduledMessages.empty());
+  mDispatchState = DispatchState::NotDispatched;
+  mScheduledMessages.clear();
+  mMessages.clear();
+}
+
+void MessageContext::discard()
+{
+  mDispatchState = DispatchState::Discarded;
+  mScheduledMessages.clear();
   mMessages.clear();
 }
 
@@ -157,7 +166,7 @@ void MessageContext::schedule(Messages::value_type&& message)
         }
         mDispatchControl.dispatch(std::move(parts), ChannelIndex{ci}, DefaultChannelIndex);
       }
-      mDidDispatch = mScheduledMessages.empty() == false;
+      mDispatchState = DispatchState::Dispatched;
       mScheduledMessages.clear();
     }
   }