From 4384f577807f9275d3b2e6544bfc0f3fbd840d35 Mon Sep 17 00:00:00 2001
From: Gigon Bae <gbae@nvidia.com>
Date: Fri, 27 Oct 2023 01:45:47 -0700
Subject: [PATCH] Fix memory corruption issue in ThreadBatchDataLoader

- Return a shared pointer of CuImageIterator, taking ownership of
  CuImageIterator instance, instead of returning a copy of
  CuImageIterator
- Do not create new instance for CuImageIterator when __iter__ is
  called. Instead, return self.
- Wait until all threads are finished in ThreadBatchDataLoader
  destructor

Signed-off-by: Gigon Bae <gbae@nvidia.com>
---
 cpp/src/loader/thread_batch_data_loader.cpp |  5 ++++-
 python/pybind11/cucim_py.cpp                | 10 +++++-----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/cpp/src/loader/thread_batch_data_loader.cpp b/cpp/src/loader/thread_batch_data_loader.cpp
index 89de8e677..0b3ba3af4 100644
--- a/cpp/src/loader/thread_batch_data_loader.cpp
+++ b/cpp/src/loader/thread_batch_data_loader.cpp
@@ -1,6 +1,6 @@
 /*
  * Apache License, Version 2.0
- * Copyright 2021 NVIDIA Corporation
+ * Copyright 2021-2023 NVIDIA Corporation
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -86,6 +86,9 @@ ThreadBatchDataLoader::ThreadBatchDataLoader(LoadFunc load_func,
 
 ThreadBatchDataLoader::~ThreadBatchDataLoader()
 {
+    // Wait until all tasks are done.
+    while (wait_batch() > 0);
+
     cucim::io::DeviceType device_type = out_device_.type();
     for (auto& raster_ptr : raster_data_)
     {
diff --git a/python/pybind11/cucim_py.cpp b/python/pybind11/cucim_py.cpp
index a79ae103f..bb8c37af8 100644
--- a/python/pybind11/cucim_py.cpp
+++ b/python/pybind11/cucim_py.cpp
@@ -200,7 +200,7 @@ PYBIND11_MODULE(_cucim, m)
             },
             py::call_guard<py::gil_scoped_release>());
 
-    py::class_<CuImageIterator<CuImage>>(m, "CuImageIterator") //
+    py::class_<CuImageIterator<CuImage>, std::shared_ptr<CuImageIterator<CuImage>>>(m, "CuImageIterator") //
         .def(py::init<std::shared_ptr<CuImage>, bool>(), doc::CuImageIterator::doc_CuImageIterator,
              py::arg("cuimg"), //
              py::arg("ending") = false, py::call_guard<py::gil_scoped_release>())
@@ -212,8 +212,8 @@ PYBIND11_MODULE(_cucim, m)
             py::call_guard<py::gil_scoped_release>())
         .def(
             "__iter__", //
-            [](CuImageIterator<CuImage>& it) { //
-                return CuImageIterator<CuImage>(it); //
+            [](CuImageIterator<CuImage>* it) { //
+                return it; //
             }, //
             py::call_guard<py::gil_scoped_release>())
         .def("__next__", &py_cuimage_iterator_next, py::call_guard<py::gil_scoped_release>())
@@ -517,11 +517,11 @@ py::object py_read_region(const CuImage& cuimg,
     auto loader = region_ptr->loader();
     if (batch_size > 1 || (loader && loader->size() > 1))
     {
-        auto iter_ptr = region_ptr->begin();
+        auto iter_ptr = std::make_shared<CuImageIterator<CuImage>>(region_ptr->shared_from_this());
 
         py::gil_scoped_acquire scope_guard;
 
-        py::object iter = py::cast(iter_ptr);
+        py::object iter = py::cast(iter_ptr, py::return_value_policy::take_ownership);
 
         return iter;
     }