fastHadd.cc

Go to the documentation of this file.

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <vector>
#include <set>
#include <string>
#include <iostream>
#include <memory>
#include <thread>
#include <mutex>
#include <list>
#include "DQMServices/Core/src/ROOTFilePB.pb.h"
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/gzip_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <TROOT.h>
#include <TFile.h>
#include <TBufferFile.h>
#include <TObject.h>
#include <TObjString.h>
#include <TH1.h>
#include <TProfile.h>
#include <TKey.h>
#include <TClass.h>

#include <sys/prctl.h>
#include <sys/wait.h>
#include <csignal>

#define DEBUG(x, msg) if (debug >= x) std::cout << "DEBUG: " << msg << std::flush

int debug = 0;

struct MicroME {
  MicroME(
          TObject *o,
          const std::string& dir,
          const std::string& obj,
          uint32_t flags = 0)
      : obj(o), dirname(dir), objname(obj), flags(flags) {}

  mutable TObject * obj;

  const std::string dirname;
  const std::string objname;

  uint32_t flags;

  bool operator<(const MicroME &rhs) const {
    const MicroME &lhs = *this;
    int diff = lhs.dirname.compare(rhs.dirname);
    return (diff < 0 ? true
          : diff == 0 ? lhs.objname < rhs.objname : false);
  };

  void add(TObject *obj_to_add) const {
      DEBUG(1, "Merging: " << obj->GetName() <<
        " << " << obj_to_add->GetName() << std::endl);

      if (dynamic_cast<TH1 *>(obj) && dynamic_cast<TH1 *>(obj_to_add)) {
        dynamic_cast<TH1 *>(obj)->Add(dynamic_cast<TH1 *>(obj_to_add));
      } else if (dynamic_cast<TObjString *>(obj) && dynamic_cast<TObjString *>(obj_to_add)) {


      } else {
        DEBUG(1, "Cannot merge (different types): " << obj->GetName() <<
          " << " << obj_to_add->GetName() << std::endl);
      }
  };

  const std::string fullname() const {
    return dirname + '/' + objname;
  };

};

using MEStore = std::set<MicroME>;

enum TaskType {
  TASK_ADD,
  TASK_DUMP,
  TASK_CONVERT,
  TASK_ENCODE
};

enum ErrType {
  ERR_BADCFG=1,
  ERR_NOFILE
};

using google::protobuf::io::FileInputStream;
using google::protobuf::io::FileOutputStream;
using google::protobuf::io::GzipInputStream;
using google::protobuf::io::GzipOutputStream;
using google::protobuf::io::CodedInputStream;
using google::protobuf::io::ArrayInputStream;


inline TObject * extractNextObject(TBufferFile &buf) {
  if (buf.Length() == buf.BufferSize())
    return nullptr;

  buf.InitMap();
  return reinterpret_cast<TObject *>(buf.ReadObjectAny(nullptr));
}

static void get_info(const dqmstorepb::ROOTFilePB::Histo &h,
                     std::string &dirname,
                     std::string &objname,
                     TObject ** obj) {

  size_t slash = h.full_pathname().rfind('/');
  size_t dirpos = (slash == std::string::npos ? 0 : slash);
  size_t namepos = (slash == std::string::npos ? 0 : slash+1);
  dirname.assign(h.full_pathname(), 0, dirpos);
  objname.assign(h.full_pathname(), namepos, std::string::npos);
  TBufferFile buf(TBufferFile::kRead, h.size(),
                  (void*)h.streamed_histo().data(),
                  kFALSE);
  buf.Reset();
  *obj = extractNextObject(buf);
  if (!*obj) {
    std::cerr << "Error reading element: " << h.full_pathname() << std::endl;
  }
}

