Write.cpp

Go to the documentation of this file.

#include "cafe/Write.hpp"
#include "cafe/Config.hpp"
#include "cafe/Event.hpp"

#include <cstdlib>
#include <iostream>
#include <sstream>
#include <algorithm>
#include <cassert>
#include "TFile.h"
#include "TTree.h"
#include "TBranchRef.h"
#include "TRefTable.h"
#include "TObjArray.h"
#include "TBranchElement.h"

namespace {

    // enable this branch and all sub-branches
    void enableBranch(TBranch *br)
    {
        if(br != 0) {
            br->ResetBit(kDoNotProcess);
            TIter iter(br->GetListOfBranches());
            while(TBranch *sub = (TBranch *)iter.Next()) {
                enableBranch(sub);
            }
        } 
    }


    void disableBranch(TBranch *br)
    {
        if(br != 0) {
            br->SetBit(kDoNotProcess);
            TIter iter(br->GetListOfBranches());
            while(TBranch *sub = (TBranch *)iter.Next()) {
                disableBranch(sub);
            }
        }
    }


}

namespace cafe {

    Write::Write(const char *name)
        : Processor(name),
          _file(0),
          _tree(0),
          _fileOwner(false),
          _updateObjectVersion(true)
    {
        Config config(name);
        _pattern.set(config.get("File", ""));
        _branches = config.getVString("Branches", " ,");
        if(_branches.empty()) _branches = config.getVString("Enable", " ,");
        _disable  = config.getVString("Disable", " ,");
        _treeName = config.get("Tree", "TMBTree");
        _autosave = config.get("AutoSave", -1);
        _destination = config.get("Destination", "");
        _filesPerOutput = config.get("FilesPerOutput", 0);
        _filesPerOutputDone = _filesPerOutput;
        _updateObjectVersion = config.get("UpdateObjectVersion", 1);
    }

    Write::~Write()
    {
    }


    void Write::finish()
    {
        if(_file && _fileOwner) {
            _file->Write();
            _file->Close();
            if(!_destination.empty()) {
                copy_file();
            }
            delete _file;
        }
    }


    bool Write::processEvent(Event& event)
    {
        if(_tree) {
            TTree *tree = event.getTree();
            Long64_t current = tree->GetReadEntry();

            // Loop over enabled branches of output tree
            // Read only if necessary (readBranch keeps track
            // of the top-level branches).
            TIter iter(_tree->GetListOfBranches());
            while(TBranch *br = (TBranch *)iter.Next()) {

                // TODO: this test is no longer necessary ?
                if(!br->TestBit(kDoNotProcess)) {
                    event.readBranch(br->GetName());
                }
            }

            _tree->Fill();

            if(_autosave > 0 && (eventCount() % _autosave == 0)) {
                _tree->AutoSave();
            }
        }
        return true;
    }

