PruneDataStoreModule.cc

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#include <framework/modules/core/PruneDataStoreModule.h>
#include <framework/datastore/DataStore.h>
#include <framework/datastore/StoreEntry.h>
#include <framework/logging/LogSystem.h>
#include <framework/logging/Logger.h>
 
using namespace std;
using namespace Belle2;
 
REG_MODULE(PruneDataStore);
 
PruneDataStoreModule::PruneDataStoreModule() :
  Module()
{
  setDescription(R"DOC(
    Clears the content of the DataStore while it keeps entries listed in the matchEntries option.
    The EventMetaData object will always be kept, as it is required by the framework to properly
    work with the DataStore.
 
    This logic can be inverted to remove only the entries matched by the regex
    in the matchEntries parameter by setting the parameter to
    keepMatchedEntries to False.
 
    Note:
        Also all Relations will be cleared if they are not matched by one entry
        in the matchEntries list. You have to ensure the objects referenced by
        kept relations are also matched by one entry in the keepEntries list so
        a relation does not point to nirvana.
  )DOC");
  addParam("matchEntries", m_matchEntries,
           "name of all DataStore entries (with regular expression syntax) to match. "
           "For example, you can use 'Raw.*' to match all Raw-Objects.",
           m_matchEntries);
  addParam("keepMatchedEntries", m_keepMatchedEntries,
           "If true, all entries matched by the regular expression are kept. "
           "If false, matched entries will be removed.",
           m_keepMatchedEntries);
  setPropertyFlags(c_ParallelProcessingCertified);
 
}
 
void PruneDataStoreModule::initialize()
{
  // prepare the regex_matchers, otherwise this nede to be done for each DataStore item
  for (auto& kEntry : m_matchEntries) {
    m_compiled_regex.push_back(compileAndCatch(kEntry));
  }
  // also get the regex for the implicit keeps
  for (auto& kEntry : m_keepEntriesImplicit) {
    m_compiled_regex_implicit.push_back(compileAndCatch(kEntry));
  }
  // the cached prune decisions of a previous run are not valid any more
  m_entriesToPrune.clear();
  m_entriesToKeep.clear();
  m_cacheValid = false;
}
 
std::regex PruneDataStoreModule::compileAndCatch(std::string& regexString) const
{
  try {
    return std::regex(regexString);
  } catch (const std::regex_error& e) {
    B2FATAL("Regex '" << regexString << "' cannot be compiled: " <<  e.what());
    // keep the compiler happy and return something
    return std::regex();
  }
}
 
void PruneDataStoreModule::updateCache(DataStore::StoreEntryMap& storemap)
{
  m_entriesToPrune.clear();
  m_entriesToKeep.clear();
 
  // iterate through all StoreEntries and check the Regex expression for each entry
  for (auto& datastore_item : storemap) {
    std::string const& datastore_key = datastore_item.first;
 
    // check if this entry is in our to keep list
    bool toKeep = !m_keepMatchedEntries;
    for (auto const& regx : m_compiled_regex) {
      if (std::regex_match(datastore_key, regx)) {
        toKeep = m_keepMatchedEntries;
      }
    }
 
    // check for implicit keeps
    for (auto const& regx : m_compiled_regex_implicit) {
      if (std::regex_match(datastore_key, regx)) {
        // always keep, no matter of the keepMatchedEntries configuration
        toKeep = true;
      }
    }
 
    if (toKeep) {
      m_entriesToKeep.push_back(&datastore_key);
    } else {
      m_entriesToPrune.push_back(&datastore_item.second);
    }
  }
 
  m_cacheValid = true;
}
 
void PruneDataStoreModule::event()
{
  auto& storemap = DataStore::Instance().getStoreEntryMap(DataStore::c_Event);
 
  // Matching the regular expressions against every DataStore key is expensive,
  // while the set of DataStore entries is frozen once the initialize() methods
  // of all modules have been called: everything which registers an entry
  // (DataStore::registerEntry, DataStore::copyEntriesTo) is only reached from
  // there. So the regular expressions are evaluated on the first event only.
  if (not m_cacheValid) {
    updateCache(storemap);
  }
 
  if (LogSystem::Instance().isLevelEnabled(LogConfig::c_Debug, 100, PACKAGENAME())) {
    for (auto const* datastore_key : m_entriesToKeep) {
      B2DEBUG(100,
              "StoreArray entry " << *datastore_key << " will be not pruned from the datastore");
    }
    for (auto const* entry : m_entriesToPrune) {
      B2DEBUG(100,
              "StoreArray entry " << entry->name << " will be pruned from the datastore");
    }
  }
 
  for (auto* entry : m_entriesToPrune) {
    // recreate to clear its content
    entry->invalidate();
  }
}