void writeMessageFD(const dqmstorepb::ROOTFilePB &dqmstore_output_msg, int out_fd) {
  FileOutputStream out_stream(out_fd);
  GzipOutputStream::Options options;
  options.format = GzipOutputStream::GZIP;
  options.compression_level = 2;
  GzipOutputStream gzip_stream(&out_stream,
                               options);
  dqmstore_output_msg.SerializeToZeroCopyStream(&gzip_stream);

  // make sure we flush before close
  gzip_stream.Close();
  out_stream.Close();
}

void writeMessage(const dqmstorepb::ROOTFilePB &dqmstore_output_msg,
                  const std::string &output_filename) {

  DEBUG(1, "Writing file" << std::endl);

  int out_fd = ::open(output_filename.c_str(),
                      O_WRONLY | O_CREAT | O_TRUNC,
                      S_IRUSR | S_IWUSR |
                      S_IRGRP | S_IWGRP |
                      S_IROTH);

  writeMessageFD(dqmstore_output_msg, out_fd);
  ::close(out_fd);
}


void fillMessage(dqmstorepb::ROOTFilePB &dqmstore_output_msg,
                 const MEStore & micromes) {
  auto mi = micromes.begin();
  auto me = micromes.end();

  DEBUG(1, "Streaming ROOT objects" << std::endl);
  for (; mi != me; ++mi) {
    dqmstorepb::ROOTFilePB::Histo* h = dqmstore_output_msg.add_histo();
    DEBUG(2, "Streaming ROOT object " << mi->fullname() << "\n");
    h->set_full_pathname(mi->fullname());
    TBufferFile buffer(TBufferFile::kWrite);
    buffer.WriteObject(mi->obj);
    h->set_size(buffer.Length());
    h->set_flags(mi->flags);
    h->set_streamed_histo((const void*)buffer.Buffer(),
                          buffer.Length());
    delete mi->obj;
  }
}


void processDirectory(TFile *file,
                      const std::string& curdir,
                      MEStore& micromes) {
  DEBUG(1, "Processing directory " << curdir << "\n");
  file->cd(curdir.c_str());
  TKey *key;
  TIter next (gDirectory->GetListOfKeys());
  while ((key = (TKey *) next())) {
    TObject * obj = key->ReadObj();
    if (dynamic_cast<TDirectory *>(obj)) {
      std::string subdir;
      subdir.reserve(curdir.size() + strlen(obj->GetName()) + 2);
      subdir += curdir;
      if (! curdir.empty())
        subdir += '/';
      subdir += obj->GetName();
      processDirectory(file, subdir, micromes);
    } else if ((dynamic_cast<TH1 *>(obj)) || (dynamic_cast<TObjString *>(obj))) {
      if (dynamic_cast<TH1 *>(obj)) {
        dynamic_cast<TH1 *>(obj)->SetDirectory(nullptr);
      }

      DEBUG(2, curdir << "/" << obj->GetName() << "\n");
      MicroME mme(obj, curdir, obj->GetName());

      micromes.insert(mme);
    }
  }
}


int encodeFile(const std::string &output_filename,
               const std::vector<std::string> &filenames) {
  assert(filenames.size() == 1);
  TFile input(filenames[0].c_str());
  DEBUG(0, "Encoding file " << filenames[0] << std::endl);
  MEStore micromes;
  dqmstorepb::ROOTFilePB dqmstore_message;

  processDirectory(&input, "", micromes);
  fillMessage(dqmstore_message, micromes);
  writeMessage(dqmstore_message, output_filename);

  return 0;
}

