feat: track unlimited usage in memory manager

[fengjiachun]

I hereby agree to the terms of the GreptimeDB CLA.

Refer to a related PR or issue link (optional)

What's changed and what's your intention?

As gemini said

PR Checklist

Please convert it to a draft if some of the following conditions are not met.

• I have written the necessary rustdoc comments.
• I have added the necessary unit tests and integration tests.
• This PR requires documentation updates.
• API changes are backward compatible.
• Schema or data changes are backward compatible.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the memory manager by introducing the capability to track memory usage even when no hard limit is set. This change provides a more comprehensive view of memory consumption across different operational modes, improving system observability and resource management without imposing artificial constraints.

Highlights

• Unlimited Memory Quota Tracking: Introduced a new UnlimitedMemoryQuota structure and integrated it into the MemoryManager to accurately track memory usage even when no explicit memory limit is configured. This allows for better visibility into memory consumption in unconstrained environments.
• Refactored Memory Quota State: The MemoryManager now uses a MemoryQuotaState enum to explicitly differentiate between Limited and Unlimited memory management modes, simplifying logic and improving clarity.
• Enhanced MemoryGuard Functionality: The MemoryGuard has been updated to correctly manage and report granted bytes for unlimited allocations, including support for acquiring additional memory, partially releasing memory, and proper cleanup upon dropping the guard.
• Updated Metrics and Assertions: Memory usage metrics are now consistently updated for both limited and unlimited quotas. Unit tests have been expanded and adjusted to validate the new unlimited memory tracking behavior, ensuring correctness.

Changelog

• src/common/memory-manager/src/guard.rs
  • Removed unused common_telemetry::debug import.
  • Added UnlimitedMemoryQuota import.
  • Refactored GuardState enum to include Released and a new Unlimited variant that tracks UnlimitedMemoryQuota and granted_bytes.
  • Implemented a release method for GuardState to centralize memory release logic.
  • Modified MemoryGuard::unlimited and MemoryGuard::limited constructors to align with new GuardState variants and argument order.
  • Updated granted_bytes method to correctly report allocated bytes for unlimited guards.
  • Adjusted acquire_additional and try_acquire_additional methods to handle unlimited memory guards by updating tracked bytes.
  • Removed debug! calls from acquisition methods.
  • Updated MemoryLimitExceededSnafu error handling to account for new GuardState variants.
  • Modified release_partial to correctly decrease tracked bytes for unlimited guards.
  • Simplified the Drop implementation for MemoryGuard to use the new GuardState::release method.
• src/common/memory-manager/src/manager.rs
  • Added AtomicU64, Ordering, and OwnedSemaphorePermit imports for atomic operations and semaphore management.
  • Changed MemoryManager's internal quota field to use the new MemoryQuotaState enum.
  • Introduced UnlimitedMemoryQuota struct to manage current bytes using Arc<AtomicU64> and metrics.
  • Defined MemoryQuotaState enum to encapsulate Unlimited and Limited quota types.
  • Modified with_granularity to initialize UnlimitedMemoryQuota when limit_bytes is 0.
  • Updated limit_bytes, used_bytes, and available_bytes methods to correctly report values based on the MemoryQuotaState, with u64::MAX for available bytes in unlimited mode.
  • Adjusted acquire and try_acquire methods to return appropriate MemoryGuard instances based on the MemoryQuotaState.
  • Added release_permit method to MemoryQuota for consistent permit release.
  • Implemented add_in_use, sub_in_use, and release_bytes methods for UnlimitedMemoryQuota to atomically manage memory usage.
• src/common/memory-manager/src/tests.rs
  • Updated test_try_acquire_unlimited to assert correct granted_bytes, used_bytes, and available_bytes for unlimited managers.
  • Modified test_request_additional_unlimited to verify that granted_bytes and used_bytes are tracked and released correctly for unlimited managers.
  • Adjusted test_early_release_partial_unlimited to confirm that partial releases are accurately reflected in granted_bytes and used_bytes for unlimited guards.
  • Updated test_acquire_additional_unlimited to assert that granted_bytes and used_bytes are correctly updated after additional acquisitions in unlimited mode.

Activity

• The author, fengjiachun, created this pull request to implement the feature of tracking unlimited usage in the memory manager.
• The author has acknowledged the GreptimeDB CLA.
• The author has provided a standard PR checklist, indicating the status of documentation, tests, and API compatibility.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces tracking for memory usage in 'unlimited' mode within the memory manager. Previously, unlimited guards did not track their granted bytes. The changes refactor GuardState and MemoryManager to handle limited and unlimited quotas explicitly, using an atomic counter for unlimited usage. The related methods and tests are updated accordingly. The implementation is solid. I've added a few suggestions to improve code clarity and maintainability by refactoring some of the memory release logic.