1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

/* Copyright 2023 Architect Financial Technologies LLC. This is free
 * software released under the GNU Affero Public License version 3. */
//! set and query limits

use crate::{protocol::limits::{GlobalLimitName, Limit, LimitScope, LimitSet}, config::Common};
use anyhow::{bail, Result};
use netidx::{
    pack::Pack, pool::Pooled, protocol::value::Value,
    utils::pack,
};
use netidx_protocols::{call_rpc, rpc::client::Proc};

/// Manage limits in the OMS
pub struct OmsLimitsApi {
    set_limits: Proc,
    get_limits: Proc,
    get_global_limits: Proc,
}

impl OmsLimitsApi {
    pub async fn new(common: &Common) -> Result<Self> {
        let base = common.paths.core_limit();
        Ok(Self {
            set_limits: Proc::new(&common.subscriber, base.append("set-limits")).await?,
            get_limits: Proc::new(&common.subscriber, base.append("get-limits")).await?,
            get_global_limits: Proc::new(&common.subscriber, base.append("get-global-limits")).await?,
        })
    }

    /// Set the specified limits
    pub async fn set_limits(&self, limits: &Pooled<Vec<Limit>>) -> Result<()> {
        let limits = Value::Bytes(pack(limits)?.freeze());
        let res = call_rpc!(&self.set_limits, limits: limits).await?;
        match res {
            Value::Error(e) => bail!(e.to_string()),
            Value::Ok => Ok(()),
            _ => bail!("unexpected response"),
        }
    }

    /// Get the specified limits. If no limit is set for a set and
    /// scope then it will not appear in the result.
    pub async fn get_limits(
        &self,
        query: &Pooled<Vec<(LimitSet, LimitScope)>>,
    ) -> Result<Pooled<Vec<Limit>>> {
        let query = Value::Bytes(pack(query)?.freeze());
        let res = call_rpc!(&self.get_limits, query: query).await?;
        match res {
            Value::Error(e) => bail!(e.to_string()),
            Value::Bytes(b) => Ok(Pack::decode(&mut &*b)?),
            _ => bail!("unexpected response"),
        }
    }

    /// Get the value of an OMS global limit
    pub async fn get_global_limits(
        &self,
        query: &Pooled<Vec<GlobalLimitName>>,
    ) -> Result<Pooled<Vec<Limit>>> {
        let query = Value::Bytes(pack(query)?.freeze());
        let res = call_rpc!(&self.get_global_limits, query: query).await?;
        match res {
            Value::Error(e) => bail!(e.to_string()),
            Value::Bytes(b) => Ok(Pack::decode(&mut &*b)?),
            _ => bail!("unexpected response"),
        }
    }
}