[QC-260] Support EndOfStream in QC (#280)

* Fix data producer cases when minSize == maxSize

* max number of messages parameter for data-producer

* Finish cycle upon an EndOfStream signal

Author: knopers8

Framework/include/QualityControl/DataProducer.h

@@ -26,20 +26,30 @@ namespace o2::quality_control::core
 /// \param minSize  Minimum size of a message in bytes
 /// \param maxSize  Maximum size of a message in bytes
 /// \param rate     How much messages to produce in one second
-/// \param fill     Should it fill messages with random data
+/// \param amount   How many messages should be produce in total (0 for inf). EndOfStream is sent at the end.
 /// \param index    SubSpecification of the data producer (useful when more than one needed)
+/// \param monitoringUrl Where monitoring metrics should be sent
+/// \param fill     Should it fill messages with random data
+///
 /// \return         A random data producer specification
-framework::DataProcessorSpec getDataProducerSpec(size_t minSize, size_t maxSize, double rate, bool fill = true, size_t index = 0, std::string monitoringUrl = "");
+framework::DataProcessorSpec
+  getDataProducerSpec(size_t minSize, size_t maxSize, double rate, uint64_t amount = 0, size_t index = 0,
+                      std::string monitoringUrl = "", bool fill = true);
 
 /// \brief Returns an algorithm generating random messages
 ///
 /// \param output   Origin, Description and SubSpecification of data to be produced
 /// \param minSize  Minimum size of a message in bytes
 /// \param maxSize  Maximum size of a message in bytes
 /// \param rate     How much messages to produce in one second
+/// \param amount   How many messages should be produce in total (0 for inf). EndOfStream is sent at the end.
+/// \param monitoringUrl Where monitoring metrics should be sent
 /// \param fill     Should it fill messages with random data
+///
 /// \return         A random data producer algorithm
-framework::AlgorithmSpec getDataProducerAlgorithm(framework::ConcreteDataMatcher output, size_t minSize, size_t maxSize, double rate, bool fill = true, std::string monitoringUrl = "");
+framework::AlgorithmSpec
+  getDataProducerAlgorithm(framework::ConcreteDataMatcher output, size_t minSize, size_t maxSize, double rate,
+                           uint64_t amount = 0, std::string monitoringUrl = "", bool fill = true);
 
 } // namespace o2::quality_control::core
 

Framework/include/QualityControl/TaskRunner.h

@@ -17,13 +17,12 @@
 #ifndef QC_CORE_TASKRUNNER_H
 #define QC_CORE_TASKRUNNER_H
 
-//#include <boost/accumulators/accumulators.hpp>
-//#include <boost/accumulators/statistics.hpp>
 // O2
 #include <Common/Timer.h>
 #include <Framework/Task.h>
 #include <Framework/DataProcessorSpec.h>
 #include <Framework/CompletionPolicy.h>
+#include <Framework/EndOfStreamContext.h>
 #include <Headers/DataHeader.h>
 // QC
 #include "QualityControl/TaskConfig.h"
@@ -98,6 +97,9 @@ class TaskRunner : public framework::Task
   /// \brief Unified DataDescription naming scheme for all tasks
   static header::DataDescription createTaskDataDescription(const std::string& taskName);
 
+  /// \brief Callback for CallbackService::Id::EndOfStream
+  void endOfStream(framework::EndOfStreamContext& eosContext) override;
+
  private:
   /// \brief Callback for CallbackService::Id::Start (DPL) a.k.a. RUN transition (FairMQ)
   void start();
@@ -132,6 +134,7 @@ class TaskRunner : public framework::Task
   framework::Options mOptions;
 
   bool mCycleOn = false;
+  bool mNoMoreCycles = false;
   int mCycleNumber = 0;
 
   // stats

Framework/src/DataProducer.cxx

@@ -13,10 +13,12 @@
 /// \author Piotr Konopka
 ///
 #include "QualityControl/DataProducer.h"
+#include "QualityControl/QcInfoLogger.h"
 
 #include <random>
 #include <Common/Timer.h>
 #include <Monitoring/MonitoringFactory.h>
