pingora_load_balancing/
discovery.rs

1// Copyright 2026 Cloudflare, Inc.
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
15//! Service discovery interface and implementations
16
17use arc_swap::ArcSwap;
18use async_trait::async_trait;
19use http::Extensions;
20use pingora_core::protocols::l4::socket::SocketAddr;
21use pingora_error::Result;
22use std::io::Result as IoResult;
23use std::net::ToSocketAddrs;
24use std::{
25    collections::{BTreeSet, HashMap},
26    sync::Arc,
27};
28
29use crate::Backend;
30
31/// [ServiceDiscovery] is the interface to discover [Backend]s.
32#[async_trait]
33pub trait ServiceDiscovery {
34    /// Return the discovered collection of backends and, optionally, whether
35    /// individual backends are enabled to serve.
36    ///
37    /// Enablement map keys are hashes of the corresponding [`Backend`] values,
38    /// produced with [`std::collections::hash_map::DefaultHasher`]. A backend
39    /// omitted from the map is considered enabled.
40    ///
41    /// Background services drop this future to cancel discovery on shutdown, so
42    /// an implementation must not rely on running to completion once polled.
43    async fn discover(&self) -> Result<(BTreeSet<Backend>, HashMap<u64, bool>)>;
44}
45
46// TODO: add DNS base discovery
47
48/// A static collection of [Backend]s for service discovery.
49#[derive(Default)]
50pub struct Static {
51    backends: ArcSwap<BTreeSet<Backend>>,
52}
53
54impl Static {
55    /// Create a new boxed [Static] service discovery with the given backends.
56    pub fn new(backends: BTreeSet<Backend>) -> Box<Self> {
57        Box::new(Static {
58            backends: ArcSwap::new(Arc::new(backends)),
59        })
60    }
61
62    /// Create a new boxed [Static] from a given iterator of items that implements [ToSocketAddrs].
63    pub fn try_from_iter<A, T: IntoIterator<Item = A>>(iter: T) -> IoResult<Box<Self>>
64    where
65        A: ToSocketAddrs,
66    {
67        let mut upstreams = BTreeSet::new();
68        for addrs in iter.into_iter() {
69            let addrs = addrs.to_socket_addrs()?.map(|addr| Backend {
70                addr: SocketAddr::Inet(addr),
71                weight: 1,
72                ext: Extensions::new(),
73            });
74            upstreams.extend(addrs);
75        }
76        Ok(Self::new(upstreams))
77    }
78
79    /// return the collection to backends
80    pub fn get(&self) -> BTreeSet<Backend> {
81        BTreeSet::clone(&self.backends.load())
82    }
83
84    // Concurrent set/add/remove might race with each other
85    // TODO: use a queue to avoid racing
86
87    // TODO: take an impl iter
88    #[allow(dead_code)]
89    pub(crate) fn set(&self, backends: BTreeSet<Backend>) {
90        self.backends.store(backends.into())
91    }
92
93    #[allow(dead_code)]
94    pub(crate) fn add(&self, backend: Backend) {
95        let mut new = self.get();
96        new.insert(backend);
97        self.set(new)
98    }
99
100    #[allow(dead_code)]
101    pub(crate) fn remove(&self, backend: &Backend) {
102        let mut new = self.get();
103        new.remove(backend);
104        self.set(new)
105    }
106}
107
108#[async_trait]
109impl ServiceDiscovery for Static {
110    async fn discover(&self) -> Result<(BTreeSet<Backend>, HashMap<u64, bool>)> {
111        // no readiness
112        let health = HashMap::new();
113        Ok((self.get(), health))
114    }
115}
pingora_load_balancing/discovery.rs

pingora_load_balancing/
discovery.rs
