foyer_storage/io/device/
throttle.rs

1// Copyright 2026 foyer Project Authors
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::Display, num::NonZeroUsize, str::FromStr};
16
17/// Device iops counter.
18#[derive(Debug, Clone, PartialEq, Eq, Default)]
19#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
20pub enum IopsCounter {
21    /// Count 1 iops for each read/write.
22    #[default]
23    PerIo,
24    /// Count 1 iops for each read/write with the size of the i/o.
25    PerIoSize(NonZeroUsize),
26}
27
28impl IopsCounter {
29    /// Create a new iops counter that count 1 iops for each io.
30    pub fn per_io() -> Self {
31        Self::PerIo
32    }
33
34    /// Create a new iops counter that count 1 iops for every io size in bytes among ios.
35    ///
36    /// NOTE: `io_size` must NOT be zero.
37    pub fn per_io_size(io_size: usize) -> Self {
38        Self::PerIoSize(NonZeroUsize::new(io_size).expect("io size must be non-zero"))
39    }
40
41    /// Count io(s) by io size in bytes.
42    pub fn count(&self, bytes: usize) -> usize {
43        match self {
44            IopsCounter::PerIo => 1,
45            IopsCounter::PerIoSize(size) => bytes / *size + if bytes % *size != 0 { 1 } else { 0 },
46        }
47    }
48}
49
50impl Display for IopsCounter {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        match self {
53            IopsCounter::PerIo => write!(f, "PerIo"),
54            IopsCounter::PerIoSize(size) => write!(f, "PerIoSize({size})"),
55        }
56    }
57}
58
59impl FromStr for IopsCounter {
60    type Err = anyhow::Error;
61
62    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
63        let s = s.trim();
64        match s {
65            "PerIo" => Ok(IopsCounter::PerIo),
66            _ if s.starts_with("PerIoSize(") && s.ends_with(')') => {
67                let num = &s[10..s.len() - 1];
68                let v = num.parse::<NonZeroUsize>()?;
69                Ok(IopsCounter::PerIoSize(v))
70            }
71            _ => Err(anyhow::anyhow!("Invalid IopsCounter format: {}", s)),
72        }
73    }
74}
75
76/// Throttle config for the device.
77#[derive(Debug, Clone, PartialEq, Eq)]
78#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
79#[cfg_attr(feature = "clap", derive(clap::Args))]
80pub struct Throttle {
81    /// The maximum write iops for the device.
82    #[cfg_attr(feature = "clap", clap(long))]
83    pub write_iops: Option<NonZeroUsize>,
84    /// The maximum read iops for the device.
85    #[cfg_attr(feature = "clap", clap(long))]
86    pub read_iops: Option<NonZeroUsize>,
87    /// The maximum write throughput for the device.
88    #[cfg_attr(feature = "clap", clap(long))]
89    pub write_throughput: Option<NonZeroUsize>,
90    /// The maximum read throughput for the device.
91    #[cfg_attr(feature = "clap", clap(long))]
92    pub read_throughput: Option<NonZeroUsize>,
93    /// The iops counter for the device.
94    #[cfg_attr(feature = "clap", clap(long, default_value_t))]
95    pub iops_counter: IopsCounter,
96}
97
98impl Default for Throttle {
99    fn default() -> Self {
100        Throttle::new()
101    }
102}
103
104impl Throttle {
105    /// Create a new unlimited throttle config.
106    pub const fn new() -> Self {
107        Self {
108            write_iops: None,
109            read_iops: None,
110            write_throughput: None,
111            read_throughput: None,
112            iops_counter: IopsCounter::PerIo,
113        }
114    }
115
116    /// Set the maximum write iops for the device.
117    pub const fn with_write_iops(mut self, iops: usize) -> Self {
118        self.write_iops = NonZeroUsize::new(iops);
119        self
120    }
121
122    /// Set the maximum read iops for the device.
123    pub const fn with_read_iops(mut self, iops: usize) -> Self {
124        self.read_iops = NonZeroUsize::new(iops);
125        self
126    }
127
128    /// Set the maximum write throughput for the device.
129    pub const fn with_write_throughput(mut self, throughput: usize) -> Self {
130        self.write_throughput = NonZeroUsize::new(throughput);
131        self
132    }
133
134    /// Set the maximum read throughput for the device.
135    pub const fn with_read_throughput(mut self, throughput: usize) -> Self {
136        self.read_throughput = NonZeroUsize::new(throughput);
137        self
138    }
139
140    /// Set the iops counter for the device.
141    pub const fn with_iops_counter(mut self, counter: IopsCounter) -> Self {
142        self.iops_counter = counter;
143        self
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn test_throttle_default() {
153        assert!(matches!(
154            Throttle::new(),
155            Throttle {
156                write_iops: None,
157                read_iops: None,
158                write_throughput: None,
159                read_throughput: None,
160                iops_counter: IopsCounter::PerIo,
161            }
162        ));
163    }
164
165    #[test]
166    fn test_iops_counter_from_str() {
167        assert!(matches!(IopsCounter::from_str("PerIo"), Ok(IopsCounter::PerIo)));
168        assert!(matches!(IopsCounter::from_str(" PerIo "), Ok(IopsCounter::PerIo)));
169        assert!(matches!(IopsCounter::from_str("PerIo "), Ok(IopsCounter::PerIo)));
170        assert!(matches!(IopsCounter::from_str(" PerIo"), Ok(IopsCounter::PerIo)));
171
172        let _num = NonZeroUsize::new(1024).unwrap();
173
174        assert!(matches!(
175            IopsCounter::from_str("PerIoSize(1024)"),
176            Ok(IopsCounter::PerIoSize(_num))
177        ));
178        assert!(matches!(
179            IopsCounter::from_str(" PerIoSize(1024) "),
180            Ok(IopsCounter::PerIoSize(_num))
181        ));
182        assert!(matches!(
183            IopsCounter::from_str("PerIoSize(1024) "),
184            Ok(IopsCounter::PerIoSize(_num))
185        ));
186        assert!(matches!(
187            IopsCounter::from_str(" PerIoSize(1024)"),
188            Ok(IopsCounter::PerIoSize(_num))
189        ));
190
191        assert!(IopsCounter::from_str("PerIoSize(0)").is_err());
192        assert!(IopsCounter::from_str("PerIoSize(1024a)").is_err());
193
194        assert!(IopsCounter::from_str("invalid_string").is_err());
195    }
196}
foyer_storage/io/device/throttle.rs

foyer_storage/io/device/
throttle.rs
