(Experimental) Pickle: Consider Series identity when unpickling

CMakeLists.txt

@@ -605,6 +605,7 @@ if(openPMD_HAVE_PYTHON)
         src/binding/python/ParticleSpecies.cpp
         src/binding/python/PatchRecord.cpp
         src/binding/python/PatchRecordComponent.cpp
+        src/binding/python/Pickle.cpp
         src/binding/python/Record.cpp
         src/binding/python/RecordComponent.cpp
         src/binding/python/MeshRecordComponent.cpp

include/openPMD/Iteration.hpp

@@ -174,8 +174,8 @@ class Iteration
     friend class Container;
     friend class Series;
     friend class internal::AttributableData;
-    template <typename T>
-    friend T &internal::makeOwning(T &self, Series);
+    template <typename T, typename Series_type>
+    friend T &internal::makeOwning(T &self, Series_type);
     friend class Writable;
     friend class StatefulIterator;
     friend class StatefulSnapshotsContainer;

include/openPMD/ParticleSpecies.hpp

@@ -40,8 +40,8 @@ class ParticleSpecies
     friend class Container<ParticleSpecies>;
     friend class Container<Record>;
     friend class Iteration;
-    template <typename T>
-    friend T &internal::makeOwning(T &self, Series);
+    template <typename T, typename Series_type>
+    friend T &internal::makeOwning(T &self, Series_type);
     friend class internal::ScientificDefaults;
     friend class Attributable;
 

include/openPMD/RecordComponent.hpp

@@ -130,8 +130,8 @@ class RecordComponent
     friend class DynamicMemoryView;
     friend class internal::RecordComponentData;
     friend class MeshRecordComponent;
-    template <typename T>
-    friend T &internal::makeOwning(T &self, Series);
+    template <typename T, typename Series_type>
+    friend T &internal::makeOwning(T &self, Series_type);
     friend class internal::ScientificDefaults;
     friend class Attributable;
 

include/openPMD/Series.hpp

@@ -294,6 +294,8 @@ class Series : public Attributable
     friend class internal::SeriesData;
     friend class internal::AttributableData;
     friend class StatefulSnapshotsContainer;
+    template <typename T, typename Series_type>
+    friend T &internal::makeOwning(T &self, Series_type);
 
 public:
     explicit Series();
@@ -778,6 +780,8 @@ class Series : public Attributable
      */
     void close();
 
+    [[nodiscard]] bool closed() const;
+
     void visitHierarchy(HierarchyVisitor &v, bool recursive) override;
 
     /**
@@ -804,6 +808,11 @@ OPENPMD_private
     using Data_t = internal::SeriesData;
     std::shared_ptr<Data_t> m_series = nullptr;
 
+    inline std::shared_ptr<Data_t> getShared()
+    {
+        return m_series;
+    }
+
     inline Data_t &get()
     {
         if (m_series)

include/openPMD/backend/Attributable.hpp

@@ -207,8 +207,8 @@ namespace internal
      * Instantiations for T exist for types RecordComponent,
      * MeshRecordComponent, Mesh, Record, ParticleSpecies, Iteration.
      */
-    template <typename T>
-    T &makeOwning(T &self, Series);
+    template <typename T, typename Series_type>
+    T &makeOwning(T &self, Series_type);
 } // namespace internal
 
 namespace debug
@@ -241,8 +241,8 @@ class Attributable
     friend class Writable;
     friend class internal::RecordComponentData;
     friend void debug::printDirty(Series const &);
-    template <typename T>
-    friend T &internal::makeOwning(T &self, Series);
+    template <typename T, typename Series_type>
+    friend T &internal::makeOwning(T &self, Series_type);
     friend class StatefulSnapshotsContainer;
     friend class internal::AttributableData;
     friend class Snapshots;
@@ -457,6 +457,13 @@ class Attributable
 
     [[nodiscard]] OpenpmdStandard openPMDStandard() const;
 
+    /** Returns the persistent immutable memory ID of the underlying Series.
+     *
+     * Useful when trying to determine which API handles refer to the same IO
+     * instance.
+     */
+    [[nodiscard]] uintptr_t memoryID() const;
+
     // clang-format off
 OPENPMD_protected
     // clang-format on

include/openPMD/backend/BaseRecord.hpp

@@ -199,8 +199,8 @@ class BaseRecord
     friend class internal::BaseRecordData;
     template <typename, typename, typename>
     friend class internal::ScalarIterator;
-    template <typename T>
-    friend T &internal::makeOwning(T &self, Series);
+    template <typename T, typename Series_type>
+    friend T &internal::makeOwning(T &self, Series_type);
     friend class internal::ScientificDefaults;
 
     using Data_t =

include/openPMD/binding/python/Pickle.hpp

@@ -27,13 +27,55 @@
 
 #include "Common.hpp"
 
+#include <cstdint>
 #include <exception>
+#include <memory>
+#include <shared_mutex>
 #include <string>