int convertFile(const std::string &output_filename,
                const std::vector<std::string> &filenames) {
  assert(filenames.size() == 1);
  TFile output(output_filename.c_str(), "RECREATE");
  DEBUG(0, "Converting file " << filenames[0] << std::endl);
  dqmstorepb::ROOTFilePB dqmstore_message;

  int filedescriptor = ::open(filenames[0].c_str(), O_RDONLY);
  FileInputStream fin(filedescriptor);
  GzipInputStream input(&fin);
  CodedInputStream input_coded(&input);
  input_coded.SetTotalBytesLimit(1024*1024*1024, -1);
  if (!dqmstore_message.ParseFromCodedStream(&input_coded)) {
    std::cout << "Fatal Error opening file "
              << filenames[0] << std::endl;
    return ERR_NOFILE;
  }
  ::close(filedescriptor);

  for (int i = 0; i < dqmstore_message.histo_size(); i++) {
    const dqmstorepb::ROOTFilePB::Histo& h = dqmstore_message.histo(i);
    DEBUG(1, h.full_pathname() << std::endl);
    DEBUG(1, h.size() << std::endl);
    TBufferFile buf(TBufferFile::kRead, h.size(),
                    (void*)h.streamed_histo().data(),
                    kFALSE);
    buf.Reset();
    TObject *obj = extractNextObject(buf);
    std::string path,objname;
    get_info(h, path, objname, &obj);
    gDirectory->cd("/");
    // Find the first path component.
    size_t start = 0;
    size_t end = path.find('/', start);
    if (end == std::string::npos)
    end = path.size();
    while (true)
    {
      std::string part(path, start, end-start);
      if (! gDirectory->Get(part.c_str()))
        gDirectory->mkdir(part.c_str());
      gDirectory->cd(part.c_str());
      // Stop if we reached the end, ignoring any trailing '/'.
      if (end+1 >= path.size())
        break;
      // Find the next path component.
      start = end+1;
      end = path.find('/', start);
      if (end == std::string::npos)
        end = path.size();
    }
    obj->Write();
    DEBUG(1, obj->GetName() << std::endl);
  }
  output.Close();
  return 0;
}

int dumpFiles(const std::vector<std::string> &filenames) {
  assert(!filenames.empty());
  for (int i = 0, e = filenames.size(); i != e; ++i) {
    DEBUG(0, "Dumping file " << filenames[i] << std::endl);
    dqmstorepb::ROOTFilePB dqmstore_message;

    int filedescriptor = ::open(filenames[0].c_str(), O_RDONLY);
    FileInputStream fin(filedescriptor);
    GzipInputStream input(&fin);
    CodedInputStream input_coded(&input);
    input_coded.SetTotalBytesLimit(1024*1024*1024, -1);
    if (!dqmstore_message.ParseFromCodedStream(&input_coded)) {
      std::cout << "Fatal Error opening file "
                << filenames[0] << std::endl;
      return ERR_NOFILE;
    }
    ::close(filedescriptor);

    for (int i = 0; i < dqmstore_message.histo_size(); i++) {
      const dqmstorepb::ROOTFilePB::Histo& h = dqmstore_message.histo(i);
      DEBUG(1, h.full_pathname() << std::endl);
      DEBUG(1, h.size() << std::endl);
      TBufferFile buf(TBufferFile::kRead, h.size(),
                      (void*)h.streamed_histo().data(),
                      kFALSE);
      buf.Reset();
      TObject *obj = extractNextObject(buf);
      DEBUG(1, obj->GetName() << std::endl);
      DEBUG(1, "Flags: " << h.flags() << std::endl);
    }
  }

  return 0;
}

int addFile(MEStore& micromes, int fd) {
  dqmstorepb::ROOTFilePB dqmstore_msg;

  FileInputStream fin(fd);
  GzipInputStream input(&fin);
  CodedInputStream input_coded(&input);
  input_coded.SetTotalBytesLimit(1024*1024*1024, -1);
  if (!dqmstore_msg.ParseFromCodedStream(&input_coded)) {
    std::cout << "Fatal decoding stream: "
              << fd << std::endl;
    return ERR_NOFILE;
  }

  auto hint = micromes.begin();
  for (int i = 0; i < dqmstore_msg.histo_size(); i++) {
    std::string path;
    std::string objname;
    TObject *obj = nullptr;
    const dqmstorepb::ROOTFilePB::Histo &h = dqmstore_msg.histo(i);
    get_info(h, path, objname, &obj);

    MicroME mme(nullptr, path, objname, h.flags());
    auto ir = micromes.insert(hint, mme);
    if (ir->obj != nullptr) {
      // new element was not added
      // so we merge

      ir->add(obj);
      delete obj;
      DEBUG(2, "Merged MicroME " << mme.fullname() << std::endl);
    } else {
      ir->obj = obj;
      DEBUG(2, "Inserted MicroME " << mme.fullname() << std::endl);
    }

    hint = ir;
    ++hint;
  }

  return 0;
}

