Fix memory leak in ThreadBatchDataLoader

- Throw StopIteration exception early
- Do not set nullptr for 'raster_data_' in
  ThreadBatchDataLoader::next_data()
  - Setting nullptr prevents freeing memory in the destructor
    of ThreadBatchDataLoader

Signed-off-by: Gigon Bae <[email protected]>

cpp/src/loader/thread_batch_data_loader.cpp

@@ -210,19 +210,16 @@ uint8_t* ThreadBatchDataLoader::next_data()
 {
     if (num_workers_ == 0) // (location_len == 1 && batch_size == 1)
     {
-        // If it reads entire image with multi threads (using loader), release raster memory from batch data loader.
+        // If it reads entire image with multi threads (using loader), release raster memory from batch data loader
+        // by setting it to nullptr so that it will not be freed by ~ThreadBatchDataLoader (destructor).
         uint8_t* batch_raster_ptr = raster_data_[0];
         raster_data_[0] = nullptr;
         return batch_raster_ptr;
     }
 
     if (processed_batch_count_ * batch_size_ >= location_len_)
     {
-        // Remove buffer items that are no longer needed.
-        for (size_t i = 0; i < buffer_item_len_; ++i)
-        {
-            raster_data_[i] = nullptr;
-        }
+        // If all batches are processed, return nullptr.
         return nullptr;
     }
 

python/pybind11/cucim_py.cpp

@@ -557,6 +557,10 @@ py::object py_associated_image(const CuImage& cuimg, const std::string& name, co
 py::object py_cuimage_iterator_next(CuImageIterator<CuImage>& it)
 {
     bool stop_iteration = (it.index() == it.size());
+    if (stop_iteration)
+    {
+        throw py::stop_iteration();
+    }
 
     // Get the next batch of images.
     ++it;
@@ -573,10 +577,6 @@ py::object py_cuimage_iterator_next(CuImageIterator<CuImage>& it)
         {
             _set_array_interface(cuimg_obj);
         }
-        if (stop_iteration)
-        {
-            throw py::stop_iteration();
-        }
         return cuimg_obj;
     }
 }