rama_proxy/proxydb/
update.rs

1use super::ProxyDB;
2use arc_swap::ArcSwap;
3use rama_core::error::BoxErrorExt as _;
4use rama_core::error::{BoxError, ErrorContext};
5use rama_utils::collections::NonEmptyVec;
6use std::{fmt, num::NonZeroUsize, ops::Deref, sync::Arc};
7
8/// Create a new [`ProxyDB`] updater which allows you to have a (typically in-memory) [`ProxyDB`]
9/// which you can update live.
10///
11/// This construct returns a pair of:
12///
13/// - [`LiveUpdateProxyDB`]: to be used as the [`ProxyDB`] instead of the inner `T`, dubbed the "reader";
14/// - [`LiveUpdateProxyDBSetter`]: to be used as the _only_ way to set the inner `T` as many time as you wish, dubbed the "writer".
15///
16/// Note that the inner `T` is not yet created when this construct returns this pair.
17/// Until you actually called [`LiveUpdateProxyDBSetter::set`] with the inner `T` [`ProxyDB`],
18/// any [`ProxyDB`] trait method call to [`LiveUpdateProxyDB`] will fail.
19///
20/// It is therefore recommended that you immediately set the inner `T` [`ProxyDB`] upon
21/// receiving the reader/writer pair, prior to starting to actually use the [`ProxyDB`]
22/// in your rama service stack.
23///
24/// This goal of this updater is to be fast for reading (getting proxies),
25/// and slow for the infrequent updates (setting the proxy db). As such it is recommended
26/// to not update the [`ProxyDB`] to frequent. An example use case for this updater
27/// could be to update your in-memory proxy database every 15 minutes, by populating it from
28/// a shared external database (e.g. MySQL`). Failures to create a new `T` ProxyDB should be handled
29/// by the Writer, and can be as simple as just logging it and move on without an update.
30pub fn proxy_db_updater<T>() -> (LiveUpdateProxyDB<T>, LiveUpdateProxyDBSetter<T>)
31where
32    T: ProxyDB<Error: Into<BoxError>>,
33{
34    let data = Arc::new(ArcSwap::from_pointee(None));
35    let reader = LiveUpdateProxyDB(data.clone());
36    let writer = LiveUpdateProxyDBSetter(data);
37    (reader, writer)
38}
39
40/// A wrapper around a `T` [`ProxyDB`] which can be updated
41/// through the _only_ linked writer [`LiveUpdateProxyDBSetter`].
42///
43/// See [`proxy_db_updater`] for more details.
44pub struct LiveUpdateProxyDB<T>(Arc<ArcSwap<Option<T>>>);
45
46impl<T: fmt::Debug> fmt::Debug for LiveUpdateProxyDB<T> {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        f.debug_tuple("LiveUpdateProxyDB").field(&self.0).finish()
49    }
50}
51
52impl<T> Clone for LiveUpdateProxyDB<T> {
53    fn clone(&self) -> Self {
54        Self(self.0.clone())
55    }
56}
57
58impl<T> ProxyDB for LiveUpdateProxyDB<T>
59where
60    T: ProxyDB<Error: Into<BoxError>>,
61{
62    type Error = BoxError;
63
64    async fn get_proxies_if(
65        &self,
66        ctx: super::ProxyContext,
67        filter: super::ProxyFilter,
68        predicate: impl super::ProxyQueryPredicate,
69        limit: Option<NonZeroUsize>,
70    ) -> Result<NonEmptyVec<super::Proxy>, Self::Error> {
71        match self.0.load().deref().deref() {
72            Some(db) => db
73                .get_proxies_if(ctx, filter, predicate, limit)
74                .await
75                .into_box_error(),
76            None => Err(BoxError::from_static_str(
77                "live proxy db: proxy db is None: get_proxies_if unable to proceed",
78            )),
79        }
80    }
81
82    async fn get_proxies(
83        &self,
84        ctx: super::ProxyContext,
85        filter: super::ProxyFilter,
86        limit: Option<NonZeroUsize>,
87    ) -> Result<NonEmptyVec<super::Proxy>, Self::Error> {
88        match self.0.load().deref().deref() {
89            Some(db) => db.get_proxies(ctx, filter, limit).await.into_box_error(),
90            None => Err(BoxError::from_static_str(
91                "live proxy db: proxy db is None: get_proxies unable to proceed",
92            )),
93        }
94    }
95
96    async fn get_proxy_if(
97        &self,
98        ctx: super::ProxyContext,
99        filter: super::ProxyFilter,
100        predicate: impl super::ProxyQueryPredicate,
101    ) -> Result<super::Proxy, Self::Error> {
102        match self.0.load().deref().deref() {
103            Some(db) => db
104                .get_proxy_if(ctx, filter, predicate)
105                .await
106                .into_box_error(),
107            None => Err(BoxError::from_static_str(
108                "live proxy db: proxy db is None: get_proxy_if unable to proceed",
109            )),
110        }
111    }
112
113    async fn get_proxy(
114        &self,
115        ctx: super::ProxyContext,
116        filter: super::ProxyFilter,
117    ) -> Result<super::Proxy, Self::Error> {
118        match self.0.load().deref().deref() {
119            Some(db) => db.get_proxy(ctx, filter).await.into_box_error(),
120            None => Err(BoxError::from_static_str(
121                "live proxy db: proxy db is None: get_proxy unable to proceed",
122            )),
123        }
124    }
125}
126
127/// Writer to set a new [`ProxyDB`] in the linked [`LiveUpdateProxyDB`].
128///
129/// There can only be one writer [`LiveUpdateProxyDBSetter`] for each
130/// collection of [`LiveUpdateProxyDB`] linked to the same internal data `T`.
131///
132/// See [`proxy_db_updater`] for more details.
133pub struct LiveUpdateProxyDBSetter<T>(Arc<ArcSwap<Option<T>>>);
134
135impl<T> LiveUpdateProxyDBSetter<T> {
136    /// Set the new `T` [`ProxyDB`] to be used for future [`ProxyDB`]
137    /// calls made to the linked [`LiveUpdateProxyDB`] instances.
138    pub fn set(&self, db: T) {
139        self.0.store(Arc::new(Some(db)))
140    }
141}
142
143impl<T: fmt::Debug> fmt::Debug for LiveUpdateProxyDBSetter<T> {
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        f.debug_tuple("LiveUpdateProxyDBSetter")
146            .field(&self.0)
147            .finish()
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use crate::{Proxy, ProxyFilter, proxydb::ProxyContext};
154    use rama_net::{asn::Asn, transport::TransportProtocol};
155    use rama_utils::str::non_empty_str;
156
157    use super::*;
158
159    #[tokio::test]
160    async fn test_empty_live_update_db() {
161        let (reader, _) = proxy_db_updater::<Proxy>();
162        reader
163            .get_proxies(
164                ProxyContext {
165                    protocol: TransportProtocol::Tcp,
166                },
167                ProxyFilter::default(),
168                None,
169            )
170            .await
171            .unwrap_err();
172        reader
173            .get_proxy(
174                ProxyContext {
175                    protocol: TransportProtocol::Tcp,
176                },
177                ProxyFilter::default(),
178            )
179            .await
180            .unwrap_err();
181    }
182
183    #[tokio::test]
184    async fn test_live_update_db_updated() {
185        let (reader, writer) = proxy_db_updater();
186
187        reader
188            .get_proxy(
189                ProxyContext {
190                    protocol: TransportProtocol::Tcp,
191                },
192                ProxyFilter::default(),
193            )
194            .await
195            .unwrap_err();
196
197        writer.set(Proxy {
198            id: non_empty_str!("id"),
199            address: "authority:80".parse().unwrap(),
200            tcp: true,
201            udp: false,
202            http: false,
203            https: true,
204            socks5: false,
205            socks5h: false,
206            datacenter: true,
207            residential: false,
208            mobile: true,
209            pool_id: Some("pool_id".into()),
210            continent: Some("continent".into()),
211            country: Some("country".into()),
212            state: Some("state".into()),
213            city: Some("city".into()),
214            carrier: Some("carrier".into()),
215            asn: Some(Asn::from_static(1)),
216        });
217
218        assert_eq!(
219            "id",
220            reader
221                .get_proxy(
222                    ProxyContext {
223                        protocol: TransportProtocol::Tcp,
224                    },
225                    ProxyFilter::default(),
226                )
227                .await
228                .unwrap()
229                .id
230        );
231
232        let proxies = reader
233            .get_proxies(
234                ProxyContext {
235                    protocol: TransportProtocol::Tcp,
236                },
237                ProxyFilter::default(),
238                None,
239            )
240            .await
241            .unwrap();
242        assert_eq!(proxies.len(), 1);
243        assert_eq!(proxies.head.id, "id");
244
245        reader
246            .get_proxy(
247                ProxyContext {
248                    protocol: TransportProtocol::Udp,
249                },
250                ProxyFilter::default(),
251            )
252            .await
253            .unwrap_err();
254
255        assert_eq!(
256            "id",
257            reader
258                .get_proxy(
259                    ProxyContext {
260                        protocol: TransportProtocol::Tcp,
261                    },
262                    ProxyFilter::default(),
263                )
264                .await
265                .unwrap()
266                .id
267        );
268    }
269}
rama_proxy/proxydb/update.rs

rama_proxy/proxydb/
update.rs
