Fix double free in loader destructor and race condition in consumer/producer (#678)

Signed-off-by: Joaquin Anton <[email protected]>

Author: jantonguirao

dali/pipeline/operators/reader/caffe2_reader_op.h

@@ -30,13 +30,7 @@ class Caffe2Reader : public DataReader<CPUBackend, Tensor<CPUBackend>> {
   }
 
   void RunImpl(SampleWorkspace* ws, const int i) override {
-    const int idx = ws->data_idx();
-
-    auto* raw_data = GetSample(idx);
-
-    parser_->Parse(*raw_data, ws);
-
-    return;
+    parser_->Parse(GetSample(ws->data_idx()), ws);
   }
 
  protected:

dali/pipeline/operators/reader/caffe_reader_op.h

@@ -30,13 +30,7 @@ class CaffeReader : public DataReader<CPUBackend, Tensor<CPUBackend>> {
   }
 
   void RunImpl(SampleWorkspace* ws, const int i) override {
-    const int idx = ws->data_idx();
-
-    auto* raw_data = GetSample(idx);
-
-    parser_->Parse(*raw_data, ws);
-
-    return;
+    parser_->Parse(GetSample(ws->data_idx()), ws);
   }
 
  protected:

dali/pipeline/operators/reader/coco_reader_op.h

@@ -61,13 +61,7 @@ class COCOReader : public DataReader<CPUBackend, ImageLabelWrapper> {
   }
 
   void RunImpl(SampleWorkspace* ws, const int i) override {
-    const int idx = ws->data_idx();
-
-    auto* image_label = GetSample(idx);
-
-    parser_->Parse(*image_label, ws);
-
-    return;
+    parser_->Parse(GetSample(ws->data_idx()), ws);
   }
 
  protected:

dali/pipeline/operators/reader/file_reader_op.h

@@ -30,25 +30,24 @@ class FileReader : public DataReader<CPUBackend, ImageLabelWrapper> {
   void RunImpl(SampleWorkspace *ws, const int i) override {
     const int idx = ws->data_idx();
 
-    auto* image_label = GetSample(idx);
+    const auto& image_label = GetSample(idx);
 
     // copy from raw_data -> outputs directly
     auto &image_output = ws->Output<CPUBackend>(0);
     auto &label_output = ws->Output<CPUBackend>(1);
 
-    Index image_size = image_label->image.size();
+    Index image_size = image_label.image.size();
 
     image_output.Resize({image_size});
     image_output.mutable_data<uint8_t>();
     label_output.Resize({1});
 
     std::memcpy(image_output.raw_mutable_data(),
-                image_label->image.raw_data(),
+                image_label.image.raw_data(),
                 image_size);
-    image_output.SetSourceInfo(image_label->image.GetSourceInfo());
+    image_output.SetSourceInfo(image_label.image.GetSourceInfo());
 
-    label_output.mutable_data<int>()[0] = image_label->label;
-    return;
+    label_output.mutable_data<int>()[0] = image_label.label;
   }
 
  protected:

dali/pipeline/operators/reader/loader/file_loader.cc

@@ -90,11 +90,11 @@ vector<std::pair<string, int>> filesystem::traverse_directories(const std::strin
   return file_label_pairs;
 }
 
-void FileLoader::PrepareEmpty(ImageLabelWrapper *image_label) {
-  PrepareEmptyTensor(&image_label->image);
+void FileLoader::PrepareEmpty(ImageLabelWrapper &image_label) {
+  PrepareEmptyTensor(image_label.image);
 }
 
-void FileLoader::ReadSample(ImageLabelWrapper* image_label) {
+void FileLoader::ReadSample(ImageLabelWrapper &image_label) {
   auto image_pair = image_label_pairs_[current_index_++];
 
   // handle wrap-around
@@ -104,25 +104,22 @@ void FileLoader::ReadSample(ImageLabelWrapper* image_label) {
   Index image_size = current_image->Size();
 
   if (copy_read_data_) {
-    image_label->image.Resize({image_size});
+    image_label.image.Resize({image_size});
     // copy the image
-    current_image->Read(image_label->image.mutable_data<uint8_t>(), image_size);
+    current_image->Read(image_label.image.mutable_data<uint8_t>(), image_size);
   } else {
     auto p = current_image->Get(image_size);
     // Wrap the raw data in the Tensor object.
-    image_label->image.ShareData(p, image_size, {image_size});
-
-    TypeInfo type;
-    type.SetType<uint8_t>();
-    image_label->image.set_type(type);
+    image_label.image.ShareData(p, image_size, {image_size});
+    image_label.image.set_type(TypeInfo::Create<uint8_t>());
   }
 
-  image_label->image.SetSourceInfo(image_pair.first);
+  image_label.image.SetSourceInfo(image_pair.first);
   // close the file handle
   current_image->Close();
 
   // copy the label
-  image_label->label = image_pair.second;
+  image_label.label = image_pair.second;
 }
 
 Index FileLoader::Size() {

dali/pipeline/operators/reader/loader/file_loader.h

@@ -89,8 +89,8 @@ class FileLoader : public Loader<CPUBackend, ImageLabelWrapper> {
     Reset(true);
   }
 
-  void PrepareEmpty(ImageLabelWrapper *tensor) override;
-  void ReadSample(ImageLabelWrapper *tensor) override;
+  void PrepareEmpty(ImageLabelWrapper &tensor) override;
+  void ReadSample(ImageLabelWrapper &tensor) override;
 
   Index Size() override;
 

dali/pipeline/operators/reader/loader/indexed_file_loader.h

@@ -37,7 +37,7 @@ class IndexedFileLoader : public Loader<CPUBackend, Tensor<CPUBackend>> {
         Init(options);
     }
 
-  void ReadSample(Tensor<CPUBackend>* tensor) override {
+  void ReadSample(Tensor<CPUBackend>& tensor) override {
     MoveToNextShard(current_index_);
 
     int64 seek_pos, size;
@@ -53,21 +53,18 @@ class IndexedFileLoader : public Loader<CPUBackend, Tensor<CPUBackend>> {
       auto p = current_file_->Get(size);
       DALI_ENFORCE(p != nullptr, "Error reading from a file " + uris_[current_file_index_]);
       // Wrap the raw data in the Tensor object.
-      tensor->ShareData(p, size, {size});
-
-      TypeInfo type;
-      type.SetType<uint8_t>();
-      tensor->set_type(type);
+      tensor.ShareData(p, size, {size});
+      tensor.set_type(TypeInfo::Create<uint8_t>());
     } else {
-      tensor->Resize({size});
-      tensor->mutable_data<uint8_t>();
+      tensor.set_type(TypeInfo::Create<uint8_t>());
+      tensor.Resize({size});
 
-      int64 n_read = current_file_->Read(reinterpret_cast<uint8_t*>(tensor->raw_mutable_data()),
+      int64 n_read = current_file_->Read(reinterpret_cast<uint8_t*>(tensor.raw_mutable_data()),
                           size);
       DALI_ENFORCE(n_read == size, "Error reading from a file " + uris_[current_file_index_]);
     }
 
-    tensor->SetSourceInfo(uris_[current_file_index_] + " at index " + to_string(seek_pos));
+    tensor.SetSourceInfo(uris_[current_file_index_] + " at index " + to_string(seek_pos));
     ++current_index_;
     return;
   }

dali/pipeline/operators/reader/loader/lmdb.h

@@ -91,18 +91,18 @@ class LMDBReader : public Loader<CPUBackend, Tensor<CPUBackend>> {
     mdb_env_ = nullptr;
   }
 
-  void ReadSample(Tensor<CPUBackend>* tensor) override {
+  void ReadSample(Tensor<CPUBackend>& tensor) override {
     // assume cursor is valid, read next, loop to start if necessary
     lmdb::SeekLMDB(mdb_cursor_, MDB_NEXT, &key_, &value_);
     ++current_index_;
 
     MoveToNextShard(current_index_);
 
-    tensor->Resize({static_cast<Index>(value_.mv_size)});
-    tensor->mutable_data<uint8_t>();
-    tensor->SetSourceInfo(db_path_ + " at key " + to_string(reinterpret_cast<char*>(key_.mv_data)));
+    tensor.Resize({static_cast<Index>(value_.mv_size)});
+    tensor.set_type(TypeInfo::Create<uint8_t>());
+    tensor.SetSourceInfo(db_path_ + " at key " + to_string(reinterpret_cast<char*>(key_.mv_data)));
 
-    std::memcpy(tensor->raw_mutable_data(),
+    std::memcpy(tensor.raw_mutable_data(),
                 reinterpret_cast<uint8_t*>(value_.mv_data),
                 value_.mv_size*sizeof(uint8_t));
 

dali/pipeline/operators/reader/loader/loader.h

@@ -15,14 +15,14 @@
 #ifndef DALI_PIPELINE_OPERATORS_READER_LOADER_LOADER_H_
 #define DALI_PIPELINE_OPERATORS_READER_LOADER_LOADER_H_
 
-#include <list>
 #include <map>
+#include <memory>
 #include <mutex>
 #include <random>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <vector>
-#include <type_traits>
 
 #include "dali/common.h"
 #include "dali/error_handling.h"
@@ -44,7 +44,7 @@ DLL_PUBLIC size_t start_index(const size_t shard_id,
 template <typename Backend, typename LoadTarget>
 class Loader {
  public:
-  using LoadTarget_t = LoadTarget;
+  using LoadTargetPtr = std::unique_ptr<LoadTarget>;
   explicit Loader(const OpSpec& options)
     : shuffle_(options.GetArgument<bool>("random_shuffle")),
       initial_buffer_fill_(shuffle_ ? options.GetArgument<int>("initial_fill") : 1),
@@ -65,105 +65,93 @@ class Loader {
     e_ = std::default_random_engine(seq);
   }
 
-  virtual ~Loader() {
-    // delete all the temporary tensors
-    while (!sample_buffer_.empty()) {
-      LoadTarget * t = sample_buffer_.back();
-      delete t;
-      sample_buffer_.pop_back();
-    }
-    while (!empty_tensors_.empty()) {
-      LoadTarget * t = empty_tensors_.back();
-      delete t;
-      empty_tensors_.pop_back();
-    }
-  }
+  virtual ~Loader() = default;
 
-  virtual void PrepareEmpty(LoadTarget *tensor) {
+  virtual void PrepareEmpty(LoadTarget& tensor) {
     PrepareEmptyTensor(tensor);
   }
 
   template <typename T>
   typename std::enable_if<std::is_same<T, Tensor<CPUBackend>>::value>::type
-  PrepareEmptyTensor(T *tensor) {
-    tensor->set_pinned(false);
+  PrepareEmptyTensor(T& tensor) {
+    tensor.set_pinned(false);
     // Initialize tensors to a set size to limit expensive reallocations
-    tensor->Resize({tensor_init_bytes_});
-    tensor->template mutable_data<uint8_t>();
+    tensor.Resize({tensor_init_bytes_});
+    tensor.template mutable_data<uint8_t>();
   }
 
   template <typename T>
   typename std::enable_if<!std::is_same<T, Tensor<CPUBackend>>::value>::type
-  PrepareEmptyTensor(T *) {
+  PrepareEmptyTensor(T&) {
     constexpr bool T_is_Tensor = std::is_same<T, Tensor<CPUBackend>>::value;
     DALI_ENFORCE(T_is_Tensor,
       "Please overload PrepareEmpty for custom LoadTarget type other than Tensor");
   }
 
 
   // Get a random read sample
-  LoadTarget* ReadOne() {
+  LoadTargetPtr ReadOne() {
     TimeRange tr("[Loader] ReadOne", TimeRange::kGreen1);
     // perform an iniital buffer fill if it hasn't already happened
     if (!initial_buffer_filled_) {
       TimeRange tr("[Loader] Filling initial buffer", TimeRange::kBlue1);
+
       // Read an initial number of samples to fill our
       // sample buffer
       for (int i = 0; i < initial_buffer_fill_; ++i) {
-        LoadTarget* tensor = new LoadTarget();
-        PrepareEmpty(tensor);
-
-        ReadSample(tensor);
-        sample_buffer_.push_back(tensor);
+        auto tensor_ptr = LoadTargetPtr(new LoadTarget());
+        PrepareEmpty(*tensor_ptr);
+        ReadSample(*tensor_ptr);
+        sample_buffer_.push_back(std::move(tensor_ptr));
       }
 
-      TimeRange tr2("[Loader] Filling empty list", TimeRange::kOrange);
       // need some entries in the empty_tensors_ list
+      TimeRange tr2("[Loader] Filling empty list", TimeRange::kOrange);
+      std::lock_guard<std::mutex> lock(empty_tensors_mutex_);
       for (int i = 0; i < initial_empty_size_; ++i) {
-        LoadTarget* tensor = new LoadTarget();
-        PrepareEmpty(tensor);
-
-        empty_tensors_.push_back(tensor);
+        auto tensor_ptr = LoadTargetPtr(new LoadTarget());
+        PrepareEmpty(*tensor_ptr);
+        empty_tensors_.push_back(std::move(tensor_ptr));
       }
 
       initial_buffer_filled_ = true;
     }
     // choose the random index
     int idx = shuffle_ ? dis(e_) % sample_buffer_.size() : 0;
-    LoadTarget* elem = sample_buffer_[idx];
 
     // swap end and idx, return the tensor to empties
-    std::swap(sample_buffer_[idx], sample_buffer_[sample_buffer_.size()-1]);
+    std::swap(sample_buffer_[idx], sample_buffer_.back());
     // remove last element
+    LoadTargetPtr sample_ptr = std::move(sample_buffer_.back());
     sample_buffer_.pop_back();
 
     // now grab an empty tensor, fill it and add to filled buffers
-    // empty_tensors_ needs to be thread-safe w.r.t. ReturnTensor()
+    // empty_tensors_ needs to be thread-safe w.r.t. RecycleTensor()
     // being called by multiple consumer threads
-    LoadTarget* t;
+    LoadTargetPtr tensor_ptr;
     {
-      std::lock_guard<std::mutex> lock(return_mutex_);
+      std::lock_guard<std::mutex> lock(empty_tensors_mutex_);
       DALI_ENFORCE(empty_tensors_.size() > 0, "No empty tensors - did you forget to return them?");
-      t = empty_tensors_.back();
+      tensor_ptr = std::move(empty_tensors_.back());
       empty_tensors_.pop_back();
     }
-    ReadSample(t);
-    sample_buffer_.push_back(t);
+    ReadSample(*tensor_ptr);
+    sample_buffer_.push_back(std::move(tensor_ptr));
 
-    return elem;
+    return sample_ptr;
   }
 
   // return a tensor to the empty pile
   // called by multiple consumer threads
-  void ReturnTensor(LoadTarget* tensor) {
-    std::lock_guard<std::mutex> lock(return_mutex_);
-    empty_tensors_.push_back(tensor);
+  void RecycleTensor(LoadTargetPtr&& tensor_ptr) {
+    std::lock_guard<std::mutex> lock(empty_tensors_mutex_);
+    empty_tensors_.push_back(std::move(tensor_ptr));
   }
 
   // Read an actual sample from the FileStore,
   // used to populate the sample buffer for "shuffled"
   // reads.
-  virtual void ReadSample(LoadTarget* tensor) = 0;
+  virtual void ReadSample(LoadTarget& tensor) = 0;
 
   // Give the size of the data accessed through the Loader
   virtual Index Size() = 0;
@@ -183,9 +171,9 @@ class Loader {
             (stick_to_shard_ && shard_id_ + 1 < num_shards_ &&
             current_index >= static_cast<Index>(start_index(shard_id_ + 1, num_shards_, Size())));
   }
-  std::vector<LoadTarget*> sample_buffer_;
+  std::vector<LoadTargetPtr> sample_buffer_;
 
-  std::list<LoadTarget*> empty_tensors_;
+  std::vector<LoadTargetPtr> empty_tensors_;
 
   // number of samples to initialize buffer with
   // ~1 minibatch seems reasonable
@@ -201,7 +189,7 @@ class Loader {
   Index seed_;
 
   // control return of tensors
-  std::mutex return_mutex_;
+  std::mutex empty_tensors_mutex_;
 
   // sharding
   const int shard_id_;