common_memory_manager/
granularity.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt;
16
17/// Memory permit granularity for different use cases.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
19pub enum PermitGranularity {
20    /// 1 KB per permit
21    ///
22    /// Use for:
23    /// - HTTP/gRPC request limiting (small, high-concurrency operations)
24    /// - Small batch operations
25    /// - Scenarios requiring fine-grained fairness
26    Kilobyte,
27
28    /// 1 MB per permit (default)
29    ///
30    /// Use for:
31    /// - Query execution memory management
32    /// - Compaction memory control
33    /// - Large, long-running operations
34    #[default]
35    Megabyte,
36}
37
38impl PermitGranularity {
39    /// Returns the number of bytes per permit.
40    #[inline]
41    pub const fn bytes(self) -> u64 {
42        match self {
43            Self::Kilobyte => 1024,
44            Self::Megabyte => 1024 * 1024,
45        }
46    }
47
48    /// Returns a human-readable string representation.
49    pub const fn as_str(self) -> &'static str {
50        match self {
51            Self::Kilobyte => "1KB",
52            Self::Megabyte => "1MB",
53        }
54    }
55
56    /// Converts bytes to permits based on this granularity.
57    ///
58    /// Rounds up to ensure the requested bytes are fully covered.
59    /// Clamped to Semaphore::MAX_PERMITS.
60    #[inline]
61    pub fn bytes_to_permits(self, bytes: u64) -> u32 {
62        use tokio::sync::Semaphore;
63
64        let granularity_bytes = self.bytes();
65        bytes
66            .saturating_add(granularity_bytes - 1)
67            .saturating_div(granularity_bytes)
68            .min(Semaphore::MAX_PERMITS as u64)
69            .min(u32::MAX as u64) as u32
70    }
71
72    /// Converts permits to bytes based on this granularity.
73    #[inline]
74    pub fn permits_to_bytes(self, permits: u32) -> u64 {
75        (permits as u64).saturating_mul(self.bytes())
76    }
77}
78
79impl fmt::Display for PermitGranularity {
80    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
81        write!(f, "{}", self.as_str())
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn test_bytes_to_permits_kilobyte() {
91        let granularity = PermitGranularity::Kilobyte;
92
93        // Exact multiples
94        assert_eq!(granularity.bytes_to_permits(1024), 1);
95        assert_eq!(granularity.bytes_to_permits(2048), 2);
96        assert_eq!(granularity.bytes_to_permits(10 * 1024), 10);
97
98        // Rounds up
99        assert_eq!(granularity.bytes_to_permits(1), 1);
100        assert_eq!(granularity.bytes_to_permits(1025), 2);
101        assert_eq!(granularity.bytes_to_permits(2047), 2);
102    }
103
104    #[test]
105    fn test_bytes_to_permits_megabyte() {
106        let granularity = PermitGranularity::Megabyte;
107
108        // Exact multiples
109        assert_eq!(granularity.bytes_to_permits(1024 * 1024), 1);
110        assert_eq!(granularity.bytes_to_permits(2 * 1024 * 1024), 2);
111
112        // Rounds up
113        assert_eq!(granularity.bytes_to_permits(1), 1);
114        assert_eq!(granularity.bytes_to_permits(1024), 1);
115        assert_eq!(granularity.bytes_to_permits(1024 * 1024 + 1), 2);
116    }
117
118    #[test]
119    fn test_bytes_to_permits_zero_bytes() {
120        assert_eq!(PermitGranularity::Kilobyte.bytes_to_permits(0), 0);
121        assert_eq!(PermitGranularity::Megabyte.bytes_to_permits(0), 0);
122    }
123
124    #[test]
125    fn test_bytes_to_permits_clamps_to_maximum() {
126        use tokio::sync::Semaphore;
127
128        let max_permits = (Semaphore::MAX_PERMITS as u64).min(u32::MAX as u64) as u32;
129
130        assert_eq!(
131            PermitGranularity::Kilobyte.bytes_to_permits(u64::MAX),
132            max_permits
133        );
134        assert_eq!(
135            PermitGranularity::Megabyte.bytes_to_permits(u64::MAX),
136            max_permits
137        );
138    }
139
140    #[test]
141    fn test_permits_to_bytes() {
142        assert_eq!(PermitGranularity::Kilobyte.permits_to_bytes(1), 1024);
143        assert_eq!(PermitGranularity::Kilobyte.permits_to_bytes(10), 10 * 1024);
144
145        assert_eq!(PermitGranularity::Megabyte.permits_to_bytes(1), 1024 * 1024);
146        assert_eq!(
147            PermitGranularity::Megabyte.permits_to_bytes(10),
148            10 * 1024 * 1024
149        );
150    }
151
152    #[test]
153    fn test_round_trip_conversion() {
154        // Kilobyte: bytes -> permits -> bytes (should round up)
155        let kb = PermitGranularity::Kilobyte;
156        let permits = kb.bytes_to_permits(1500);
157        let bytes = kb.permits_to_bytes(permits);
158        assert!(bytes >= 1500); // Must cover original request
159        assert_eq!(bytes, 2048); // 2KB
160
161        // Megabyte: bytes -> permits -> bytes (should round up)
162        let mb = PermitGranularity::Megabyte;
163        let permits = mb.bytes_to_permits(1500);
164        let bytes = mb.permits_to_bytes(permits);
165        assert!(bytes >= 1500);
166        assert_eq!(bytes, 1024 * 1024); // 1MB
167    }
168}
common_memory_manager/granularity.rs

common_memory_manager/
granularity.rs