// The idea is to preload root library (before forking).
// Which is significant for performance and especially memory usage,
// because root aakes a long time to init (and somehow manages to launch a subshell).
void tryRootPreload() {
    // write a single histogram
    TH1F obj_th1f("preload_th1f", "preload_th1f", 2, 0, 1);

    TBufferFile write_buffer(TBufferFile::kWrite);
    write_buffer.WriteObject(&obj_th1f);

    dqmstorepb::ROOTFilePB preload_file;
    dqmstorepb::ROOTFilePB::Histo* hw = preload_file.add_histo();
    hw->set_size(write_buffer.Length());
    hw->set_flags(0);
    hw->set_streamed_histo((const void*)write_buffer.Buffer(), write_buffer.Length());

    // now load this th1f
    const dqmstorepb::ROOTFilePB::Histo &hr = preload_file.histo(0);
    std::string path;
    std::string objname;
    TObject *obj = nullptr;
    get_info(hr, path, objname, &obj);
    delete obj;

    // all done
}

/* fork_id represents the position in a node (node number). */
void addFilesWithFork(int parent_fd, const int fork_id, const int fork_total, const std::vector<std::string>& filenames) {
  DEBUG(1, "Start process: " << fork_id << " parent: " << (fork_id / 2) << std::endl);

  std::list<std::pair<int, int> > children;

  // if this node has a subtree, start it
  for (int i = 0; i < 2; ++i) {
    int child_id = fork_id*2 + i;
    if (child_id > fork_total)
      continue;

    int fd[2];
    ::pipe(fd);

    int child_pid = ::fork();
    if (child_pid == 0) {
      ::prctl(PR_SET_PDEATHSIG, SIGKILL);
      ::close(fd[0]); // close read end

      addFilesWithFork(fd[1], child_id, fork_total, filenames);
      ::close(fd[1]);

      ::_exit(0);
    } else {
      ::close(fd[1]); // close write end
      children.push_back(std::make_pair(fd[0], child_pid));
    }
  }

  // merge all my files
  MEStore microme;

  // select the filenames to process
  // with threads=1, this just selects all the files
  for (unsigned int fi = fork_id - 1; fi < filenames.size(); fi += fork_total) {
    const std::string& file = filenames[fi];
    DEBUG(1, "Adding file " << file << std::endl);

    int filedescriptor;
    if ((filedescriptor = ::open(file.c_str(), O_RDONLY)) == -1) {
      std::cout << "Fatal Error opening file "
                << file << std::endl;

      exit(ERR_NOFILE);
    }

    addFile(microme, filedescriptor);
    ::close(filedescriptor);
  }

  // merge all children
  for (auto& chpair : children) {
    int fd = chpair.first;
    addFile(microme, fd);
    ::close(fd);

    // collect the child, not necessary, but avoids <defunct>
    int status;
    ::waitpid(chpair.second, &status, 0);
  }

  // output everything to fd
  dqmstorepb::ROOTFilePB dqmstore_output_msg;
  fillMessage(dqmstore_output_msg, microme);
  writeMessageFD(dqmstore_output_msg, parent_fd);
};

