From aedf6b0e49b55111ed03b7342cc9e8968e594c66 Mon Sep 17 00:00:00 2001
From: Rui Mo <rui@apache.org>
Date: Sat, 22 Nov 2025 02:12:45 +0800
Subject: [PATCH] [OAP][14722] Fix memory leak caused by asynchronous prefetch

---
 velox/common/caching/AsyncDataCache.h   |  2 +-
 velox/dwio/common/DirectBufferedInput.h | 11 +++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/velox/common/caching/AsyncDataCache.h b/velox/common/caching/AsyncDataCache.h
index 0baf904deb03..a7bc30951912 100644
--- a/velox/common/caching/AsyncDataCache.h
+++ b/velox/common/caching/AsyncDataCache.h
@@ -445,7 +445,7 @@ class CoalescedLoad {
     return state_;
   }
 
-  void cancel() {
+  virtual void cancel() {
     setEndState(State::kCancelled);
   }
 
diff --git a/velox/dwio/common/DirectBufferedInput.h b/velox/dwio/common/DirectBufferedInput.h
index ae1ff124ef73..1afcf93d979f 100644
--- a/velox/dwio/common/DirectBufferedInput.h
+++ b/velox/dwio/common/DirectBufferedInput.h
@@ -102,6 +102,17 @@ class DirectCoalescedLoad : public cache::CoalescedLoad {
     return size;
   }
 
+  void cancel() override {
+    folly::SemiFuture<bool> waitFuture(false);
+    if (state() == State::kLoading && !loadOrFuture(&waitFuture)) {
+      waitFuture.wait();
+    }
+    for (auto& request : requests_) {
+      pool_->freeNonContiguous(request.data);
+    }
+    CoalescedLoad::cancel();
+  }
+
  private:
   const std::shared_ptr<IoStatistics> ioStats_;
   const std::shared_ptr<filesystems::File::IoStats> fsStats_;