+#include <sys/types.h>
 #include <tuple>
 #include <vector>
 
 namespace openPMD
 {
+/*
+ * unpickled_series, as in "plural series"; this is a cache structure for series
+ * objects that have been unpickled. This cache structure fixes the issue
+ * described in https://github.com/openPMD/openPMD-api/issues/1919.
+ * Idea: One single Series instance may have multiple handles referencing it.
+ * When pickling and unpickling these references, the underlying Series must be
+ * restored once only, in order to keep the reference structure. Otherwise
+ * something like `data = E_x[:]; series.flush();` will not work, because `E_x`
+ * no longer references the same Series instance as `series`.
+ *
+ * For this, the pickle structure contains as first entry the internal
+ * (immutable) SharedAttributable pointer address of the `Series` object
+ * referenced by any handle. When unpickling, this is used to restore shared
+ * handles in accordance with their original reference structure. The pointers
+ * themselves are not restored (this would not be possible), but they are used
+ * as equivalence classes.
+ */
+struct unpickled_series
+{
+    // Cache restored object by original Series ID (i.e. internal immutable
+    // pointer address). IDs are not restored equivalently, but this does not
+    // matter. They are necessary only for figuring out which handles point to
+    // the same objects.
+    // The cached Series objects are stored as weak_ptr, since they are memory
+    // managed by the Python side. The C++ side just needs to check if the
+    // weak_ptr is still valid when handing out a new reference. If not, reopen.
+    std::map<uintptr_t, std::weak_ptr<Series>> m_series_by_former_id;
+    std::shared_mutex m_mutex;
+
+    auto get(uintptr_t id, std::string const &filename)
+        -> std::shared_ptr<Series>;
+};
+
+/*
+ * Cache the Series per thread.
+ */
+extern unpickled_series cache;
+
 /** Helper to Pickle Attributable Classes
  *
  * @tparam T_Args the types in pybind11::class_ - the first type will be pickled
@@ -48,62 +90,35 @@ add_pickle(pybind11::class_<T_Args...> &cl, T_SeriesAccessor &&seriesAccessor)
 {
     // helper: get first class in py::class_ - that's the type we pickle
     using PickledClass =
-        typename std::tuple_element<0, std::tuple<T_Args...> >::type;
+        typename std::tuple_element<0, std::tuple<T_Args...>>::type;
 
     cl.def(
         py::pickle(
             // __getstate__
             [](const PickledClass &a) {
                 // Return a tuple that fully encodes the state of the object
                 Attributable::MyPath const myPath = a.myPath();
-                return py::make_tuple(myPath.filePath(), myPath.group);
+                // retrieve Series even though retrieveSeries is protected...
+                return py::make_tuple(
+                    a.memoryID(), myPath.filePath(), myPath.group);
             },
 
             // __setstate__
             [&seriesAccessor](py::tuple const &t) {
-                // our tuple has exactly two elements: filePath & group
-                if (t.size() != 2)
+                // Our tuple has exactly three elements: Series ID, filePath &
+                // group.
+                //  Check the documentation of unpickled_series above for
+                // the reasoning behind Series ID.
+                if (t.size() != 3)
                     throw std::runtime_error("Invalid state!");
 
-                std::string const filename = t[0].cast<std::string>();
+                auto id = t[0].cast<uintptr_t>();
+                std::string const filename = t[1].cast<std::string>();
                 std::vector<std::string> const group =
-                    t[1].cast<std::vector<std::string> >();
-
-                /*
-                 * Cache the Series per thread.
-                 */
-                thread_local std::optional<openPMD::Series> series;
-                bool re_initialize = [&]() {
-                    try
-                    {
-                        return !series.has_value() ||
-                            !series->operator bool() ||
-                            auxiliary::replace_all(
-                                series->myPath().filePath(), "\\", "/") !=
-                            auxiliary::replace_all(filename, "\\", "/");
-                    }
-                    /*
-                     * Better safe than sorry, if anything goes wrong because
-                     * the Series is in a weird state, just reinitialize it.
-                     */
-                    catch (...)
-                    {
-                        return true;
-                    }
-                }();
-                if (re_initialize)
-                {
-                    /*
-                     * Do NOT close the old Series, it might still be active in
-                     * terms of handed-out handles.
-                     */
-                    series = std::make_optional<Series>(
-                        filename,
-                        Access::READ_ONLY,
-                        "defer_iteration_parsing = true");
-                }
+                    t[2].cast<std::vector<std::string>>();
 
-                return seriesAccessor(*series, group);
+                auto series = cache.get(id, filename);
+                return seriesAccessor(std::move(series), group);
             }));
 }
 } // namespace openPMD

src/Series.cpp

@@ -3564,6 +3564,20 @@ void Series::close()
     m_attri.reset();
 }
 
+bool Series::closed() const
+{
+    if (!operator bool())
+    {
+        return true;
+    }
+    auto &w = writable();
+    if (!w.IOHandler)
+    {
+        throw error::Internal("Series went into illegal state");
+    }
+    return !w.IOHandler->has_value();
+}
+
 void Series::visitHierarchy(HierarchyVisitor &v, bool recursive)
 {
     if (recursive)

src/backend/Attributable.cpp

@@ -343,6 +343,11 @@ OpenpmdStandard Attributable::openPMDStandard() const
     return IOHandler()->m_standard;
 }
 