    void Write::inputFileOpened(TFile *file)
    {
        TTree *tree  = (TTree *)file->Get(_treeName.c_str());
        if(tree == 0) tree = (TTree *)getDirectory()->Get(_treeName.c_str());
        if(tree == 0) tree = (TTree *)gROOT->Get(_treeName.c_str());
        if(tree == 0) {
            _tree = 0;
            err() << "Write[" << name() << "] No tree in file named: " << _treeName << std::endl;
            return;
        }

        // check this before we generate a new name
        if(_file && _filesPerOutput && (++_filesPerOutputDone < _filesPerOutput)) {
            // Same output file, just update tree addresses
            tree->CopyAddresses(_tree);
            return;
        }

        // Now generate a new name if necessory
        std::string new_name = _pattern.replace(file->GetName());
        if(_file) { 
            if(new_name == _file->GetName()) {
                // Same output file, just update tree addresses
                tree->CopyAddresses(_tree);
                return;
            }
            if(_fileOwner) {
                _file->Write();
                _file->Close();
                if(!_destination.empty()) {
                    copy_file();
                }
                delete _file;
            }
            _file = 0;
        }

        _filesPerOutputDone = 0;

        _file = dynamic_cast<TFile*>(gROOT->GetFile(new_name.c_str()));


        if(_file == 0) {
            out() << "Write[" << name() << "] Opening new output file: " << new_name << std::endl;
            _file = new TFile(new_name.c_str(), "RECREATE");
            _fileOwner = true;
        } else {
            out() << "Write[" << name() << "] Using existing file: " << new_name << std::endl;
            _fileOwner = false;
        }

        // We allow to enable/disable top-level branches only.

        if(!_branches.empty()) {
            // Disable all branches in input tree
            tree->SetBranchStatus("*",0);

            for(std::vector<std::string>::iterator it = _branches.begin();
                it != _branches.end();
                ++it) {
                out() << "Write[" << name() << "] Enabling branch: " << *it << std::endl;
                if(TBranch *br = tree->GetBranch((*it).c_str())) {
                    enableBranch(br);
                } else {
                    err() << "Write[" << name() << "] Branch not found: " << *it << std::endl;
                    _branches.erase(it--);
                }
            }

        } else if(!_disable.empty()) {

            tree->SetBranchStatus("*", 1);

            // Disable only branches in list
            for(std::vector<std::string>::iterator it = _disable.begin();
                it != _disable.end();
                ++it) {
                out() << "Write[" << name() << "] Disabling branch: " << *it << std::endl;
                if(TBranch *br = tree->GetBranch((*it).c_str())) {
                    disableBranch(br);
                } else {
                    err() << "Write[" << name() << "] Branch not found: " << *it << std::endl;
                    _disable.erase(it--);
                }
            }
        } 

        //
        // Disable branches in updatedBranches
        // 
        // If requested, look for objects in the original tree 
        // that have old object format versions.
        _updatedBranches.clear();

        if (_updateObjectVersion) {
            TIter branchNext (tree->GetListOfBranches());
            while (TBranch *br = static_cast<TBranch*>(branchNext.Next())) {
                if (!br->TestBit(kDoNotProcess) && containsUpdatedClass (br)) {
                    _updatedBranches.push_back (br->GetName());
                    disableBranch(br);
                }
            }
        }

        // Clone all active branches
        _tree = tree->CloneTree(0);

        // Return everything to normal
        tree->SetBranchStatus("*",1);

        for(std::vector<std::string>::iterator it = _updatedBranches.begin();
            it != _updatedBranches.end();
            ++it) {
            out() << "Write[" << name() << "] Updating branch: " << *it << std::endl;
            updateBranch(tree, *it);
        }

        // This is for buggy ROOT versions
        if(TBranchRef *ref = _tree->GetBranchRef()) {
            // BranchRef is not properly reset after cloning
            ref->Reset();
            if(TRefTable *table = ref->GetRefTable()) {
                // TRefTable is not cleared; crashes when only some branches
                // are written out.
                if(TObjArray *parents = table->GetParents()) {
                    parents->Clear();
                }
            }
        } else {
            // Always create a BranchRef in the new tree, even
            // if it didn't have one before.
            _tree->BranchRef();
        }

        // Some friends work, others don't, e.g. if they are only
        // in memory.
        // Remove them and have he user write them out explicitly
        // with a separate Write processor (can go in same file).
        if(TList *fr = _tree->GetListOfFriends()) {
            fr->Clear();
        }

    }

    bool Write::containsUpdatedClass (TBranch *br)
    {
        TBranchElement *be = static_cast<TBranchElement*>(br);

        int tree_version = be->GetInfo()->GetClassVersion();
        int class_version = TClass::GetClass(be->GetClassName())->GetStreamerInfo()->GetClassVersion();
        if (tree_version < class_version) {
            return true;
        }

        TIter branchNext (be->GetListOfBranches());
        while (TBranchElement *be = static_cast<TBranchElement*>(branchNext.Next())) {
            if (containsUpdatedClass(be)) {
                return true;
            }
        }
        return false;
    }

    void Write::updateBranch(TTree *input_tree, const std::string& branch_name)
    {
        TBranch *br = input_tree->GetBranch(branch_name.c_str());
        assert(br != 0);

        // We want to store this data anyway, so let's load the first
        // entry to set the address of the various objects.
        br->GetEntry(0);

        void* obj = br->GetAddress();
        if (obj != 0) {
            if (std::string(br->GetClassName()) == "TClonesArray") {
                _tree->Branch(branch_name.c_str(), "TClonesArray", obj);
            } else {
                _tree->Branch(branch_name.c_str(), br->GetClassName(), obj);
            }
        } // else what ??
    }

    void Write::copy_file()
    {
        std::ostringstream s;

        if(_destination.find(':') != std::string::npos) {
            s << "/usr/krb5/bin/kbatch && /usr/krb5/bin/rcp " 
              << _file->GetName() <<  " "
              << _destination;
        } else {
            s << "cp " 
              << _file->GetName() <<  " "
              << _destination;
        }
        out() << "Write[" << name() << "] : Copied " << _file->GetName() 
              << " to " << _destination << std::endl;

        if(system(s.str().c_str()) == 0) {
            remove(_file->GetName());
            out() << "Write[" << name() << "] : Removed " << _file->GetName() << std::endl;
        } else {
            err() << "Write[" << name() << "] : Cannot copy " << _file->GetName() 
                  << "to" << _destination << std::endl;
        }
    }


}

ClassImp(cafe::Write)

---