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, Default)]
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 = "PerIo"))]
95    pub iops_counter: IopsCounter,
96}
97
98impl Throttle {
99    /// Create a new unlimited throttle config.
100    pub fn new() -> Self {
101        Self::default()
102    }
103
104    /// Set the maximum write iops for the device.
105    pub fn with_write_iops(mut self, iops: usize) -> Self {
106        self.write_iops = NonZeroUsize::new(iops);
107        self
108    }
109
110    /// Set the maximum read iops for the device.
111    pub fn with_read_iops(mut self, iops: usize) -> Self {
112        self.read_iops = NonZeroUsize::new(iops);
113        self
114    }
115
116    /// Set the maximum write throughput for the device.
117    pub fn with_write_throughput(mut self, throughput: usize) -> Self {
118        self.write_throughput = NonZeroUsize::new(throughput);
119        self
120    }
121
122    /// Set the maximum read throughput for the device.
123    pub fn with_read_throughput(mut self, throughput: usize) -> Self {
124        self.read_throughput = NonZeroUsize::new(throughput);
125        self
126    }
127
128    /// Set the iops counter for the device.
129    pub fn with_iops_counter(mut self, counter: IopsCounter) -> Self {
130        self.iops_counter = counter;
131        self
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn test_throttle_default() {
141        assert!(matches!(
142            Throttle::new(),
143            Throttle {
144                write_iops: None,
145                read_iops: None,
146                write_throughput: None,
147                read_throughput: None,
148                iops_counter: IopsCounter::PerIo,
149            }
150        ));
151    }
152
153    #[test]
154    fn test_iops_counter_from_str() {
155        assert!(matches!(IopsCounter::from_str("PerIo"), Ok(IopsCounter::PerIo)));
156        assert!(matches!(IopsCounter::from_str(" PerIo "), Ok(IopsCounter::PerIo)));
157        assert!(matches!(IopsCounter::from_str("PerIo "), Ok(IopsCounter::PerIo)));
158        assert!(matches!(IopsCounter::from_str(" PerIo"), Ok(IopsCounter::PerIo)));
159
160        let _num = NonZeroUsize::new(1024).unwrap();
161
162        assert!(matches!(
163            IopsCounter::from_str("PerIoSize(1024)"),
164            Ok(IopsCounter::PerIoSize(_num))
165        ));
166        assert!(matches!(
167            IopsCounter::from_str(" PerIoSize(1024) "),
168            Ok(IopsCounter::PerIoSize(_num))
169        ));
170        assert!(matches!(
171            IopsCounter::from_str("PerIoSize(1024) "),
172            Ok(IopsCounter::PerIoSize(_num))
173        ));
174        assert!(matches!(
175            IopsCounter::from_str(" PerIoSize(1024)"),
176            Ok(IopsCounter::PerIoSize(_num))
177        ));
178
179        assert!(IopsCounter::from_str("PerIoSize(0)").is_err());
180        assert!(IopsCounter::from_str("PerIoSize(1024a)").is_err());
181
182        assert!(IopsCounter::from_str("invalid_string").is_err());
183    }
184}
foyer_storage/io/device/throttle.rs

foyer_storage/io/device/
throttle.rs