+#include <Framework/ControlService.h>
 
 using namespace o2::framework;
 using namespace o2::monitoring;
@@ -27,28 +29,29 @@ using namespace AliceO2::Common;
 namespace o2::quality_control::core
 {
 
-DataProcessorSpec getDataProducerSpec(size_t minSize, size_t maxSize, double rate, bool fill, size_t index, std::string monitoringUrl)
+DataProcessorSpec getDataProducerSpec(size_t minSize, size_t maxSize, double rate, uint64_t amount, size_t index,
+                                      std::string monitoringUrl, bool fill)
 {
   return DataProcessorSpec{
     "producer-" + std::to_string(index),
     Inputs{},
     Outputs{
       { { "out" }, "TST", "RAWDATA", static_cast<SubSpec>(index) } },
-    getDataProducerAlgorithm({ "TST", "RAWDATA", static_cast<SubSpec>(index) }, minSize, maxSize, rate, fill, monitoringUrl)
+    getDataProducerAlgorithm({ "TST", "RAWDATA", static_cast<SubSpec>(index) }, minSize, maxSize, rate, amount,
+                             monitoringUrl, fill)
   };
 }
 
-framework::AlgorithmSpec
-  getDataProducerAlgorithm(ConcreteDataMatcher output, size_t minSize, size_t maxSize, double rate, bool fill,
-                           std::string monitoringUrl)
+AlgorithmSpec getDataProducerAlgorithm(ConcreteDataMatcher output, size_t minSize, size_t maxSize, double rate,
+                                       uint64_t amount, std::string monitoringUrl, bool fill)
 {
   return AlgorithmSpec{
     [=](InitContext&) {
       // this is the initialization code
       std::default_random_engine generator(time(nullptr));
       std::shared_ptr<Timer> timer = nullptr;
 
-      int messageCounter = 0;
+      uint64_t messageCounter = 0;
       std::shared_ptr<monitoring::Monitoring> collector;
       if (!monitoringUrl.empty()) {
         collector = MonitoringFactory::Get(monitoringUrl);
@@ -59,6 +62,14 @@ framework::AlgorithmSpec
       return [=](ProcessingContext& processingContext) mutable {
         // everything inside this lambda function is invoked in a loop, because it this Data Processor has no inputs
 
+        // checking if we have reached the maximum amount of messages
+        if (amount != 0 && messageCounter >= amount) {
+          ILOG(Info) << "Reached the maximum number of messages, requesting to quit the producer and sending an EndOfStream" << ENDM;
+          processingContext.services().get<ControlService>().endOfStream();
+          processingContext.services().get<ControlService>().readyToQuit(QuitRequest::Me);
+          return;
+        }
+
         // setting up the timer
         if (!timer) {
           timer = std::make_shared<Timer>();
@@ -72,17 +83,19 @@ framework::AlgorithmSpec
         timer->increment();
 
         // generating data
-        size_t length = minSize + (generator() % (maxSize - minSize));
+        size_t length = (minSize == maxSize) ? minSize : (minSize + (generator() % (maxSize - minSize)));
         auto data = processingContext.outputs().make<char>({ output.origin, output.description, output.subSpec },
                                                            length);
+        ++messageCounter;
         if (fill) {
           for (auto&& item : data) {
             item = static_cast<char>(generator());
           }
         }
 
+        // send metrics
         if (collector) {
-          collector->send({ messageCounter++, "Data_producer_" + std::to_string(output.subSpec) + "_message_" },
+          collector->send({ messageCounter, "Data_producer_" + std::to_string(output.subSpec) + "_message_" },
                           DerivedMetricMode::RATE);
         }
       };

Framework/src/TaskRunner.cxx

@@ -80,12 +80,16 @@ void TaskRunner::init(InitContext& iCtx)
   // init user's task
   mTask->loadCcdb(mTaskConfig.conditionUrl);
   mTask->initialize(iCtx);
+
+  mNoMoreCycles = false;
+  mCycleNumber = 0;
 }
 
 void TaskRunner::run(ProcessingContext& pCtx)
 {
-  if (mTaskConfig.maxNumberCycles >= 0 && mCycleNumber >= mTaskConfig.maxNumberCycles) {
-    LOG(INFO) << "The maximum number of cycles (" << mTaskConfig.maxNumberCycles << ") has been reached.";
+  if (mNoMoreCycles) {
+    ILOG(Info) << "The maximum number of cycles (" << mTaskConfig.maxNumberCycles << ") has been reached"
+               << " or the device has received an EndOfStream signal. Won't start a new cycle." << ENDM;
     return;
   }
 
@@ -107,6 +111,8 @@ void TaskRunner::run(ProcessingContext& pCtx)
     }
     if (mTaskConfig.maxNumberCycles < 0 || mCycleNumber < mTaskConfig.maxNumberCycles) {
       startCycle();
+    } else {
+      mNoMoreCycles = true;
     }
   }
 }
@@ -117,7 +123,6 @@ CompletionPolicy::CompletionOp TaskRunner::completionPolicyCallback(o2::framewor
   //  added, this will break.
   size_t dataInputsExpected = inputs.size() - 1;
   size_t dataInputsPresent = 0;
-  size_t allInputs = 0;
 
   CompletionPolicy::CompletionOp action = CompletionPolicy::CompletionOp::Wait;
 
@@ -138,7 +143,6 @@ CompletionPolicy::CompletionOp TaskRunner::completionPolicyCallback(o2::framewor
 
   LOG(DEBUG) << "Completion policy callback. "
              << "Total inputs possible: " << inputs.size()
-             << ", inputs present: " << allInputs
              << ", data inputs: " << dataInputsPresent
              << ", timer inputs: " << (action == CompletionPolicy::CompletionOp::Consume);
 
@@ -173,12 +177,20 @@ header::DataDescription TaskRunner::createTaskDataDescription(const std::string&
   return description;
 }
 
+void TaskRunner::endOfStream(framework::EndOfStreamContext& eosContext)
+{
+  ILOG(Info) << "Received an EndOfStream, finishing the current cycle" << ENDM;
+  finishCycle(eosContext.outputs());
+  mNoMoreCycles = true;
+}
+
 void TaskRunner::start()
 {
   startOfActivity();
 
-  if (mTaskConfig.maxNumberCycles >= 0 && mCycleNumber >= mTaskConfig.maxNumberCycles) {
-    LOG(INFO) << "The maximum number of cycles (" << mTaskConfig.maxNumberCycles << ") has been reached.";
+  if (mNoMoreCycles) {
+    ILOG(Info) << "The maximum number of cycles (" << mTaskConfig.maxNumberCycles << ") has been reached"
+               << " or the device has received an EndOfStream signal. Won't start a new cycle." << ENDM;
     return;
   }
 

Framework/src/runDataProducer.cxx

@@ -41,6 +41,8 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
     ConfigParamSpec{ "empty", VariantType::Bool, false, { "Don't fill messages with random data." } });
   workflowOptions.push_back(
     ConfigParamSpec{ "message-rate", VariantType::Double, 10.0, { "Rate of messages per second." } });
+  workflowOptions.push_back(
+    ConfigParamSpec{ "message-amount", VariantType::Int, 0, { "Amount of messages to be produced in total (0 for inf)." } });
   workflowOptions.push_back(
     ConfigParamSpec{ "producers", VariantType::Int, 1, { "Number of producers. Each will have unique SubSpec, counting from 0." } });
   workflowOptions.push_back(
@@ -58,12 +60,13 @@ WorkflowSpec defineDataProcessing(const ConfigContext& config)
   size_t maxSize = config.options().get<int>("max-size");
   bool fill = !config.options().get<bool>("empty");
   double rate = config.options().get<double>("message-rate");
+  uint64_t amount = config.options().get<int>("message-amount");
   size_t producers = config.options().get<int>("producers");
   std::string monitoringUrl = config.options().get<std::string>("monitoring-url");
 
   WorkflowSpec specs;
   for (size_t i = 0; i < producers; i++) {
-    specs.push_back(getDataProducerSpec(minSize, maxSize, rate, fill, i, monitoringUrl));
+    specs.push_back(getDataProducerSpec(minSize, maxSize, rate, amount, i, monitoringUrl, fill));
   }
   return specs;
 }