ROOTTests/test/write_threaded_t.cc

0001 #include "TFile.h"
0002 #include "TTree.h"
0003 #include "TBranch.h"
0004 #include "TClass.h"
0005 #include "TThread.h"
0006 #include "TVirtualStreamerInfo.h"
0007 #include "TList.h"
0008 #include "TMap.h"
0009 #include "TObjString.h"
0010 #include "TH1F.h"
0011
0012 #include <memory>
0013 #include <cassert>
0014 #include <iostream>
0015 #include <thread>
0016 #include <atomic>
0017 #include <sstream>
0018
0019 std::atomic<bool> waitToStart{true};
0020 std::atomic<unsigned int> countdownToWrite{0};
0021
0022 void printHelp(const char* iName, int iDefaultNThreads) {
0023   std::cout << iName << " [number of threads] [gDebug value]\n\n"
0024             << "[number of threads] number of threads to use in test\n"
0025             << "[gDebug value] value of gDebug to pass to ROOT (gDebug=1 is useful)\n"
0026             << "If no arguments are given " << iDefaultNThreads << " threads will be used" << std::endl;
0027 }
0028
0029 std::pair<int, int> parseOptionsForNumberOfThreadsAndgDebug(int argc, char** argv) {
0030   constexpr int kDefaultNThreads = 4;
0031   int kDefaultgDebug = gDebug;
0032   int nThreads = kDefaultNThreads;
0033   int newGDebug = kDefaultgDebug;
0034   if (argc >= 2) {
0035     if (strcmp("-h", argv[1]) == 0) {
0036       printHelp(argv[0], kDefaultNThreads);
0037       exit(0);
0038     }
0039
0040     nThreads = atoi(argv[1]);
0041   }
0042   if (argc == 3) {
0043     newGDebug = atoi(argv[2]);
0044   }
0045
0046   if (argc > 3) {
0047     printHelp(argv[0], kDefaultNThreads);
0048     exit(1);
0049   }
0050   return std::make_pair(nThreads, newGDebug);
0051 }
0052
0053 TList* createList() {
0054   auto returnValue = new TList();
0055
0056   for (unsigned int i = 0; i < 10; ++i) {
0057     returnValue->Add(new TList());
0058     returnValue->Add(new TMap());
0059     returnValue->Add(new TObjString());
0060     returnValue->Add(new TH1F());
0061     returnValue->Add(new TH1D());
0062   }
0063
0064   return returnValue;
0065 }
0066
0067 int main(int argc, char** argv) {
0068   auto values = parseOptionsForNumberOfThreadsAndgDebug(argc, argv);
0069
0070   const int kNThreads = values.first;
0071   countdownToWrite = kNThreads;
0072
0073   gDebug = values.second;
0074
0075   //Tell Root we want to be multi-threaded
0076   TThread::Initialize();
0077   //When threading, also have to keep ROOT from logging all TObjects into a list
0078   TObject::SetObjectStat(false);
0079   //Have to avoid having Streamers modify themselves after they have been used
0080   TVirtualStreamerInfo::Optimize(false);
0081
0082   std::vector<std::shared_ptr<std::thread>> threads;
0083   threads.reserve(kNThreads);
0084
0085   for (int i = 0; i < kNThreads; ++i) {
0086     threads.push_back(std::make_shared<std::thread>(std::thread([i]() {
0087       static thread_local TThread s_thread_guard;
0088
0089       std::stringstream nameStream;
0090       nameStream << "write_thread_" << i << ".root";
0091
0092       auto theList = createList();
0093
0094       while (waitToStart)
0095         ;
0096       TFile f(nameStream.str().c_str(), "RECREATE", "test");
0097
0098       auto listTree = new TTree("TheList", "TheList");
0099       listTree->Branch("theList", "TList", &theList);
0100
0101       --countdownToWrite;
0102       while (countdownToWrite != 0)
0103         ;
0104
0105       for (unsigned int i = 0; i < 100; ++i) {
0106         listTree->Fill();
0107       }
0108       f.Write();
0109       f.Close();
0110     })));
0111   }
0112   waitToStart = false;
0113   for (auto& t : threads) {
0114     t->join();
0115   }
0116
0117   return 0;
0118 }