Skip to main content

pingap_upstream/
lib.rs

1// Copyright 2024-2025 Tree xie.
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 ahash::AHashMap;
16use snafu::Snafu;
17use std::collections::HashMap;
18use std::sync::Arc;
19
20mod backend_circuit_state;
21mod backend_stats;
22mod hash_strategy;
23mod peer_tracer;
24mod upstream;
25static LOG_TARGET: &str = "pingap::upstream";
26
27#[derive(Debug, Snafu)]
28pub enum Error {
29    #[snafu(display("Common error, category: {category}, {message}"))]
30    Common { message: String, category: String },
31}
32
33pub type Upstreams = AHashMap<String, Arc<Upstream>>;
34
35/// Upstream provider trait
36pub trait UpstreamProvider: Send + Sync {
37    /// Get an upstream by name
38    ///
39    /// # Arguments
40    /// * `name` - The name of the upstream to get
41    ///
42    /// # Returns
43    /// * `Option<Arc<Upstream>>` - The upstream if found, None otherwise
44    fn get(&self, name: &str) -> Option<Arc<Upstream>>;
45    /// Get the list of upstreams
46    ///
47    /// # Returns
48    /// * `Vec<(String, Arc<Upstream>)>` - The list of upstreams
49    fn list(&self) -> Vec<(String, Arc<Upstream>)>;
50
51    /// Get the healthy status of all upstreams
52    ///
53    /// # Returns
54    /// * `HashMap<String, UpstreamHealthyStatus>` - Healthy status of all upstreams
55    ///
56    /// This function iterates through all upstreams and checks their health status.
57    fn healthy_status(&self) -> HashMap<String, UpstreamHealthyStatus> {
58        let upstreams = self.list();
59        let mut healthy_status = HashMap::with_capacity(upstreams.len());
60        upstreams.iter().for_each(|(k, v)| {
61            let mut total = 0;
62            let mut healthy = 0;
63            let mut unhealthy_backends = vec![];
64            if let Some(backends) = v.get_backends() {
65                let backend_set = backends.get_backend();
66                total = backend_set.len();
67                backend_set.iter().for_each(|backend| {
68                    if backends.ready(backend) {
69                        healthy += 1;
70                    } else {
71                        unhealthy_backends.push(backend.to_string());
72                    }
73                });
74            }
75            healthy_status.insert(
76                k.to_string(),
77                UpstreamHealthyStatus {
78                    healthy,
79                    total: total as u32,
80                    unhealthy_backends,
81                },
82            );
83        });
84        healthy_status
85    }
86
87    /// Get the stats of all upstreams
88    ///
89    /// # Returns
90    /// * `HashMap<String, UpstreamStats>` - Stats of all upstreams
91    fn get_all_stats(&self) -> HashMap<String, UpstreamStats> {
92        let upstreams = self.list();
93        let mut stats_list = HashMap::with_capacity(upstreams.len());
94        upstreams.iter().for_each(|(k, v)| {
95            stats_list.insert(k.to_string(), v.stats());
96        });
97        stats_list
98    }
99}
100
101pub use hash_strategy::HashStrategy;
102pub use upstream::*;