Proper handling of GRP objects for Run3 runs

Author: victor-gonzalez

Common/TableProducer/trackextension.cxx

@@ -57,6 +57,7 @@ struct TrackExtension {
   Produces<aod::TracksExtended> extendedTrackQuantities;
   Service<o2::ccdb::BasicCCDBManager> ccdb;
 
+  o2::base::MatLayerCylSet* lut;
   int mRunNumber;
   float mMagField;
 
@@ -68,15 +69,10 @@ struct TrackExtension {
     ccdb->setCaching(true);
     ccdb->setLocalObjectValidityChecking();
 
-    auto lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->get<o2::base::MatLayerCylSet>(ccdbpath_lut));
+    lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->get<o2::base::MatLayerCylSet>(ccdbpath_lut));
 
     if (!o2::base::GeometryManager::isGeometryLoaded()) {
       ccdb->get<TGeoManager>(ccdbpath_geo);
-      /* it seems this is needed at this level for the material LUT to work properly */
-      /* but what happens if the run changes while doing the processing?             */
-      o2::parameters::GRPObject* grpo = ccdb->getForTimeStamp<o2::parameters::GRPObject>(ccdbpath_grp, analysis::trackextension::run3grp_timestamp);
-      o2::base::Propagator::initFieldFromGRP(grpo);
-      o2::base::Propagator::Instance()->setMatLUT(lut);
     }
     mRunNumber = 0;
     mMagField = 0.0;
@@ -118,7 +114,7 @@ struct TrackExtension {
   }
   PROCESS_SWITCH(TrackExtension, processRun2, "Process Run2 track extension task", true);
 
-  void processRun3(aod::FullTracks const& tracks, aod::Collisions const&)
+  void processRun3(aod::FullTracks const& tracks, aod::Collisions const&, aod::BCsWithTimestamps const&)
   {
     using namespace analysis::trackextension;
 
@@ -131,10 +127,26 @@ struct TrackExtension {
     /* when a new run (Run3) is processed                           */
     o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT;
 
+    int lastCollId = -1;
     for (auto& track : tracks) {
       std::array<float, 2> dca{1e10f, 1e10f};
       if (track.has_collision()) {
         if (track.trackType() == o2::aod::track::TrackTypeEnum::Track) {
+          if (lastCollId != track.collisionId()) {
+            auto bc = track.collision_as<aod::Collisions>().bc_as<aod::BCsWithTimestamps>();
+            if (mRunNumber != bc.runNumber()) {
+              auto grpo = ccdb->getForTimeStamp<o2::parameters::GRPObject>(ccdbpath_grp, bc.timestamp());
+              if (grpo != nullptr) {
+                o2::base::Propagator::initFieldFromGRP(grpo);
+                o2::base::Propagator::Instance()->setMatLUT(lut);
+                LOGF(info, "Setting magnetic field to %d kG for run %d from its GRP CCDB object", grpo->getNominalL3Field(), bc.runNumber());
+              } else {
+                LOGF(fatal, "GRP object is not available in CCDB for run=%d at timestamp=%llu", bc.runNumber(), bc.timestamp());
+              }
+              mRunNumber = bc.runNumber();
+            }
+            lastCollId = track.collisionId();
+          }
           auto trackPar = getTrackPar(track);
           auto const& collision = track.collision();
           gpu::gpustd::array<float, 2> dcaInfo;