diff --git a/core/base/inc/TObject.h b/core/base/inc/TObject.h index 7aaf4f6f00cf1..456c21ad10fc3 100644 --- a/core/base/inc/TObject.h +++ b/core/base/inc/TObject.h @@ -36,9 +36,11 @@ class TTimer; namespace ROOT { namespace Internal { bool DeleteChangesMemoryImpl(); + void MarkTObjectAsNotOnHeap(TObject &obj); }} class TObject { + friend void ROOT::Internal::MarkTObjectAsNotOnHeap(TObject &obj); private: UInt_t fUniqueID; ///< object unique identifier diff --git a/core/base/src/TObject.cxx b/core/base/src/TObject.cxx index f02db598949bb..3decd1db81926 100644 --- a/core/base/src/TObject.cxx +++ b/core/base/src/TObject.cxx @@ -1265,3 +1265,8 @@ void TObject::operator delete[](void *ptr, void *vp) { TStorage::ObjectDealloc(ptr, vp); } + +void ROOT::Internal::MarkTObjectAsNotOnHeap(TObject &obj) +{ + obj.fBits &= ~TObject::kIsOnHeap; +} diff --git a/io/io/src/RFile.cxx b/io/io/src/RFile.cxx index d333810b2b283..2d7f78120a8c1 100644 --- a/io/io/src/RFile.cxx +++ b/io/io/src/RFile.cxx @@ -304,9 +304,18 @@ void *RFile::GetUntyped(std::string_view path, void *obj = key ? key->ReadObjectAny(cls) : nullptr; if (obj) { - // Disavow any ownership on `obj` - if (auto autoAddFunc = cls->GetDirectoryAutoAdd(); autoAddFunc) { - autoAddFunc(obj, nullptr); + // Disavow any ownership on `obj` unless the object is a TTree, in which case we need to link it to our internal + // file for it to be usable. + if (auto autoAddFunc = cls->GetDirectoryAutoAdd()) { + if (cls->InheritsFrom("TTree")) { + autoAddFunc(obj, fFile.get()); + // NOTE(gparolini): this is a hacky but effective way of preventing the Tree from being deleted by + // the internal TFile once we close it. We need to avoid that because this TTree will be returned inside + // a unique_ptr and would end up being double-freed if we allowed ROOT to do its own memory management. + ROOT::Internal::MarkTObjectAsNotOnHeap(*static_cast(obj)); + } else { + autoAddFunc(obj, nullptr); + } } } else if (key) { R__LOG_INFO(RFileLog()) << "Tried to get object '" << path << "' of type " << cls->GetName() diff --git a/io/io/test/CMakeLists.txt b/io/io/test/CMakeLists.txt index 641b34b1ec939..2e9efbd370683 100644 --- a/io/io/test/CMakeLists.txt +++ b/io/io/test/CMakeLists.txt @@ -18,7 +18,19 @@ if(uring AND NOT DEFINED ENV{ROOTTEST_IGNORE_URING}) ROOT_ADD_GTEST(RIoUring RIoUring.cxx LIBRARIES RIO) endif() -ROOT_ADD_GTEST(rfile rfile.cxx LIBRARIES RIO Hist ROOTNTuple) +ROOT_STANDARD_LIBRARY_PACKAGE(RFileTestIncludes + NO_INSTALL_HEADERS + HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/RFileTestIncludes.hxx + SOURCES RFileTestIncludes.cxx + LINKDEF RFileTestIncludesLinkDef.h + DEPENDENCIES RIO Tree) +configure_file(RFileTestIncludes.hxx . COPYONLY) +if(MSVC AND NOT CMAKE_GENERATOR MATCHES Ninja) + add_custom_command(TARGET RFileTestIncludes POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/$/libRFileTestIncludes.dll + ${CMAKE_CURRENT_BINARY_DIR}/libRFileTestIncludes.dll) +endif() +ROOT_ADD_GTEST(rfile rfile.cxx LIBRARIES RIO Hist ROOTNTuple Tree RFileTestIncludes) if(pyroot) ROOT_ADD_PYUNITTEST(rfile_py rfile.py) endif() diff --git a/io/io/test/RFileTestIncludes.cxx b/io/io/test/RFileTestIncludes.cxx new file mode 100644 index 0000000000000..0f97b13489efa --- /dev/null +++ b/io/io/test/RFileTestIncludes.cxx @@ -0,0 +1,17 @@ +#include "RFileTestIncludes.hxx" + +TTreeDestructorCounter::~TTreeDestructorCounter() { + ++fgTimesDestructed; +} + +void TTreeDestructorCounter::ResetTimesDestructed() +{ + fgTimesDestructed = 0; +} + +int TTreeDestructorCounter::GetTimesDestructed() +{ + return fgTimesDestructed; +} + +int TTreeDestructorCounter::fgTimesDestructed = 0; diff --git a/io/io/test/RFileTestIncludes.hxx b/io/io/test/RFileTestIncludes.hxx new file mode 100644 index 0000000000000..c8e39510db574 --- /dev/null +++ b/io/io/test/RFileTestIncludes.hxx @@ -0,0 +1,21 @@ +#ifndef ROOT_RFILE_TEST_INCLUDES +#define ROOT_RFILE_TEST_INCLUDES + +#include + +// WARNING: this class is used in some tests that check its `fgTimesDestructed`. +// ResetTimesDestructed() must be called before using this class for consistency. +class TTreeDestructorCounter : public TTree { + static int fgTimesDestructed; + +public: + static void ResetTimesDestructed(); + static int GetTimesDestructed(); + + using TTree::TTree; + ~TTreeDestructorCounter(); + + ClassDefOverride(TTreeDestructorCounter, 2); +}; + +#endif diff --git a/io/io/test/RFileTestIncludesLinkDef.h b/io/io/test/RFileTestIncludesLinkDef.h new file mode 100644 index 0000000000000..3254c0d1b7a1c --- /dev/null +++ b/io/io/test/RFileTestIncludesLinkDef.h @@ -0,0 +1,5 @@ +#ifdef __CLING__ + +#pragma link C++ class TTreeDestructorCounter+; + +#endif diff --git a/io/io/test/rfile.cxx b/io/io/test/rfile.cxx index 47db58fb8b293..51407b7a95459 100644 --- a/io/io/test/rfile.cxx +++ b/io/io/test/rfile.cxx @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -15,6 +16,8 @@ #include #include +#include "RFileTestIncludes.hxx" + using ROOT::Experimental::RFile; using ROOT::TestSupport::FileRaii; @@ -721,3 +724,89 @@ TEST(RFile, RNTuple) EXPECT_EQ(reader->GetNEntries(), 1); EXPECT_FLOAT_EQ(reader->GetView("x")(0), 42); } + +TEST(RFile, TTreeRead) +{ + FileRaii fileGuard("test_rfile_ttree_read.root"); + + { + auto file = std::unique_ptr(TFile::Open(fileGuard.GetPath().c_str(), "RECREATE")); + auto tree = std::make_unique("tree", "tree"); + int x; + tree->Branch("x", &x); + for (int i = 0; i < 10; ++i) { + x = i; + tree->Fill(); + } + tree->Write(); + } + + { + auto file = ROOT::Experimental::RFile::Open(fileGuard.GetPath()); + auto tree = file->Get("tree"); + ASSERT_NE(tree, nullptr); + EXPECT_EQ(tree->GetEntries(), 10); + int x; + tree->SetBranchAddress("x", &x); + for (auto i = 0u; i < tree->GetEntries(); ++i) { + tree->GetEntry(i); + EXPECT_EQ(x, i); + } + } +} + +TEST(RFile, TTreeReadAfterClose) +{ + FileRaii fileGuard("test_rfile_ttree_read_after_close.root"); + + { + auto file = std::unique_ptr(TFile::Open(fileGuard.GetPath().c_str(), "RECREATE")); + auto tree = std::make_unique("tree", "tree"); + int x; + tree->Branch("x", &x); + for (int i = 0; i < 10; ++i) { + x = i; + tree->Fill(); + } + tree->Write(); + } + + { + auto file = ROOT::Experimental::RFile::Open(fileGuard.GetPath()); + auto tree = file->Get("tree"); + file.reset(); // close the file + ASSERT_NE(tree, nullptr); + EXPECT_EQ(tree->GetEntries(), 10); + EXPECT_EQ(tree->GetEntry(0), -1); // -1 means "I/O error" + } +} + +TEST(RFile, TTreeNoDoubleFree) +{ + FileRaii fileGuard("test_rfile_ttree_read_no_double_free.root"); + + TTreeDestructorCounter::ResetTimesDestructed(); + + { + auto file = std::unique_ptr(TFile::Open(fileGuard.GetPath().c_str(), "RECREATE")); + auto tree = std::make_unique("tree", "tree"); + int x; + tree->Branch("x", &x); + for (int i = 0; i < 10; ++i) { + x = i; + tree->Fill(); + } + tree->Write(); + } + EXPECT_EQ(TTreeDestructorCounter::GetTimesDestructed(), 1); + + { + auto file = ROOT::Experimental::RFile::Open(fileGuard.GetPath()); + auto tree = file->Get("tree"); + file.reset(); // close the file (does not delete the three) + // tree is deleted here + } + + // destructed once during writing, once during reading. + EXPECT_EQ(TTreeDestructorCounter::GetTimesDestructed(), 2); +} diff --git a/io/io/test/rfile.py b/io/io/test/rfile.py index 135e6afac9132..4dfd06e8d0d4e 100644 --- a/io/io/test/rfile.py +++ b/io/io/test/rfile.py @@ -144,5 +144,29 @@ def test_putUnsupportedType(self): finally: os.remove(fileName) + def test_getTree(self): + from array import array + + fileName = "test_rfile_gettree_py.root" + try: + with ROOT.TFile.Open(fileName, "RECREATE") as file: + tree = ROOT.TTree("tree", "tree") + x = array("i", [0]) + tree.Branch("myColumn", x, "myColumn/I") + for i in range(10): + x[0] = i + tree.Fill() + tree.Write() + + with RFile.Open(fileName) as rfile: + tree = rfile.Get("tree") + self.assertIsNot(tree, None) + self.assertEqual(tree.GetEntries(), 10) + xs = [entry.myColumn for entry in tree] + self.assertEqual(xs, [x for x in range(10)]) + finally: + os.remove(fileName) + + if __name__ == "__main__": unittest.main()