FWStorage/Catalog/src/FileLocator.cc - master - githubmirror/cmssw

FWStorage/Catalog/src/FileLocator.cc
167 строк6 KB
W. David Dagenhart
Reduce the memory usage of the InputFileCatalog
18 май 2026, 23:53
18 май 2026, 23:537e6433c
#include <algorithm>
#include <string_view>
#include <utility>

#include <boost/algorithm/string/replace.hpp>
#include <boost/property_tree/json_parser.hpp>

#include "FWCore/Utilities/interface/Exception.h"
#include "FWStorage/Catalog/interface/FileLocator.h"
#include "FWStorage/Catalog/interface/SiteLocalConfig.h"

namespace {

  std::string replaceWithRegexp(std::smatch const& matches, std::string const& outputFormat) {
    std::string result = outputFormat;

    for (size_t i = 1; i < matches.size(); ++i) {
      std::string_view const matchedString(matches[i].first, matches[i].second);
      if (!matchedString.empty()) {
        std::string placeholder = "$" + std::to_string(i);
        boost::algorithm::replace_all(result, placeholder, matchedString);
      }
    }
    return result;
  }

  constexpr char const* const kEmptyString = "";
  constexpr char const* const kLFNPrefix = "/store/";

}  // namespace

namespace pt = boost::property_tree;

namespace edm {

  FileLocator::FileLocator(CatalogAttributes const& catalogAttributes, std::filesystem::path const& filename_storage) {
    // now read json (usually from file named storage.json
    // whose location is from the SiteLocalConfig service)
    pt::ptree json;
    try {
      pt::read_json(filename_storage.string(), json);
    } catch (std::exception& e) {
      cms::Exception ex("FileCatalog");
      ex << "Can not open storage.json (" << filename_storage.string()
         << "). Check SITECONFIG_PATH and site-local-config.xml <data-access>";
      ex.addContext("edm::FileLocator:init()");
      throw ex;
    }

    // Find entry in storage.json that matches storageSite and volume from catalogAttributes.
    // Usually the catalogAttributes are from site-local-config.xml.
    auto found_site = std::find_if(json.begin(), json.end(), [&](pt::ptree::value_type const& site) {
      //get site name
      std::string siteName = site.second.get("site", kEmptyString);
      //get volume name
      std::string volName = site.second.get("volume", kEmptyString);
      return catalogAttributes.storageSite == siteName && catalogAttributes.volume == volName;
    });
    if (found_site == json.end()) {
      cms::Exception ex("FileCatalog");
      ex << "Can not find storage site \"" << catalogAttributes.storageSite << "\" and volume \""
         << catalogAttributes.volume
         << "\" in storage.json. Check site-local-config.xml <data-access> and storage.json";
      ex.addContext("edm::FileLocator:init()");
      throw ex;
    }

    // The matching top level entry should contain a list of protocols.
    // Find the protocol that matches the protocol from catalogAttributes.
    const pt::ptree& protocols = found_site->second.find("protocols")->second;
    auto found_protocol = std::find_if(protocols.begin(), protocols.end(), [&](pt::ptree::value_type const& protocol) {
      std::string protName = protocol.second.get("protocol", kEmptyString);
      return catalogAttributes.protocol == protName;
    });
    if (found_protocol == protocols.end()) {
      cms::Exception ex("FileCatalog");
      ex << "Can not find protocol \"" << catalogAttributes.protocol << "\" for the storage site \""
         << catalogAttributes.storageSite << "\" and volume \"" << catalogAttributes.volume
         << "\" in storage.json. Check site-local-config.xml <data-access> and storage.json";
      ex.addContext("edm::FileLocator:init()");
      throw ex;
    }

    m_protocol = found_protocol->second.get("protocol", kEmptyString);

    // Store prefixes and rules in m_directRules for all protocols
    // for this site and volume. We need to store for all protocols so
    // that "applyRules" can find them in the case that chaining is used.
    for (pt::ptree::value_type const& protocol : protocols) {
      std::string protName = protocol.second.get("protocol", kEmptyString);
      // if there is not a prefix, loop over rules
      std::string prefixTmp = protocol.second.get("prefix", kEmptyString);
      if (prefixTmp == kEmptyString) {
        const pt::ptree& rules = protocol.second.find("rules")->second;
        for (pt::ptree::value_type const& storageRule : rules) {
          parseRule(storageRule, protName, m_directRules);
        }
      }
      // If there is a prefix, convert it to a single rule and save it
      else {
        Rule rule;
        rule.pathMatch.assign("/?(.*)");
        rule.result = prefixTmp + "/$1";
        rule.chain = kEmptyString;
        m_directRules[protName].emplace_back(std::move(rule));
      }
    }
  }

  std::string FileLocator::pfn(std::string const& lfn) const {
    // check if lfn is an authentic LFN
    if (lfn.compare(0, 7, kLFNPrefix) != 0) {
      return "";
    }
    return applyRules(m_directRules, m_protocol, lfn);
  }

  void FileLocator::parseRule(pt::ptree::value_type const& storageRule,
                              std::string const& protocol,
                              ProtocolRules& rules) {
    if (storageRule.second.empty()) {
      throw cms::Exception("RucioFileCatalog", "edm::FileLocator::parseRule Malformed storage rule");
    }
    auto const pathMatchRegexp = storageRule.second.get<std::string>("lfn");
    auto const result = storageRule.second.get<std::string>("pfn");
    auto const chain = storageRule.second.get("chain", kEmptyString);
    Rule rule;
    rule.pathMatch.assign(pathMatchRegexp);
    rule.result = result;
    rule.chain = chain;
    rules[protocol].emplace_back(std::move(rule));
  }

  std::string FileLocator::applyRules(ProtocolRules const& protocolRules,
                                      std::string const& protocol,
                                      std::string name) const {
    ProtocolRules::const_iterator const rulesIterator = protocolRules.find(protocol);
    if (rulesIterator == protocolRules.end()) {
      return "";
    }

    Rules const& rules = rulesIterator->second;

    std::smatch nameMatches;

    // Look for a matching rule
    for (auto const& rule : rules) {
      if (!std::regex_match(name, nameMatches, rule.pathMatch)) {
        continue;
      }

      if (!rule.chain.empty()) {
        // Chain is the protocol in the recursive call to applyRules
        std::string const& protocol = rule.chain;
        name = applyRules(protocolRules, protocol, name);

        if (name.empty()) {
          return "";
        }
        std::regex_match(name, nameMatches, rule.pathMatch);
      }

      return replaceWithRegexp(nameMatches, rule.result);
    }
    return "";
  }
}  // namespace edm