promql_parser/util/
duration.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 lazy_static::lazy_static;
16use regex::Regex;
17use std::fmt::Write;
18use std::time::Duration;
19
20lazy_static! {
21    static ref DURATION_RE: Regex = Regex::new(
22        r"(?x)
23^
24((?P<y>[0-9]+)y)?
25((?P<w>[0-9]+)w)?
26((?P<d>[0-9]+)d)?
27((?P<h>[0-9]+)h)?
28((?P<m>[0-9]+)m)?
29((?P<s>[0-9]+)s)?
30((?P<ms>[0-9]+)ms)?
31$",
32    )
33    .unwrap();
34}
35
36pub const MILLI_DURATION: Duration = Duration::from_millis(1);
37pub const SECOND_DURATION: Duration = Duration::from_secs(1);
38pub const MINUTE_DURATION: Duration = Duration::from_secs(60);
39pub const HOUR_DURATION: Duration = Duration::from_secs(60 * 60);
40pub const DAY_DURATION: Duration = Duration::from_secs(60 * 60 * 24);
41pub const WEEK_DURATION: Duration = Duration::from_secs(60 * 60 * 24 * 7);
42pub const YEAR_DURATION: Duration = Duration::from_secs(60 * 60 * 24 * 365);
43
44const ALL_CAPS: [(&str, Duration); 7] = [
45    ("y", YEAR_DURATION),
46    ("w", WEEK_DURATION),
47    ("d", DAY_DURATION),
48    ("h", HOUR_DURATION),
49    ("m", MINUTE_DURATION),
50    ("s", SECOND_DURATION),
51    ("ms", MILLI_DURATION),
52];
53
54/// parses a string into a Duration, assuming that a year
55/// always has 365d, a week always has 7d, and a day always has 24h.
56///
57/// # Examples
58///
59/// Basic usage:
60///
61/// ```
62/// use std::time::Duration;
63/// use promql_parser::util;
64///
65/// assert_eq!(util::parse_duration("1h").unwrap(), Duration::from_secs(3600));
66/// assert_eq!(util::parse_duration("4d").unwrap(), Duration::from_secs(3600 * 24 * 4));
67/// assert_eq!(util::parse_duration("4d1h").unwrap(), Duration::from_secs(3600 * 97));
68/// ```
69pub fn parse_duration(ds: &str) -> Result<Duration, String> {
70    if ds.is_empty() {
71        return Err("empty duration string".into());
72    }
73
74    if ds == "0" {
75        return Err("duration must be greater than 0".into());
76    }
77
78    // the duration is float number of seconds
79    if let Ok(float_duration) = ds.parse::<f64>() {
80        return Ok(Duration::from_secs_f64(float_duration));
81    }
82
83    if !DURATION_RE.is_match(ds) {
84        return Err(format!("not a valid duration string: {ds}"));
85    }
86
87    let caps = DURATION_RE.captures(ds).unwrap();
88    let dur = ALL_CAPS
89        .into_iter()
90        // map captured string to Option<Duration> iterator
91        // FIXME: None is ignored in closure. It is better to tell users which part is wrong.
92        .map(|(title, duration)| {
93            caps.name(title)
94                .and_then(|cap| cap.as_str().parse::<u32>().ok())
95                .and_then(|v| duration.checked_mul(v))
96        })
97        .try_fold(Duration::ZERO, |acc, x| {
98            acc.checked_add(x.unwrap_or(Duration::ZERO))
99                .ok_or_else(|| "duration overflowed".into())
100        });
101
102    if matches!(dur, Ok(d) if d == Duration::ZERO) {
103        Err("duration must be greater than 0".into())
104    } else {
105        dur
106    }
107}
108
109/// display Duration in Prometheus format
110pub fn display_duration(duration: &Duration) -> String {
111    if duration.is_zero() {
112        return "0s".into();
113    }
114    let mut ms = duration.as_millis();
115    let mut ss = String::new();
116
117    let mut f = |unit: &str, mult: u128, exact: bool| {
118        if exact && ms % mult != 0 {
119            return;
120        }
121
122        let v = ms / mult;
123        if v > 0 {
124            write!(ss, "{v}{unit}").unwrap();
125            ms -= v * mult
126        }
127    };
128
129    // Only format years and weeks if the remainder is zero, as it is often
130    // easier to read 90d than 12w6d.
131    f("y", 1000 * 60 * 60 * 24 * 365, true);
132    f("w", 1000 * 60 * 60 * 24 * 7, true);
133
134    f("d", 1000 * 60 * 60 * 24, false);
135    f("h", 1000 * 60 * 60, false);
136    f("m", 1000 * 60, false);
137    f("s", 1000, false);
138    f("ms", 1, false);
139
140    ss
141}
142
143#[cfg(feature = "ser")]
144pub(crate) fn serialize_duration<S>(dur: &Duration, serializer: S) -> Result<S::Ok, S::Error>
145where
146    S: serde::Serializer,
147{
148    let duration_millis = dur.as_millis();
149    serializer.serialize_u128(duration_millis)
150}
151
152#[cfg(feature = "ser")]
153pub(crate) fn serialize_duration_opt<S>(
154    dur: &Option<Duration>,
155    serializer: S,
156) -> Result<S::Ok, S::Error>
157where
158    S: serde::Serializer,
159{
160    if let Some(dur) = dur {
161        serialize_duration(dur, serializer)
162    } else {
163        serializer.serialize_none()
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    #[test]
172    fn test_re() {
173        // valid regex
174        let res = vec![
175            "1y", "2w", "3d", "4h", "5m", "6s", "7ms", "1y2w3d", "4h30m", "3600ms",
176        ];
177        for re in res {
178            assert!(DURATION_RE.is_match(re), "{re} failed.")
179        }
180
181        // invalid regex
182        let res = vec!["1", "1y1m1d", "-1w", "1.5d", "d"];
183        for re in res {
184            assert!(!DURATION_RE.is_match(re), "{re} failed.")
185        }
186    }
187
188    #[test]
189    fn test_valid_duration() {
190        let ds = vec![
191            ("324ms", Duration::from_millis(324)),
192            ("3s", Duration::from_secs(3)),
193            ("5m", MINUTE_DURATION * 5),
194            ("1h", HOUR_DURATION),
195            ("4d", DAY_DURATION * 4),
196            ("4d1h", DAY_DURATION * 4 + HOUR_DURATION),
197            ("14d", DAY_DURATION * 14),
198            ("3w", WEEK_DURATION * 3),
199            ("3w2d1h", WEEK_DURATION * 3 + HOUR_DURATION * 49),
200            ("10y", YEAR_DURATION * 10),
201            ("0.5", Duration::from_secs_f64(0.5)),
202            (".5", Duration::from_secs_f64(0.5)),
203            ("1", Duration::from_secs_f64(1.)),
204            ("11.2345", Duration::from_secs_f64(11.2345)),
205        ];
206
207        for (s, expect) in ds {
208            let d = parse_duration(s);
209            assert!(d.is_ok());
210            assert_eq!(expect, d.unwrap(), "{s} and {expect:?} not matched");
211        }
212    }
213
214    // valid here but invalid in PromQL Go Version
215    #[test]
216    fn test_diff_with_promql() {
217        let ds = vec![
218            ("294y", YEAR_DURATION * 294),
219            ("200y10400w", YEAR_DURATION * 200 + WEEK_DURATION * 10400),
220            ("107675d", DAY_DURATION * 107675),
221            ("2584200h", HOUR_DURATION * 2584200),
222        ];
223
224        for (s, expect) in ds {
225            let d = parse_duration(s);
226            assert!(d.is_ok());
227            assert_eq!(expect, d.unwrap(), "{s} and {expect:?} not matched");
228        }
229    }
230
231    #[test]
232    fn test_invalid_duration() {
233        let ds = vec!["1y1m1d", "-1w", "1.5d", "d", "", "0", "0w", "0s"];
234        for d in ds {
235            assert!(parse_duration(d).is_err(), "{d} is invalid duration!");
236        }
237    }
238
239    #[test]
240    fn test_display_duration() {
241        let ds = vec![
242            (Duration::ZERO, "0s"),
243            (Duration::from_millis(324), "324ms"),
244            (Duration::from_secs(3), "3s"),
245            (MINUTE_DURATION * 5, "5m"),
246            (MINUTE_DURATION * 5 + MILLI_DURATION * 500, "5m500ms"),
247            (HOUR_DURATION, "1h"),
248            (DAY_DURATION * 4, "4d"),
249            (DAY_DURATION * 4 + HOUR_DURATION, "4d1h"),
250            (
251                DAY_DURATION * 4 + HOUR_DURATION * 2 + MINUTE_DURATION * 10,
252                "4d2h10m",
253            ),
254            (DAY_DURATION * 14, "2w"),
255            (WEEK_DURATION * 3, "3w"),
256            (WEEK_DURATION * 3 + HOUR_DURATION * 49, "23d1h"),
257            (YEAR_DURATION * 10, "10y"),
258        ];
259
260        for (d, expect) in ds {
261            let s = display_duration(&d);
262            assert_eq!(expect, s, "{s} and {expect:?} not matched");
263        }
264    }
265}