diff --git a/.gitignore b/.gitignore index f59e5f6f786..862eb8c5b4f 100644 --- a/.gitignore +++ b/.gitignore @@ -67,3 +67,6 @@ greptimedb_data # Claude code CLAUDE.md + +# AGENTS.md +AGENTS.md diff --git a/src/common/memory-manager/src/guard.rs b/src/common/memory-manager/src/guard.rs index 6bad763b6a2..770b6dec24a 100644 --- a/src/common/memory-manager/src/guard.rs +++ b/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 { @@ -55,11 +60,52 @@ impl MemoryGuard { } } - /// 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 MemoryGuard { 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 MemoryGuard { } } - /// 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 MemoryGuard { 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, diff --git a/src/common/memory-manager/src/tests.rs b/src/common/memory-manager/src/tests.rs index c35889c00f7..886eef9dac4 100644 --- a/src/common/memory-manager/src/tests.rs +++ b/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); +}