resize queue and test

mindspore/ccsrc/minddata/dataset/engine/connector.h

@@ -138,7 +138,7 @@ class Connector {
   // starting from the beginning.
   void Reset() {
     for (int i = 0; i < queues_.size(); ++i) {
-      queues_[i]->ResetQue();
+      queues_[i]->Reset();
     }
     expect_consumer_ = 0;
     pop_from_ = 0;

mindspore/ccsrc/minddata/dataset/util/queue.h

@@ -66,7 +66,10 @@ class Queue {
 
   bool empty() const { return head_ == tail_; }
 
-  void Reset() { ResetQue(); }
+  void Reset() {
+    std::unique_lock<std::mutex> _lock(mux_);
+    ResetQue();
+  }
 
   // Producer
   Status Add(const_reference ele) noexcept {
@@ -74,8 +77,7 @@ class Queue {
     // Block when full
     Status rc = full_cv_.Wait(&_lock, [this]() -> bool { return (size() != capacity()); });
     if (rc.IsOk()) {
-      auto k = tail_++ % sz_;
-      *(arr_[k]) = ele;
+      RETURN_IF_NOT_OK(AddWhileHoldingLock(ele));
       empty_cv_.NotifyAll();
       _lock.unlock();
     } else {
@@ -121,8 +123,7 @@ class Queue {
     // Block when empty
     Status rc = empty_cv_.Wait(&_lock, [this]() -> bool { return !empty(); });
     if (rc.IsOk()) {
-      auto k = head_++ % sz_;
-      *p = std::move(*(arr_[k]));
+      RETURN_IF_NOT_OK(PopFrontWhileHoldingLock(p));
       full_cv_.NotifyAll();
       _lock.unlock();
     } else {
@@ -131,24 +132,6 @@ class Queue {
     return rc;
   }
 
-  void ResetQue() noexcept {
-    std::unique_lock<std::mutex> _lock(mux_);
-    // If there are elements in the queue, drain them. We won't call PopFront directly
-    // because we have got the lock already. We will deadlock if we call PopFront
-    for (auto i = head_; i < tail_; ++i) {
-      auto k = i % sz_;
-      auto val = std::move(*(arr_[k]));
-      // Let val go out of scope and its destructor will be invoked automatically.
-      // But our compiler may complain val is not in use. So let's do some useless
-      // stuff.
-      MS_LOG(DEBUG) << "Address of val: " << &val;
-    }
-    empty_cv_.ResetIntrpState();
-    full_cv_.ResetIntrpState();
-    head_ = 0;
-    tail_ = 0;
-  }
-
   Status Register(TaskGroup *vg) {
     Status rc1 = empty_cv_.Register(vg->GetIntrpService());
     Status rc2 = full_cv_.Register(vg->GetIntrpService());
@@ -159,6 +142,28 @@ class Queue {
     }
   }
 
+  Status Resize(int32_t new_capacity) {
+    std::unique_lock<std::mutex> _lock(mux_);
+    CHECK_FAIL_RETURN_UNEXPECTED(
+      new_capacity >= static_cast<int32_t>(size()),
+      "New capacity: " + std::to_string(new_capacity) + ", is smaller than queue size:" + std::to_string(size()));
+    std::vector<T> queue;
+    while (head_ < tail_) {
+      T temp;
+      RETURN_IF_NOT_OK(this->PopFrontWhileHoldingLock(&temp));
+      queue.push_back(temp);
+    }
+    this->ResetQue();
+    RETURN_IF_NOT_OK(arr_.allocate(new_capacity));
+    sz_ = new_capacity;
+    for (int i = 0; i < queue.size(); ++i) {
+      RETURN_IF_NOT_OK(this->AddWhileHoldingLock(queue[i]));
+    }
+    queue.clear();
+    _lock.unlock();
+    return Status::OK();
+  }
+
  private:
   size_t sz_;
   MemGuard<T, Allocator<T>> arr_;
@@ -168,6 +173,32 @@ class Queue {
   std::mutex mux_;
   CondVar empty_cv_;
   CondVar full_cv_;
+
+  // Helper function for Add, must be called when holding a lock
+  Status AddWhileHoldingLock(const_reference ele) {
+    auto k = tail_++ % sz_;
+    *(arr_[k]) = ele;
+    return Status::OK();
+  }
+
+  // Helper function for PopFront, must be called when holding a lock
+  Status PopFrontWhileHoldingLock(pointer p) {
+    auto k = head_++ % sz_;
+    *p = std::move(*(arr_[k]));
+    return Status::OK();
+  }
+
+  void ResetQue() noexcept {
+    while (head_ < tail_) {
+      T val;
+      this->PopFrontWhileHoldingLock(&val);
+      MS_LOG(DEBUG) << "Address of val: " << &val;
+    }
+    empty_cv_.ResetIntrpState();
+    full_cv_.ResetIntrpState();
+    head_ = 0;
+    tail_ = 0;
+  }
 };
 
 // A container of queues with [] operator accessors.  Basically this is a wrapper over of a vector of queues

tests/ut/cpp/dataset/queue_test.cc

@@ -177,3 +177,51 @@ TEST_F(MindDataTestQueue, Test6) {
   MS_LOG(INFO) << "Popped value " << *pepped_value << " from queue index " << chosen_queue_index;
   ASSERT_EQ(*pepped_value, 99);
 }
+
+// Feature: Check resize is finished without changing elements and influencing operations.
+// Description: Compare elements in queue before and after resize, and test add/pop/reset.
+// Expectation: Elements in queue after resize are the same as the original queue.
+TEST_F(MindDataTestQueue, TestResize) {
+  // Create a list of queues with capacity = 3
+  Queue<TensorRow> queue(3);
+  ASSERT_EQ(3, queue.capacity());
+  // Add 3 rows into the queue
+  TensorRow a;
+  std::shared_ptr<Tensor> test_tensor1;
+  std::vector<float> input = {1.1, 0.2, 0.3, 0.4, 0.5, 0.6, 1.2, 0.7, 0.8, 0.9, 1.0, 2.0, 1.3, 3.0, 4.0};
+  ASSERT_OK(Tensor::CreateFromVector(input, TensorShape{3, 5}, &test_tensor1));
+  a.push_back(test_tensor1);
+  EXPECT_OK(queue.Add(a));
+
+  TensorRow b;
+  std::shared_ptr<Tensor> test_tensor2;
+  ASSERT_OK(Tensor::CreateScalar(true, &test_tensor2));
+  b.push_back(test_tensor2);
+  EXPECT_OK(queue.Add(b));
+
+  TensorRow c;
+  std::shared_ptr<Tensor> test_tensor3;
+  ASSERT_OK(Tensor::CreateFromVector(input, &test_tensor3));
+  c.push_back(test_tensor3);
+  EXPECT_OK(queue.Add(c));
+  ASSERT_EQ(3, queue.size());
+  // Check false if the resize is smaller than current size
+  EXPECT_ERROR(queue.Resize(2));
+  // Check true if the resize is larger than current size, and capacity is changed
+  EXPECT_OK(queue.Resize(12));
+  ASSERT_EQ(12, queue.capacity());
+  TensorRow d = a;
+  EXPECT_OK(queue.Add(d));
+  ASSERT_EQ(4, queue.size());
+  // Expect the rows after resize are the same as original input
+  TensorRow e;
+  EXPECT_OK(queue.PopFront(&e));
+  EXPECT_EQ(a.getRow(), e.getRow());
+  EXPECT_OK(queue.PopFront(&e));
+  EXPECT_EQ(b.getRow(), e.getRow());
+  EXPECT_OK(queue.PopFront(&e));
+  EXPECT_EQ(c.getRow(), e.getRow());
+  ASSERT_EQ(1, queue.size());
+  queue.Reset();
+  ASSERT_EQ(0, queue.size());
+}