feat: refine the MemoryGuard (#7466)

* feat: refine MemoryGuard

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* chore: add test

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

.gitignore

@@ -67,3 +67,6 @@ greptimedb_data
 
 # Claude code
 CLAUDE.md
+
+# AGENTS.md
+AGENTS.md

src/common/memory-manager/src/guard.rs

@@ -15,9 +15,14 @@
 use std::{fmt, mem};
 
 use common_telemetry::debug;
+use snafu::ensure;
 use tokio::sync::{OwnedSemaphorePermit, TryAcquireError};
 
+use crate::error::{
+    MemoryAcquireTimeoutSnafu, MemoryLimitExceededSnafu, MemorySemaphoreClosedSnafu, Result,
+};
 use crate::manager::{MemoryMetrics, MemoryQuota};
+use crate::policy::OnExhaustedPolicy;
 
 /// Guard representing a slice of reserved memory.
 pub struct MemoryGuard<M: MemoryMetrics> {
@@ -55,11 +60,52 @@ impl<M: MemoryMetrics> MemoryGuard<M> {
         }
     }
 
-    /// Tries to allocate additional memory during task execution.
+    /// Acquires additional memory, waiting if necessary until enough is available.
+    ///
+    /// On success, merges the new memory into this guard.
+    ///
+    /// # Errors
+    /// - Returns error if requested bytes would exceed the manager's total limit
+    /// - Returns error if the semaphore is unexpectedly closed
+    pub async fn acquire_additional(&mut self, bytes: u64) -> Result<()> {
+        match &mut self.state {
+            GuardState::Unlimited => Ok(()),
+            GuardState::Limited { permit, quota } => {
+                if bytes == 0 {
+                    return Ok(());
+                }
+
+                let additional_permits = quota.bytes_to_permits(bytes);
+                let current_permits = permit.num_permits() as u32;
+
+                ensure!(
+                    current_permits.saturating_add(additional_permits) <= quota.limit_permits,
+                    MemoryLimitExceededSnafu {
+                        requested_bytes: bytes,
+                        limit_bytes: quota.permits_to_bytes(quota.limit_permits)
+                    }
+                );
+
+                let additional_permit = quota
+                    .semaphore
+                    .clone()
+                    .acquire_many_owned(additional_permits)
+                    .await
+                    .map_err(|_| MemorySemaphoreClosedSnafu.build())?;
+
+                permit.merge(additional_permit);
+                quota.update_in_use_metric();
+                debug!("Acquired additional {} bytes", bytes);
+                Ok(())
+            }
+        }
+    }
+
+    /// Tries to acquire additional memory without waiting.
     ///
     /// On success, merges the new memory into this guard and returns true.
     /// On failure, returns false and leaves this guard unchanged.
-    pub fn request_additional(&mut self, bytes: u64) -> bool {
+    pub fn try_acquire_additional(&mut self, bytes: u64) -> bool {
         match &mut self.state {
             GuardState::Unlimited => true,
             GuardState::Limited { permit, quota } => {
@@ -77,11 +123,11 @@ impl<M: MemoryMetrics> MemoryGuard<M> {
                     Ok(additional_permit) => {
                         permit.merge(additional_permit);
                         quota.update_in_use_metric();
-                        debug!("Allocated additional {} bytes", bytes);
+                        debug!("Acquired additional {} bytes", bytes);
                         true
                     }
                     Err(TryAcquireError::NoPermits) | Err(TryAcquireError::Closed) => {
-                        quota.metrics.inc_rejected("request_additional");
+                        quota.metrics.inc_rejected("try_acquire_additional");
                         false
                     }
                 }
@@ -89,11 +135,55 @@ impl<M: MemoryMetrics> MemoryGuard<M> {
         }
     }
 
-    /// Releases a portion of granted memory back to the pool early,
-    /// before the guard is dropped.
+    /// Acquires additional memory based on the given policy.
+    ///
+    /// - For `OnExhaustedPolicy::Wait`: Waits up to the timeout duration for memory to become available
+    /// - For `OnExhaustedPolicy::Fail`: Returns immediately if memory is not available
+    ///
+    /// # Errors
+    /// - `MemoryLimitExceeded`: Requested bytes would exceed the total limit (both policies), or memory is currently exhausted (Fail policy only)
+    /// - `MemoryAcquireTimeout`: Timeout elapsed while waiting for memory (Wait policy only)
+    /// - `MemorySemaphoreClosed`: The internal semaphore is unexpectedly closed (rare, indicates system issue)
+    pub async fn acquire_additional_with_policy(
+        &mut self,
+        bytes: u64,
+        policy: OnExhaustedPolicy,
+    ) -> Result<()> {
+        match policy {
+            OnExhaustedPolicy::Wait { timeout } => {
+                match tokio::time::timeout(timeout, self.acquire_additional(bytes)).await {
+                    Ok(Ok(())) => Ok(()),
+                    Ok(Err(e)) => Err(e),
+                    Err(_elapsed) => MemoryAcquireTimeoutSnafu {
+                        requested_bytes: bytes,
+                        waited: timeout,
+                    }
+                    .fail(),
+                }
+            }
+            OnExhaustedPolicy::Fail => {
+                if self.try_acquire_additional(bytes) {
+                    Ok(())
+                } else {
+                    MemoryLimitExceededSnafu {
+                        requested_bytes: bytes,
+                        limit_bytes: match &self.state {
+                            GuardState::Unlimited => 0, // unreachable: unlimited mode always succeeds
+                            GuardState::Limited { quota, .. } => {
+                                quota.permits_to_bytes(quota.limit_permits)
+                            }
+                        },
+                    }
+                    .fail()
+                }
+            }
+        }
+    }
+
+    /// Releases a portion of granted memory back to the pool before the guard is dropped.
     ///
     /// Returns true if the release succeeds or is a no-op; false if the request exceeds granted.
-    pub fn early_release_partial(&mut self, bytes: u64) -> bool {
+    pub fn release_partial(&mut self, bytes: u64) -> bool {
         match &mut self.state {
             GuardState::Unlimited => true,
             GuardState::Limited { permit, quota } => {
@@ -109,7 +199,7 @@ impl<M: MemoryMetrics> MemoryGuard<M> {
                             quota.permits_to_bytes(released_permit.num_permits() as u32);
                         drop(released_permit);
                         quota.update_in_use_metric();
-                        debug!("Early released {} bytes from memory guard", released_bytes);
+                        debug!("Released {} bytes from memory guard", released_bytes);
                         true
                     }
                     None => false,

src/common/memory-manager/src/tests.rs

@@ -83,7 +83,7 @@ fn test_request_additional_success() {
     assert_eq!(manager.used_bytes(), base);
 
     // Request additional memory (3MB) - should succeed and merge
-    assert!(guard.request_additional(3 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.try_acquire_additional(3 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 8 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 8 * PERMIT_GRANULARITY_BYTES);
 }
@@ -98,11 +98,11 @@ fn test_request_additional_exceeds_limit() {
     let mut guard = manager.try_acquire(base).unwrap();
 
     // Request additional memory (3MB) - should succeed
-    assert!(guard.request_additional(3 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.try_acquire_additional(3 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(manager.used_bytes(), 8 * PERMIT_GRANULARITY_BYTES);
 
     // Request more (3MB) - should fail (would exceed 10MB limit)
-    let result = guard.request_additional(3 * PERMIT_GRANULARITY_BYTES);
+    let result = guard.try_acquire_additional(3 * PERMIT_GRANULARITY_BYTES);
     assert!(!result);
 
     // Still at 8MB
@@ -119,7 +119,7 @@ fn test_request_additional_auto_release_on_guard_drop() {
         let mut guard = manager.try_acquire(5 * PERMIT_GRANULARITY_BYTES).unwrap();
 
         // Request additional - memory is merged into guard
-        assert!(guard.request_additional(3 * PERMIT_GRANULARITY_BYTES));
+        assert!(guard.try_acquire_additional(3 * PERMIT_GRANULARITY_BYTES));
         assert_eq!(manager.used_bytes(), 8 * PERMIT_GRANULARITY_BYTES);
 
         // When guard drops, all memory (base + additional) is released together
@@ -135,7 +135,7 @@ fn test_request_additional_unlimited() {
     let mut guard = manager.try_acquire(5 * PERMIT_GRANULARITY_BYTES).unwrap();
 
     // Should always succeed with unlimited manager
-    assert!(guard.request_additional(100 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.try_acquire_additional(100 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 0);
     assert_eq!(manager.used_bytes(), 0);
 }
@@ -148,7 +148,7 @@ fn test_request_additional_zero_bytes() {
     let mut guard = manager.try_acquire(5 * PERMIT_GRANULARITY_BYTES).unwrap();
 
     // Request 0 bytes should succeed without affecting anything
-    assert!(guard.request_additional(0));
+    assert!(guard.try_acquire_additional(0));
     assert_eq!(guard.granted_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
 }
@@ -162,7 +162,7 @@ fn test_early_release_partial_success() {
     assert_eq!(manager.used_bytes(), 8 * PERMIT_GRANULARITY_BYTES);
 
     // Release half
-    assert!(guard.early_release_partial(4 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.release_partial(4 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 4 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 4 * PERMIT_GRANULARITY_BYTES);
 
@@ -177,7 +177,7 @@ fn test_early_release_partial_exceeds_granted() {
     let mut guard = manager.try_acquire(5 * PERMIT_GRANULARITY_BYTES).unwrap();
 
     // Try to release more than granted - should fail
-    assert!(!guard.early_release_partial(10 * PERMIT_GRANULARITY_BYTES));
+    assert!(!guard.release_partial(10 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
 }
@@ -188,7 +188,7 @@ fn test_early_release_partial_unlimited() {
     let mut guard = manager.try_acquire(100 * PERMIT_GRANULARITY_BYTES).unwrap();
 
     // Unlimited guard - release should succeed (no-op)
-    assert!(guard.early_release_partial(50 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.release_partial(50 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 0);
 }
 
@@ -200,22 +200,22 @@ fn test_request_and_early_release_symmetry() {
     let mut guard = manager.try_acquire(5 * PERMIT_GRANULARITY_BYTES).unwrap();
 
     // Request additional
-    assert!(guard.request_additional(5 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.try_acquire_additional(5 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 10 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 10 * PERMIT_GRANULARITY_BYTES);
 
     // Early release some
-    assert!(guard.early_release_partial(3 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.release_partial(3 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 7 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 7 * PERMIT_GRANULARITY_BYTES);
 
     // Request again
-    assert!(guard.request_additional(2 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.try_acquire_additional(2 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 9 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 9 * PERMIT_GRANULARITY_BYTES);
 
     // Early release again
-    assert!(guard.early_release_partial(4 * PERMIT_GRANULARITY_BYTES));
+    assert!(guard.release_partial(4 * PERMIT_GRANULARITY_BYTES));
     assert_eq!(guard.granted_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
     assert_eq!(manager.used_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
 
@@ -226,25 +226,186 @@ fn test_request_and_early_release_symmetry() {
 #[test]
 fn test_small_allocation_rounds_up() {
     // Test that allocations smaller than PERMIT_GRANULARITY_BYTES
-    // round up to 1 permit and can use request_additional()
+    // round up to 1 permit and can use try_acquire_additional()
     let limit = 10 * PERMIT_GRANULARITY_BYTES;
     let manager = MemoryManager::new(limit, NoOpMetrics);
 
     let mut guard = manager.try_acquire(512 * 1024).unwrap(); // 512KB
     assert_eq!(guard.granted_bytes(), PERMIT_GRANULARITY_BYTES); // Rounds up to 1MB
-    assert!(guard.request_additional(2 * PERMIT_GRANULARITY_BYTES)); // Can request more
+    assert!(guard.try_acquire_additional(2 * PERMIT_GRANULARITY_BYTES)); // Can request more
     assert_eq!(guard.granted_bytes(), 3 * PERMIT_GRANULARITY_BYTES);
 }
 
 #[test]
 fn test_acquire_zero_bytes_lazy_allocation() {
-    // Test that acquire(0) returns 0 permits but can request_additional() later
+    // Test that acquire(0) returns 0 permits but can try_acquire_additional() later
     let manager = MemoryManager::new(10 * PERMIT_GRANULARITY_BYTES, NoOpMetrics);
 
     let mut guard = manager.try_acquire(0).unwrap();
     assert_eq!(guard.granted_bytes(), 0); // No permits consumed
     assert_eq!(manager.used_bytes(), 0);
 
-    assert!(guard.request_additional(3 * PERMIT_GRANULARITY_BYTES)); // Lazy allocation
+    assert!(guard.try_acquire_additional(3 * PERMIT_GRANULARITY_BYTES)); // Lazy allocation
     assert_eq!(guard.granted_bytes(), 3 * PERMIT_GRANULARITY_BYTES);
 }
+
+#[tokio::test(flavor = "current_thread")]
+async fn test_acquire_additional_blocks_and_unblocks() {
+    let limit = 10 * PERMIT_GRANULARITY_BYTES;
+    let manager = MemoryManager::new(limit, NoOpMetrics);
+
+    // First guard takes 9MB, leaving only 1MB available
+    let mut guard1 = manager.try_acquire(9 * PERMIT_GRANULARITY_BYTES).unwrap();
+    assert_eq!(manager.used_bytes(), 9 * PERMIT_GRANULARITY_BYTES);
+
+    // Spawn a task that will block trying to acquire additional 5MB (needs total 10MB available)
+    let manager_clone = manager.clone();
+    let waiter = tokio::spawn(async move {
+        let mut guard2 = manager_clone.try_acquire(0).unwrap();
+        // This will block until enough memory is available
+        guard2
+            .acquire_additional(5 * PERMIT_GRANULARITY_BYTES)
+            .await
+            .unwrap();
+        guard2
+    });
+
+    sleep(Duration::from_millis(10)).await;
+
+    // Release 5MB from guard1 - this should unblock the waiter
+    assert!(guard1.release_partial(5 * PERMIT_GRANULARITY_BYTES));
+
+    // Waiter should complete now
+    let guard2 = waiter.await.unwrap();
+    assert_eq!(guard2.granted_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
+
+    // Total: guard1 has 4MB, guard2 has 5MB = 9MB
+    assert_eq!(manager.used_bytes(), 9 * PERMIT_GRANULARITY_BYTES);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn test_acquire_additional_exceeds_total_limit() {
+    let limit = 10 * PERMIT_GRANULARITY_BYTES;
+    let manager = MemoryManager::new(limit, NoOpMetrics);
+
+    let mut guard = manager.try_acquire(8 * PERMIT_GRANULARITY_BYTES).unwrap();
+
+    // Try to acquire additional 5MB - would exceed total limit of 10MB
+    let result = guard.acquire_additional(5 * PERMIT_GRANULARITY_BYTES).await;
+    assert!(result.is_err());
+
+    // Guard should remain unchanged
+    assert_eq!(guard.granted_bytes(), 8 * PERMIT_GRANULARITY_BYTES);
+    assert_eq!(manager.used_bytes(), 8 * PERMIT_GRANULARITY_BYTES);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn test_acquire_additional_success() {
+    let limit = 10 * PERMIT_GRANULARITY_BYTES;
+    let manager = MemoryManager::new(limit, NoOpMetrics);
+
+    let mut guard = manager.try_acquire(3 * PERMIT_GRANULARITY_BYTES).unwrap();
+    assert_eq!(manager.used_bytes(), 3 * PERMIT_GRANULARITY_BYTES);
+
+    // Acquire additional 4MB - should succeed
+    guard
+        .acquire_additional(4 * PERMIT_GRANULARITY_BYTES)
+        .await
+        .unwrap();
+    assert_eq!(guard.granted_bytes(), 7 * PERMIT_GRANULARITY_BYTES);
+    assert_eq!(manager.used_bytes(), 7 * PERMIT_GRANULARITY_BYTES);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn test_acquire_additional_with_policy_wait_success() {
+    use crate::policy::OnExhaustedPolicy;
+
+    let limit = 10 * PERMIT_GRANULARITY_BYTES;
+    let manager = MemoryManager::new(limit, NoOpMetrics);
+
+    let mut guard1 = manager.try_acquire(8 * PERMIT_GRANULARITY_BYTES).unwrap();
+
+    let manager_clone = manager.clone();
+    let waiter = tokio::spawn(async move {
+        let mut guard2 = manager_clone.try_acquire(0).unwrap();
+        // Wait policy with 1 second timeout
+        guard2
+            .acquire_additional_with_policy(
+                5 * PERMIT_GRANULARITY_BYTES,
+                OnExhaustedPolicy::Wait {
+                    timeout: Duration::from_secs(1),
+                },
+            )
+            .await
+            .unwrap();
+        guard2
+    });
+
+    sleep(Duration::from_millis(10)).await;
+
+    // Release memory to unblock waiter
+    assert!(guard1.release_partial(5 * PERMIT_GRANULARITY_BYTES));
+
+    let guard2 = waiter.await.unwrap();
+    assert_eq!(guard2.granted_bytes(), 5 * PERMIT_GRANULARITY_BYTES);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn test_acquire_additional_with_policy_wait_timeout() {
+    use crate::policy::OnExhaustedPolicy;
+
+    let limit = 10 * PERMIT_GRANULARITY_BYTES;
+    let manager = MemoryManager::new(limit, NoOpMetrics);
+
+    // Take all memory
+    let _guard1 = manager.try_acquire(10 * PERMIT_GRANULARITY_BYTES).unwrap();
+
+    let mut guard2 = manager.try_acquire(0).unwrap();
+
+    // Try to acquire with short timeout - should timeout
+    let result = guard2
+        .acquire_additional_with_policy(
+            5 * PERMIT_GRANULARITY_BYTES,
+            OnExhaustedPolicy::Wait {
+                timeout: Duration::from_millis(50),
+            },
+        )
+        .await;
+
+    assert!(result.is_err());
+    assert_eq!(guard2.granted_bytes(), 0);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn test_acquire_additional_with_policy_fail() {
+    use crate::policy::OnExhaustedPolicy;
+
+    let limit = 10 * PERMIT_GRANULARITY_BYTES;
+    let manager = MemoryManager::new(limit, NoOpMetrics);
+
+    let _guard1 = manager.try_acquire(8 * PERMIT_GRANULARITY_BYTES).unwrap();
+
+    let mut guard2 = manager.try_acquire(0).unwrap();
+
+    // Fail policy - should return error immediately
+    let result = guard2
+        .acquire_additional_with_policy(5 * PERMIT_GRANULARITY_BYTES, OnExhaustedPolicy::Fail)
+        .await;
+
+    assert!(result.is_err());
+    assert_eq!(guard2.granted_bytes(), 0);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn test_acquire_additional_unlimited() {
+    let manager = MemoryManager::new(0, NoOpMetrics); // Unlimited
+    let mut guard = manager.try_acquire(0).unwrap();
+
+    // Should always succeed with unlimited manager
+    guard
+        .acquire_additional(1000 * PERMIT_GRANULARITY_BYTES)
+        .await
+        .unwrap();
+    assert_eq!(guard.granted_bytes(), 0);
+    assert_eq!(manager.used_bytes(), 0);
+}