Free resources in LocalExchangeSourceOperator::close

facebookincubator · Sep 23, 2021 · 6d17dbc · 6d17dbc
1 parent 85bdb1f
commit 6d17dbc
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/velox/exec/LocalPartition.h b/velox/exec/LocalPartition.h
@@ -85,6 +85,14 @@ class LocalExchangeSource {
   /// copied into the consumers memory pool.
   BlockingReason isFinished(ContinueFuture* future);
 
+  void close() {
+    queue_.withWLock([](auto& queue) {
+      while (!queue.empty()) {
+        queue.pop();
+      }
+    });
+  }
+
  private:
   LocalExchangeMemoryManager* memoryManager_;
   const int partition_;
@@ -120,6 +128,11 @@ class LocalExchangeSourceOperator : public SourceOperator {
 
   RowVectorPtr getOutput() override;
 
+  void close() override {
+    Operator::close();
+    source_->close();
+  }
+
  private:
   const int partition_;
   const std::shared_ptr<LocalExchangeSource> source_{nullptr};