Free resources in LocalExchangeSourceOperator::close

facebookincubator · Sep 23, 2021 · 5d23016 · 5d23016
1 parent 85bdb1f
commit 5d23016
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/velox/exec/LocalPartition.h b/velox/exec/LocalPartition.h
@@ -85,6 +85,12 @@ class LocalExchangeSource {
   /// copied into the consumers memory pool.
   BlockingReason isFinished(ContinueFuture* future);
 
+  void close() {
+    queue_.withWLock([] (auto& queue)) {
+      queue.clear();
+    }
+  }
+
  private:
   LocalExchangeMemoryManager* memoryManager_;
   const int partition_;
@@ -120,6 +126,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};