+uintptr_t Attributable::memoryID() const
+{
+    return reinterpret_cast<uintptr_t>(&retrieveSeries().Attributable::get());
+}
+
 template <bool flush_entire_series>
 void Attributable::seriesFlush_impl(internal::FlushParams const &flushParams)
 {
@@ -625,8 +630,8 @@ void Attributable::linkHierarchy(Writable &w)
 
 namespace internal
 {
-    template <typename T>
-    T &makeOwning(T &self, Series s)
+    template <typename T, typename Series_type>
+    T &makeOwning(T &self, Series_type s)
     {
         /*
          * `self` is a handle object such as RecordComponent or Mesh (see
@@ -665,11 +670,15 @@ namespace internal
         return self;
     }
 
-    template RecordComponent &makeOwning(RecordComponent &, Series);
-    template MeshRecordComponent &makeOwning(MeshRecordComponent &, Series);
-    template Mesh &makeOwning(Mesh &, Series);
-    template Record &makeOwning(Record &, Series);
-    template ParticleSpecies &makeOwning(ParticleSpecies &, Series);
-    template Iteration &makeOwning(Iteration &, Series);
+    template Series &makeOwning(Series &, std::shared_ptr<Series>);
+    template RecordComponent &
+    makeOwning(RecordComponent &, std::shared_ptr<Series>);
+    template MeshRecordComponent &
+    makeOwning(MeshRecordComponent &, std::shared_ptr<Series>);
+    template Mesh &makeOwning(Mesh &, std::shared_ptr<Series>);
+    template Record &makeOwning(Record &, std::shared_ptr<Series>);
+    template ParticleSpecies &
+    makeOwning(ParticleSpecies &, std::shared_ptr<Series>);
+    template Iteration &makeOwning(Iteration &, std::shared_ptr<Series>);
 } // namespace internal
 } // namespace openPMD

src/binding/python/Attributable.cpp

@@ -657,7 +657,10 @@ void init_Attributable(py::module &m)
         .def(
             "populate_missing_metadata",
             &Attributable::populateMissingMetadata,
-            py::arg("recursive"));
+            py::arg("recursive"))
+        .def_property_readonly("memory_id", [](Attributable const &attr) {
+            return attr.memoryID();
+        });
 
     py::bind_vector<PyAttributeKeys>(m, "Attribute_Keys");
 }

src/binding/python/Iteration.cpp

@@ -117,9 +117,11 @@ void init_Iteration(py::module &m)
                 py::keep_alive<0, 1>()));
 
     add_pickle(
-        cl, [](openPMD::Series series, std::vector<std::string> const &group) {
+        cl,
+        [](std::shared_ptr<openPMD::Series> series,
+           std::vector<std::string> const &group) {
             uint64_t const n_it = std::stoull(group.at(1));
-            auto res = series.iterations[n_it];
+            auto res = series->iterations[n_it];
             return internal::makeOwning(res, std::move(series));
         });
 

src/binding/python/Mesh.cpp

@@ -196,9 +196,11 @@ Ref.: https://github.com/openPMD/openPMD-standard/pull/193)"[1])
             "set_grid_unit_SI",
             py::overload_cast<double>(&Mesh::setGridUnitSI));
     add_pickle(
-        cl, [](openPMD::Series series, std::vector<std::string> const &group) {
+        cl,
+        [](std::shared_ptr<openPMD::Series> series,
+           std::vector<std::string> const &group) {
             uint64_t const n_it = std::stoull(group.at(1));
-            auto res = series.iterations[n_it].open().meshes[group.at(3)];
+            auto res = series->iterations[n_it].open().meshes[group.at(3)];
             return internal::makeOwning(res, std::move(series));
         });
 

src/binding/python/MeshRecordComponent.cpp

@@ -82,10 +82,12 @@ void init_MeshRecordComponent(py::module &m)
             "Relative position of the component on an element "
             "(node/cell/voxel) of the mesh");
     add_pickle(
-        cl, [](openPMD::Series series, std::vector<std::string> const &group) {
+        cl,
+        [](std::shared_ptr<openPMD::Series> series,
+           std::vector<std::string> const &group) {
             uint64_t const n_it = std::stoull(group.at(1));
             auto res =
-                series.iterations[n_it]
+                series->iterations[n_it]
                     .open()
                     .meshes[group.at(3)]
                            [group.size() < 5 ? MeshRecordComponent::SCALAR

src/binding/python/ParticleSpecies.cpp

@@ -56,10 +56,12 @@ void init_ParticleSpecies(py::module &m)
                 // garbage collection: return value must be freed before Series
                 py::keep_alive<0, 1>()));
     add_pickle(
-        cl, [](openPMD::Series series, std::vector<std::string> const &group) {
+        cl,
+        [](std::shared_ptr<openPMD::Series> series,
+           std::vector<std::string> const &group) {
             uint64_t const n_it = std::stoull(group.at(1));
             ParticleSpecies res =
-                series.iterations[n_it].open().particles[group.at(3)];
+                series->iterations[n_it].open().particles[group.at(3)];
             return internal::makeOwning(res, std::move(series));
         });