From faac03a9591db739418cd781dbffa75ff57ac917 Mon Sep 17 00:00:00 2001
From: Juraj Smiesko <juraj.smiesko@cern.ch>
Date: Thu, 4 Apr 2024 16:59:43 +0200
Subject: [PATCH] Limit number of seen events

---
 k4Gen/src/components/GenEventFilter.cpp | 17 +++++++++++++++++
 k4Gen/src/components/GenEventFilter.h   |  8 ++++++++
 2 files changed, 25 insertions(+)

diff --git a/k4Gen/src/components/GenEventFilter.cpp b/k4Gen/src/components/GenEventFilter.cpp
index c0ddc70..e7d2111 100644
--- a/k4Gen/src/components/GenEventFilter.cpp
+++ b/k4Gen/src/components/GenEventFilter.cpp
@@ -142,6 +142,23 @@ StatusCode GenEventFilter::initialize() {
 
 StatusCode GenEventFilter::execute(const EventContext& evtCtx) const {
   const edm4hep::MCParticleCollection* inParticles = m_inColl.get();
+
+  if (m_nEventsSeen == m_nEventsSeenMax) {
+    debug() << "Maximum number of " << m_nEventsSeenMax.value()
+            << " seen events reached!" << endmsg;
+
+    StatusCode sc = m_eventProcessor->stopRun();
+    if (sc.isFailure()) {
+      error() << "Error when attempting to stop event processing! Aborting..."
+              << endmsg;
+      return sc;
+    }
+
+    m_incidentSvc->fireIncident(Incident(name(), IncidentType::AbortEvent));
+
+    return StatusCode::SUCCESS;
+  }
+
   m_nEventsSeen++;
 
   if (!(*m_filterRulePtr)(inParticles)) {
diff --git a/k4Gen/src/components/GenEventFilter.h b/k4Gen/src/components/GenEventFilter.h
index 7db6e34..df88481 100644
--- a/k4Gen/src/components/GenEventFilter.h
+++ b/k4Gen/src/components/GenEventFilter.h
@@ -1,6 +1,9 @@
 #ifndef GENERATION_GENEVENTFILTER_H
 #define GENERATION_GENEVENTFILTER_H
 
+// std
+#include <limits>
+
 // Gaudi
 #include "GaudiKernel/Algorithm.h"
 class IProperty;
@@ -48,6 +51,11 @@ class GenEventFilter : public Gaudi::Algorithm {
   Gaudi::Property<std::string> m_filterRulePath{
       this, "filterRulePath", "", "Path to the filter rule file"};
 
+  /// Maximum number of events seen
+  Gaudi::Property<unsigned int> m_nEventsSeenMax{
+      this, "nEventsSeenMax", std::numeric_limits<unsigned int>::max(),
+      "Limit the number of events seen"};
+
   /// Targeted number of events.
   mutable std::atomic<unsigned int> m_nEventsTarget;
   /// Keep track of how many events were already accepted.