DPL Analysis: send CCDB tables together with the objects they refer to (#16079)

Author: ktf

Framework/CCDBSupport/src/AnalysisCCDBHelpers.cxx

@@ -11,6 +11,7 @@
 
 #include "AnalysisCCDBHelpers.h"
 #include "CCDBFetcherHelper.h"
+#include "Framework/ArrowContext.h"
 #include "Framework/ArrowTypes.h"
 #include "Framework/DataProcessingStats.h"
 #include "Framework/DeviceSpec.h"
@@ -25,6 +26,8 @@
 #include "Framework/ConfigParamsHelper.h"
 #include <fairmq/Version.h>
 #include <arrow/array/builder_binary.h>
+#include <arrow/io/memory.h>
+#include <arrow/ipc/writer.h>
 #include <arrow/type.h>
 #include <arrow/type_fwd.h>
 #include <arrow/util/key_value_metadata.h>
@@ -161,7 +164,7 @@ AlgorithmSpec AnalysisCCDBHelpers::fetchFromCCDB(ConfigContext const& /*ctx*/)
     std::unordered_map<std::string, int> bindings;
     fillValidRoutes(*helper, spec.outputs, bindings);
 
-    return adaptStateless([schemas, bindings, helper, allbuilders, pathTables](InputRecord& inputs, DataTakingContext& dtc, DataAllocator& allocator, TimingInfo& timingInfo, DataProcessingStats& stats) {
+    return adaptStateless([schemas, bindings, helper, allbuilders, pathTables](InputRecord& inputs, DataTakingContext& dtc, DataAllocator& allocator, ArrowContext& arrowContext, TimingInfo& timingInfo, DataProcessingStats& stats) {
       O2_SIGNPOST_ID_GENERATE(sid, ccdb);
       O2_SIGNPOST_START(ccdb, sid, "fetchFromAnalysisCCDB", "Fetching CCDB objects for analysis%" PRIu64, (uint64_t)timingInfo.timeslice);
       std::ranges::for_each(allbuilders, [](auto& builder) { builder.second->Reset(); });
@@ -232,6 +235,7 @@ AlgorithmSpec AnalysisCCDBHelpers::fetchFromCCDB(ConfigContext const& /*ctx*/)
         }
 
         std::vector<DataAllocator::CacheId> lastIds(numBuilders, DataAllocator::CacheId{.value = -1, .handle = -1, .segment = -1});
+        std::vector<DataAllocator::CacheId> referenced;
 
         // Rows sharing a uniformity value resolve to the same objects, so the query is
         // issued once per distinct value and the resulting handles are repeated for the
@@ -321,7 +325,7 @@ AlgorithmSpec AnalysisCCDBHelpers::fetchFromCCDB(ConfigContext const& /*ctx*/)
               auto& lastId = lastIds[bi];
               if (response.id.value != lastId.value) {
                 lastId.value = response.id.value;
-                allocator.adoptFromCache(output, response.id, header::gSerializationMethodCCDB);
+                referenced.push_back(response.id);
               }
               result &= builder.second->Append();
               auto* value_builder = dynamic_cast<arrow::Int64Builder*>(builder.second->value_builder());
@@ -338,8 +342,25 @@ AlgorithmSpec AnalysisCCDBHelpers::fetchFromCCDB(ConfigContext const& /*ctx*/)
         }
         arrow::ArrayVector arrays;
         std::ranges::for_each(builders, [&arrays](auto& builder) { arrays.push_back(*builder.second->Finish()); });
-        auto outTable = arrow::Table::Make(schema, arrays);
-        allocator.adopt(output, outTable);
+        // The table is sent by the message backend, followed by the objects it refers to,
+        // so that they all leave in one multipart and the consumer relays them into the
+        // same slot. Through the arrow backend, the table would leave on its own first.
+        auto stream = arrow::io::BufferOutputStream::Create().ValueOrDie();
+        auto writer = arrow::ipc::MakeStreamWriter(stream, schema).ValueOrDie();
+        if (!writer->WriteRecordBatch(*arrow::RecordBatch::Make(schema, reserveSize, arrays)).ok()) {
+          LOGP(fatal, "Unable to serialise the table of {}", outBinding);
+        }
+        auto buffer = stream->Finish().ValueOrDie();
+        auto payload = allocator.makeVector<char>(output);
+        payload.assign(buffer->data(), buffer->data() + buffer->size());
+        allocator.adoptContainer(output, std::move(payload), DataAllocator::CacheStrategy::Never, header::gSerializationMethodArrow);
+        // Accounted for like the tables of the arrow backend, for the rate limiting.
+        allocator.findMessageHeader(output)->serialization = header::gSerializationMethodArrow;
+        arrowContext.updateBytesSent(buffer->size());
+        arrowContext.updateMessagesSent(1);
+        for (auto const& id : referenced) {
+          allocator.adoptFromCache(output, id, header::gSerializationMethodCCDB);
+        }
       }
 
       stats.updateStats({(int)ProcessingStatsId::CCDB_CACHE_FETCHED_BYTES, DataProcessingStats::Op::Set, (int64_t)helper->totalFetchedBytes});

Framework/CCDBSupport/src/CCDBFetcherHelper.cxx

@@ -256,7 +256,7 @@ auto CCDBFetcherHelper::populateCacheWith(std::shared_ptr<CCDBFetcherHelper> con
         helper->totalFetchedBytes += size;
         helper->totalRequestedBytes += size;
         api.appendFlatHeader(v, headers);
-        auto cacheId = CCDBHelpers::adoptAndReplaceCachedMessage(allocator, helper->mapURL2DPLCache, path, output, std::move(v), header::gSerializationMethodCCDB);
+        auto cacheId = CCDBHelpers::adoptAndReplaceCachedMessage(allocator, helper->mapURL2DPLCache, path, output, std::move(v), header::gSerializationMethodCCDB, DataAllocator::CacheStrategy::CacheOnly);
         helper->mapURL2DPLCache[path] = cacheId;
         responses.emplace_back(Response{.id = cacheId, .size = size, .request = nullptr});
         O2_SIGNPOST_EVENT_EMIT(ccdb, sid, "populateCacheWith", "Caching %{public}s for %{public}s (DPL id %" PRIu64 ", size %zu)", path.data(), headers["ETag"].data(), cacheId.value, size);
@@ -274,7 +274,7 @@ auto CCDBFetcherHelper::populateCacheWith(std::shared_ptr<CCDBFetcherHelper> con
         helper->totalFetchedBytes += size;
         helper->totalRequestedBytes += size;
         api.appendFlatHeader(v, headers);
-        auto cacheId = CCDBHelpers::adoptAndReplaceCachedMessage(allocator, helper->mapURL2DPLCache, path, output, std::move(v), header::gSerializationMethodCCDB);
+        auto cacheId = CCDBHelpers::adoptAndReplaceCachedMessage(allocator, helper->mapURL2DPLCache, path, output, std::move(v), header::gSerializationMethodCCDB, DataAllocator::CacheStrategy::CacheOnly);
         helper->mapURL2DPLCache[path] = cacheId;
         responses.emplace_back(Response{.id = cacheId, .size = size, .request = nullptr});
         O2_SIGNPOST_EVENT_EMIT(ccdb, sid, "populateCacheWith", "Caching %{public}s for %{public}s (DPL id %" PRIu64 ")", path.data(), headers["ETag"].data(), cacheId.value);

Framework/CCDBSupport/src/CCDBHelpers.cxx

@@ -255,10 +255,11 @@ DataAllocator::CacheId CCDBHelpers::adoptAndReplaceCachedMessage(
   std::string const& path,
   Output const& output,
   o2::pmr::vector<char>&& v,
-  o2::header::SerializationMethod method)
+  o2::header::SerializationMethod method,
+  DataAllocator::CacheStrategy strategy)
 {
   auto oldIt = cache.find(path);
-  auto cacheId = allocator.adoptContainer(output, std::move(v), DataAllocator::CacheStrategy::Always, method);
+  auto cacheId = allocator.adoptContainer(output, std::move(v), strategy, method);
   if (oldIt != cache.end()) {
     allocator.pruneFromCache(oldIt->second);
   }

Framework/CCDBSupport/src/CCDBHelpers.h

@@ -39,14 +39,16 @@ struct CCDBHelpers {
   /// @a output DPL Output matcher
   /// @a v freshly-fetched CCDB payload; consumed by the call, leaving @a v empty
   /// @a method serialization-method tag written into the message header
+  /// @a strategy CacheOnly to cache the message without sending it
   /// @return the new CacheId; the caller must record it in its map
   static DataAllocator::CacheId adoptAndReplaceCachedMessage(
     DataAllocator& allocator,
     std::unordered_map<std::string, DataAllocator::CacheId> const& cache,
     std::string const& path,
     Output const& output,
     o2::pmr::vector<char>&& v,
-    o2::header::SerializationMethod method);
+    o2::header::SerializationMethod method,
+    DataAllocator::CacheStrategy strategy = DataAllocator::CacheStrategy::Always);
 };
 
 } // namespace o2::framework

Framework/Core/include/Framework/DataAllocator.h

@@ -528,7 +528,9 @@ class DataAllocator
 
   enum struct CacheStrategy : int {
     Never = 0,
-    Always = 1
+    Always = 1,
+    /// Cache the message without sending it. Use adoptFromCache to send it.
+    CacheOnly = 2
   };
 
   template <typename ContainerT>
@@ -628,7 +630,7 @@ DataAllocator::CacheId DataAllocator::adoptContainer(const Output& spec, Contain
   );
 
   CacheId cacheId{0, 0, 0}; //
-  if (cache == CacheStrategy::Always) {
+  if (cache != CacheStrategy::Never) {
     // The message will be shallow cloned in the cache. Since the
     // clone is indistinguishable from the original, we can keep sending
     // the original.
@@ -637,6 +639,9 @@ DataAllocator::CacheId DataAllocator::adoptContainer(const Output& spec, Contain
     cacheId.handle = meta.fHandle;
     cacheId.segment = meta.fSegmentId;
   }
+  if (cache == CacheStrategy::CacheOnly) {
+    return cacheId;
+  }
 
   context.add<MessageContext::TrivialObject>(std::move(headerMessage), std::move(payloadMessage), routeIndex);
   return cacheId;

Framework/Core/test/test_DataRelayer.cxx

@@ -1105,6 +1105,91 @@ TEST_CASE("DataRelayer")
     checkContents();
   }
 
+  // The analysis CCDB fetcher sends a CCDB table followed by the objects it
+  // refers to, all for the same route. The consumer must get them in the same
+  // slot, with the table as the first part, which holds when they are relayed
+  // from one multipart, i.e. without checking for completion in between.
+  SECTION("CCDBTableAndObjectsCompleteTogether")
+  {
+    InputSpec spec0{"bcs", "AOD", "BC"};
+    InputSpec spec1{"ccdb", "AOD", "GLOCCDBOBJ", 1};
+
+    std::vector<InputRoute> inputs = {
+      InputRoute{spec0, 0, "Fake0", 0},
+      InputRoute{spec1, 1, "Fake1", 0},
+    };
+
+    std::vector<InputChannelInfo> infos{1};
+    TimesliceIndex index{1, infos};
+    ref.registerService(ServiceRegistryHelpers::handleForService<TimesliceIndex>(&index));
+
+    auto policy = CompletionPolicyHelpers::consumeWhenAll();
+    DataRelayer relayer(policy, inputs, index, {registry}, -1);
+    relayer.setPipelineLength(4);
+
+    auto transport = fair::mq::TransportFactory::CreateTransportFactory("zeromq");
+    auto channelAlloc = o2::pmr::getTransportAllocator(transport.get());
+
+    auto relayOne = [&](o2::header::DataDescription description, uint32_t subSpec, o2::header::SerializationMethod method, size_t timeslice) {
+      DataHeader dh;
+      dh.dataOrigin = "AOD";
+      dh.dataDescription = description;
+      dh.subSpecification = subSpec;
+      dh.splitPayloadIndex = 0;
+      dh.splitPayloadParts = 1;
+      dh.payloadSerializationMethod = method;
+      std::array<fair::mq::MessagePtr, 2> msgs;
+      msgs[0] = o2::pmr::getMessage(Stack{channelAlloc, dh, DataProcessingHeader{timeslice, 1}});
+      msgs[1] = transport->CreateMessage(8);
+      DataRelayer::InputInfo info{0, 2, DataRelayer::InputType::Data, {ChannelIndex::INVALID}};
+      relayer.relay(msgs[0]->GetData(), msgs.data(), info, 2);
+    };
+    auto relayCCDB = [&](size_t timeslice) {
+      relayOne("GLOCCDBOBJ", 1, o2::header::gSerializationMethodArrow, timeslice);
+      relayOne("GLOCCDBOBJ", 1, o2::header::gSerializationMethodCCDB, timeslice);
+      relayOne("GLOCCDBOBJ", 1, o2::header::gSerializationMethodCCDB, timeslice);
+    };
+
+    // Table and objects in one multipart: one slot, with the table first.
+    relayOne("BC", 0, o2::header::gSerializationMethodArrow, 0);
+    relayCCDB(0);
+    std::vector<RecordAction> ready;
+    relayer.getReadyToProcess(ready);
+    REQUIRE(ready.size() == 1);
+    auto result = relayer.consumeAllInputsForTimeslice(ready[0].slot);
+    REQUIRE((result[1] | count_parts{}) == 3);
+    for (size_t p = 0; p < 3; ++p) {
+      auto* dh = o2::header::get<DataHeader*>((result[1] | get_header{p})->GetData());
+      REQUIRE(dh->payloadSerializationMethod == (p == 0 ? o2::header::gSerializationMethodArrow : o2::header::gSerializationMethodCCDB));
+    }
+
+    // The same, with the CCDB parts arriving before the other input.
+    relayCCDB(1);
+    ready.clear();
+    relayer.getReadyToProcess(ready);
+    REQUIRE(ready.empty());
+    relayOne("BC", 0, o2::header::gSerializationMethodArrow, 1);
+    relayer.getReadyToProcess(ready);
+    REQUIRE(ready.size() == 1);
+    result = relayer.consumeAllInputsForTimeslice(ready[0].slot);
+    REQUIRE((result[1] | count_parts{}) == 3);
+
+    // What happened when the table and the objects were sent separately: the
+    // slot completes on the table alone, and the objects open a new slot which
+    // can never complete.
+    relayOne("BC", 0, o2::header::gSerializationMethodArrow, 2);
+    relayOne("GLOCCDBOBJ", 1, o2::header::gSerializationMethodArrow, 2);
+    ready.clear();
+    relayer.getReadyToProcess(ready);
+    REQUIRE(ready.size() == 1);
+    result = relayer.consumeAllInputsForTimeslice(ready[0].slot);
+    REQUIRE((result[1] | count_parts{}) == 1);
+    relayOne("GLOCCDBOBJ", 1, o2::header::gSerializationMethodCCDB, 2);
+    ready.clear();
+    relayer.getReadyToProcess(ready);
+    REQUIRE(ready.empty());
+  }
+
   // An expiring input is materialised straight into the slot, so with one
   // shared buffer per slot it lands *after* whatever the other inputs already
   // hold -- the cells are then no longer in input order. Check that the data