From 8130e287c2b96c85895e70d4fc74617c59d8ea2f Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Thu, 3 Mar 2022 17:46:04 -0800
Subject: [PATCH] Add timeout to can_grow_directly when waiting for the
 Condvar.

---
 datafusion/src/execution/memory_manager.rs | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/datafusion/src/execution/memory_manager.rs b/datafusion/src/execution/memory_manager.rs
index d39eaab3c215..a2c75f3fe3c7 100644
--- a/datafusion/src/execution/memory_manager.rs
+++ b/datafusion/src/execution/memory_manager.rs
@@ -19,6 +19,7 @@
 
 use crate::error::{DataFusionError, Result};
 use async_trait::async_trait;
+use core::time::Duration;
 use hashbrown::HashSet;
 use log::debug;
 use parking_lot::{Condvar, Mutex};
@@ -340,7 +341,13 @@ impl MemoryManager {
             } else if current < min_per_rqt {
                 // if we cannot acquire at lease 1/2n memory, just wait for others
                 // to spill instead spill self frequently with limited total mem
-                self.cv.wait(&mut rqt_current_used);
+                let timeout = self
+                    .cv
+                    .wait_for(&mut rqt_current_used, Duration::from_secs(5));
+                if timeout.timed_out() {
+                    granted = false;
+                    break;
+                }
             } else {
                 granted = false;
                 break;