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::*;