int addFiles(const std::string &output_filename,
             const std::vector<std::string> &filenames,
             int nthreads) {

  tryRootPreload();

  DEBUG(1, "Writing file" << std::endl);
  int out_fd = ::open(output_filename.c_str(),
                      O_WRONLY | O_CREAT | O_TRUNC,
                      S_IRUSR | S_IWUSR |
                      S_IRGRP | S_IWGRP |
                      S_IROTH);

  addFilesWithFork(out_fd, 1, nthreads, filenames);
  ::close(out_fd);

  return 0;
}

static int
showusage()
{
  static const std::string app_name("fasthadd");

  std::cerr << "Usage: " << app_name
            << " [--[no-]debug] TASK OPTIONS\n\n  "
            << app_name << " [OPTIONS] add [-j NUM_THREADS] -o OUTPUT_FILE [DAT FILE...]\n  "
            << app_name << " [OPTIONS] convert -o ROOT_FILE DAT_FILE\n  "
            << app_name << " [OPTIONS] encode -o DAT_FILE ROOT_FILE\n  "
            << app_name << " [OPTIONS] dump [DAT FILE...]\n  ";
  return ERR_BADCFG;
}

int main(int argc, char * argv[]) {
  int arg;
  int ret = 0;
  int jobs = 1;
  std::string output_file;
  std::vector<std::string> filenames;
  TaskType task;

  filenames.reserve(argc);

  for (arg = 1; arg < argc; ++arg) {
    if (! strcmp(argv[arg], "--no-debug"))
      debug = 0;
    else if (! strcmp(argv[arg], "--debug")
             || ! strcmp(argv[arg], "-d"))
      debug++;
    else
      break;
  }

  if (arg < argc) {
    if (! strcmp(argv[arg], "add")) {
      ++arg;
      task = TASK_ADD;
    } else if (! strcmp(argv[arg], "dump")) {
      ++arg;
      task = TASK_DUMP;
    } else if (! strcmp(argv[arg], "convert")) {
      ++arg;
      task = TASK_CONVERT;
    } else if (! strcmp(argv[arg], "encode")) {
      ++arg;
      task = TASK_ENCODE;
    } else {
      std::cerr << "Unknown action: " << argv[arg] << std::endl;
      return showusage();
    }
  } else {
    std::cerr << "Not enough arguments\n";
    return showusage();
  }

  if (task == TASK_ADD) {
    if ((arg != argc) && (strcmp(argv[arg], "-j") == 0)) {
      jobs = atoi(argv[arg+1]);

      if ((jobs < 1) || (jobs > 128)) {
        std::cerr << "Invalid argument for -j\n";
        return showusage();
      };

      arg += 2;
    }
  }

  if (task == TASK_ADD || task == TASK_CONVERT || task == TASK_ENCODE) {
    if (arg == argc) {
      std::cerr << "add|convert|encode actions requires a -o option to be set\n";
      return showusage();
    }
    if (! strcmp(argv[arg], "-o")) {
      if (arg < argc-1) {
        output_file = argv[++arg];
      } else {
        std::cerr << " -o option requires a value\n";
        return showusage();
      }
    }
  } else if (task == TASK_DUMP) {
    if (arg == argc) {
      std::cerr << "Missing input file(s)\n";
      return showusage();
    }
    for (; arg < argc; ++arg) {
      filenames.emplace_back(argv[arg]);
    }
  }

  if (task == TASK_ADD || task == TASK_CONVERT || task == TASK_ENCODE) {
    if (++arg == argc) {
      std::cerr << "Missing input file(s)\n";
      return showusage();
    }
    for (; arg < argc; ++arg) {
      filenames.emplace_back(argv[arg]);
    }
  }

  if (task == TASK_ADD)
    ret = addFiles(output_file, filenames, jobs);
  else if (task == TASK_DUMP)
    ret = dumpFiles(filenames);
  else if (task == TASK_CONVERT)
    ret = convertFile(output_file, filenames);
  else if (task == TASK_ENCODE)
    ret = encodeFile(output_file, filenames);


  google::protobuf::ShutdownProtobufLibrary();
  return ret;
}