mz_aws_util/
http_pool.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! An AWS SDK HTTP client that spreads requests over a growing set of the
11//! SDK's own clients.
12//!
13//! Over HTTP/2, one SDK client sends every request to a host over a single
14//! connection. Each client has its own connections, so spreading requests over
15//! several clients spreads them over several connections.
16
17use std::sync::{Arc, Mutex, Weak};
18
19use aws_smithy_runtime_api::box_error::BoxError;
20use aws_smithy_runtime_api::client::connector_metadata::ConnectorMetadata;
21use aws_smithy_runtime_api::client::http::{
22    HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpClient,
23    SharedHttpConnector,
24};
25use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
26use aws_smithy_runtime_api::client::runtime_components::{
27    RuntimeComponents, RuntimeComponentsBuilder,
28};
29use aws_smithy_types::body::SdkBody;
30use aws_smithy_types::config_bag::ConfigBag;
31use http_body_util::BodyExt;
32use mz_ore::metrics::IntCounter;
33
34/// The pool starts with this many clients.
35///
36/// In our benchmarks of GCS reads, one connection carried 300 to 400 MiB/s and
37/// total reads topped out around 4 GiB/s. Pools of 16 to 64 clients reached
38/// that limit, but a pool of 10 did not.
39const MIN_CLIENTS: usize = 32;
40
41/// Metrics for one pool. A pool created with a clone of these metrics replaces
42/// the pool they described before.
43#[derive(Clone, Debug)]
44pub struct HttpPoolMetrics {
45    requests: IntCounter,
46    pool: Arc<Mutex<Weak<Pool>>>,
47}
48
49impl HttpPoolMetrics {
50    /// `requests` counts the requests sent through the pool, including the
51    /// SDK's retries.
52    pub fn new(requests: IntCounter) -> Self {
53        HttpPoolMetrics {
54            requests,
55            pool: Arc::default(),
56        }
57    }
58
59    /// Metrics that no registry reports, for a pool that should not count
60    /// toward any reported metrics.
61    pub fn unreported() -> Self {
62        HttpPoolMetrics::new(IntCounter::new("unreported", "unreported").expect("valid counter"))
63    }
64
65    /// The clients in the pool, whether or not any request uses them.
66    pub fn clients(&self) -> usize {
67        self.client_in_flight().len()
68    }
69
70    /// The clients in the pool with at least one request in flight.
71    pub fn clients_in_use(&self) -> usize {
72        self.client_in_flight()
73            .into_iter()
74            .filter(|&count| count > 0)
75            .count()
76    }
77
78    /// The requests in flight through the pool.
79    pub fn in_flight(&self) -> usize {
80        self.client_in_flight().into_iter().sum()
81    }
82
83    /// The fewest and the most requests in flight on any one client in the
84    /// pool. `(0, 0)` if there is no pool.
85    pub fn client_in_flight_range(&self) -> (usize, usize) {
86        let counts = self.client_in_flight();
87        let fewest = counts.iter().min().copied().unwrap_or(0);
88        let most = counts.iter().max().copied().unwrap_or(0);
89        (fewest, most)
90    }
91
92    /// The requests in flight on each client in the pool.
93    fn client_in_flight(&self) -> Vec<usize> {
94        let Some(pool) = self.pool.lock().expect("lock poisoned").upgrade() else {
95            return Vec::new();
96        };
97        let members = pool.members.lock().expect("lock poisoned");
98        members.iter().map(Member::in_flight).collect()
99    }
100}
101
102#[derive(Debug)]
103pub(crate) struct PooledHttpClient {
104    pool: Arc<Pool>,
105}
106
107impl PooledHttpClient {
108    pub(crate) fn new(
109        new_client: fn() -> SharedHttpClient,
110        max_in_flight: impl Fn() -> usize + Send + Sync + 'static,
111        metrics: HttpPoolMetrics,
112    ) -> Self {
113        PooledHttpClient {
114            pool: Pool::new(new_client, max_in_flight, metrics),
115        }
116    }
117}
118
119impl HttpClient for PooledHttpClient {
120    fn http_connector(
121        &self,
122        settings: &HttpConnectorSettings,
123        components: &RuntimeComponents,
124    ) -> SharedHttpConnector {
125        SharedHttpConnector::new(PooledConnector {
126            pool: Arc::clone(&self.pool),
127            settings: settings.clone(),
128            components: components.clone(),
129        })
130    }
131
132    fn validate_base_client_config(
133        &self,
134        runtime_components: &RuntimeComponentsBuilder,
135        cfg: &ConfigBag,
136    ) -> Result<(), BoxError> {
137        self.pool
138            .first()
139            .validate_base_client_config(runtime_components, cfg)
140    }
141
142    fn validate_final_config(
143        &self,
144        runtime_components: &RuntimeComponents,
145        cfg: &ConfigBag,
146    ) -> Result<(), BoxError> {
147        self.pool
148            .first()
149            .validate_final_config(runtime_components, cfg)
150    }
151
152    fn connector_metadata(&self) -> Option<ConnectorMetadata> {
153        self.pool.first().connector_metadata()
154    }
155}
156
157#[derive(Debug)]
158struct Pool {
159    new_client: fn() -> SharedHttpClient,
160    max_in_flight: Getter<usize>,
161    metrics: HttpPoolMetrics,
162    /// Holds at least [`MIN_CLIENTS`]. Members are never removed. An idle member's connections
163    /// close on the SDK client's idle timeout, so it costs only memory.
164    members: Mutex<Vec<Member>>,
165}
166
167impl Pool {
168    fn new(
169        new_client: fn() -> SharedHttpClient,
170        max_in_flight: impl Fn() -> usize + Send + Sync + 'static,
171        metrics: HttpPoolMetrics,
172    ) -> Arc<Self> {
173        let pool = Arc::new(Pool {
174            new_client,
175            max_in_flight: Getter::new(max_in_flight),
176            metrics,
177            members: Mutex::new(
178                (0..MIN_CLIENTS)
179                    .map(|_| Member::new(new_client()))
180                    .collect(),
181            ),
182        });
183        *pool.metrics.pool.lock().expect("lock poisoned") = Arc::downgrade(&pool);
184        pool
185    }
186
187    fn first(&self) -> SharedHttpClient {
188        let members = self.members.lock().expect("lock poisoned");
189        members[0].client.clone()
190    }
191
192    /// Picks the member with the fewest requests in flight, adding a member
193    /// if every one is at `max_in_flight`. The request counts against the
194    /// member until the returned clone is dropped.
195    fn pick(&self) -> Member {
196        // A limit of 0 would add a member on every request, and members are never removed.
197        let max_in_flight = self.max_in_flight.get().max(1);
198        let mut members = self.members.lock().expect("lock poisoned");
199        let least = members
200            .iter()
201            .min_by_key(|member| member.in_flight())
202            .expect("never empty");
203        // Cloned while holding the lock, so concurrent picks see the count.
204        if least.in_flight() < max_in_flight {
205            least.clone()
206        } else {
207            let member = Member::new((self.new_client)());
208            members.push(member.clone());
209            member
210        }
211    }
212}
213
214#[derive(Clone, Debug)]
215struct Member {
216    client: SharedHttpClient,
217    /// Each request in flight holds a clone.
218    token: Arc<()>,
219}
220
221impl Member {
222    fn new(client: SharedHttpClient) -> Self {
223        Member {
224            client,
225            token: Arc::new(()),
226        }
227    }
228
229    fn in_flight(&self) -> usize {
230        Arc::strong_count(&self.token) - 1
231    }
232}
233
234/// Picks a member per request, so the in-flight counts stay right however long
235/// the SDK keeps a connector.
236#[derive(Debug)]
237struct PooledConnector {
238    pool: Arc<Pool>,
239    settings: HttpConnectorSettings,
240    components: RuntimeComponents,
241}
242
243impl HttpConnector for PooledConnector {
244    fn call(&self, request: HttpRequest) -> HttpConnectorFuture {
245        let Member { client, token } = self.pool.pick();
246        self.pool.metrics.requests.inc();
247        let response = client
248            .http_connector(&self.settings, &self.components)
249            .call(request);
250        HttpConnectorFuture::new(async move {
251            let mut response = response.await?;
252            // A response holds its HTTP/2 stream open until its body is read
253            // or dropped, so the request stays in flight until then.
254            //
255            // NOTE: The token rides in `map_err` because `MapErr` passes through the body's
256            // size hint and `MapFrame` drops it.
257            let body = response.take_body().map_err(move |err| {
258                let _token = &token;
259                err
260            });
261            *response.body_mut() = SdkBody::from_body_1_x(body);
262            Ok(response)
263        })
264    }
265}
266
267/// Sends each request through `pooled` or `plain`, as `use_pool` says when the
268/// request is sent.
269#[derive(Debug)]
270pub(crate) struct SwitchedHttpClient {
271    plain: SharedHttpClient,
272    pooled: SharedHttpClient,
273    use_pool: Getter<bool>,
274}
275
276impl SwitchedHttpClient {
277    pub(crate) fn new(
278        plain: SharedHttpClient,
279        pooled: SharedHttpClient,
280        use_pool: impl Fn() -> bool + Send + Sync + 'static,
281    ) -> Self {
282        SwitchedHttpClient {
283            plain,
284            pooled,
285            use_pool: Getter::new(use_pool),
286        }
287    }
288}
289
290/// A setting read each time it is used, so a change applies to existing clients.
291#[derive(Clone)]
292struct Getter<T>(Arc<dyn Fn() -> T + Send + Sync>);
293
294impl<T> Getter<T> {
295    fn new(get: impl Fn() -> T + Send + Sync + 'static) -> Self {
296        Getter(Arc::new(get))
297    }
298
299    fn get(&self) -> T {
300        (self.0)()
301    }
302}
303
304impl<T> std::fmt::Debug for Getter<T> {
305    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
306        f.write_str("Getter")
307    }
308}
309
310impl HttpClient for SwitchedHttpClient {
311    fn http_connector(
312        &self,
313        settings: &HttpConnectorSettings,
314        components: &RuntimeComponents,
315    ) -> SharedHttpConnector {
316        SharedHttpConnector::new(SwitchedConnector {
317            plain: self.plain.http_connector(settings, components),
318            pooled: self.pooled.http_connector(settings, components),
319            use_pool: self.use_pool.clone(),
320        })
321    }
322
323    fn validate_base_client_config(
324        &self,
325        runtime_components: &RuntimeComponentsBuilder,
326        cfg: &ConfigBag,
327    ) -> Result<(), BoxError> {
328        self.plain
329            .validate_base_client_config(runtime_components, cfg)?;
330        self.pooled
331            .validate_base_client_config(runtime_components, cfg)
332    }
333
334    fn validate_final_config(
335        &self,
336        runtime_components: &RuntimeComponents,
337        cfg: &ConfigBag,
338    ) -> Result<(), BoxError> {
339        self.plain.validate_final_config(runtime_components, cfg)?;
340        self.pooled.validate_final_config(runtime_components, cfg)
341    }
342
343    fn connector_metadata(&self) -> Option<ConnectorMetadata> {
344        self.plain.connector_metadata()
345    }
346}
347
348/// Asks `use_pool` per request, so a change applies however long the SDK keeps
349/// a connector.
350#[derive(Debug)]
351struct SwitchedConnector {
352    plain: SharedHttpConnector,
353    pooled: SharedHttpConnector,
354    use_pool: Getter<bool>,
355}
356
357impl HttpConnector for SwitchedConnector {
358    fn call(&self, request: HttpRequest) -> HttpConnectorFuture {
359        if self.use_pool.get() {
360            self.pooled.call(request)
361        } else {
362            self.plain.call(request)
363        }
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    use aws_smithy_runtime_api::client::orchestrator::HttpResponse;
370    use aws_smithy_runtime_api::client::result::ConnectorError;
371    use aws_smithy_runtime_api::http::StatusCode;
372
373    use super::*;
374
375    #[derive(Debug)]
376    struct UnusedClient;
377
378    impl HttpClient for UnusedClient {
379        fn http_connector(
380            &self,
381            _: &HttpConnectorSettings,
382            _: &RuntimeComponents,
383        ) -> SharedHttpConnector {
384            unimplemented!()
385        }
386    }
387
388    const MAX_IN_FLIGHT: usize = 3;
389
390    fn pool(max_in_flight: impl Fn() -> usize + Send + Sync + 'static) -> Arc<Pool> {
391        Pool::new(
392            || SharedHttpClient::new(UnusedClient),
393            max_in_flight,
394            HttpPoolMetrics::unreported(),
395        )
396    }
397
398    fn in_flight(pool: &Pool) -> Vec<usize> {
399        let members = pool.members.lock().expect("lock poisoned");
400        members.iter().map(Member::in_flight).collect()
401    }
402
403    #[mz_ore::test]
404    fn grows_when_every_member_is_full() {
405        let pool = pool(|| MAX_IN_FLIGHT);
406        let metrics = pool.metrics.clone();
407        assert_eq!(in_flight(&pool), [0; MIN_CLIENTS]);
408        assert_eq!(metrics.clients(), MIN_CLIENTS);
409        let full = [MAX_IN_FLIGHT; MIN_CLIENTS];
410        let mut held: Vec<_> = (0..MIN_CLIENTS * MAX_IN_FLIGHT)
411            .map(|_| pool.pick())
412            .collect();
413        assert_eq!(in_flight(&pool), full);
414        held.push(pool.pick());
415        assert_eq!(in_flight(&pool), [&full[..], &[1]].concat());
416        held.push(pool.pick());
417        assert_eq!(in_flight(&pool), [&full[..], &[2]].concat());
418        assert_eq!(metrics.clients(), MIN_CLIENTS + 1);
419        drop(held);
420        assert_eq!(in_flight(&pool), [0; MIN_CLIENTS + 1]);
421        drop(pool);
422        assert_eq!(metrics.clients(), 0, "a dropped pool removes its clients");
423    }
424
425    #[mz_ore::test]
426    fn picks_least_loaded_member() {
427        let pool = pool(|| MAX_IN_FLIGHT);
428        let mut held: Vec<_> = (0..MIN_CLIENTS).map(|_| pool.pick()).collect();
429        assert_eq!(in_flight(&pool), [1; MIN_CLIENTS]);
430        drop(held.remove(3));
431        held.push(pool.pick());
432        assert_eq!(in_flight(&pool), [1; MIN_CLIENTS]);
433        held.push(pool.pick());
434        assert_eq!(in_flight(&pool), [&[2][..], &[1; MIN_CLIENTS - 1]].concat());
435    }
436
437    #[mz_ore::test]
438    fn max_in_flight_change_applies_to_next_pick() {
439        use std::sync::atomic::{AtomicUsize, Ordering};
440
441        let max_in_flight = Arc::new(AtomicUsize::new(1));
442        let pool = pool({
443            let max_in_flight = Arc::clone(&max_in_flight);
444            move || max_in_flight.load(Ordering::Relaxed)
445        });
446        #[allow(clippy::collection_is_never_read)]
447        let mut held: Vec<_> = (0..MIN_CLIENTS).map(|_| pool.pick()).collect();
448        max_in_flight.store(2, Ordering::Relaxed);
449        held.push(pool.pick());
450        assert_eq!(in_flight(&pool), [&[2][..], &[1; MIN_CLIENTS - 1]].concat());
451        max_in_flight.store(1, Ordering::Relaxed);
452        held.push(pool.pick());
453        assert_eq!(
454            in_flight(&pool),
455            [&[2][..], &[1; MIN_CLIENTS - 1], &[1]].concat()
456        );
457        max_in_flight.store(0, Ordering::Relaxed);
458        held.push(pool.pick());
459        assert_eq!(
460            in_flight(&pool),
461            [&[2][..], &[1; MIN_CLIENTS - 1], &[1, 1]].concat(),
462            "0 acts as 1",
463        );
464    }
465
466    #[mz_ore::test]
467    fn metrics_follow_newest_pool() {
468        let old = pool(|| MAX_IN_FLIGHT);
469        let metrics = old.metrics.clone();
470        let _old_held = old.pick();
471        assert_eq!(metrics.in_flight(), 1);
472        let new = Pool::new(
473            || SharedHttpClient::new(UnusedClient),
474            || MAX_IN_FLIGHT,
475            metrics.clone(),
476        );
477        assert_eq!(
478            metrics.in_flight(),
479            0,
480            "the newer pool replaces the older one"
481        );
482        let mut held: Vec<_> = (0..3).map(|_| new.pick()).collect();
483        assert_eq!(metrics.clients_in_use(), 3);
484        held.extend((3..MIN_CLIENTS + 1).map(|_| new.pick()));
485        assert_eq!(metrics.client_in_flight_range(), (1, 2));
486        assert_eq!(metrics.clients_in_use(), MIN_CLIENTS);
487        drop(held);
488        assert_eq!(metrics.client_in_flight_range(), (0, 0));
489        assert_eq!(metrics.clients_in_use(), 0);
490    }
491
492    /// Answers every request with a `hello` body, or with an error if `fail`.
493    #[derive(Debug)]
494    struct FakeClient {
495        fail: bool,
496    }
497
498    impl HttpClient for FakeClient {
499        fn http_connector(
500            &self,
501            _: &HttpConnectorSettings,
502            _: &RuntimeComponents,
503        ) -> SharedHttpConnector {
504            SharedHttpConnector::new(FakeConnector { fail: self.fail })
505        }
506    }
507
508    #[derive(Debug)]
509    struct FakeConnector {
510        fail: bool,
511    }
512
513    impl HttpConnector for FakeConnector {
514        fn call(&self, _: HttpRequest) -> HttpConnectorFuture {
515            if self.fail {
516                HttpConnectorFuture::ready(Err(ConnectorError::other("fake".into(), None)))
517            } else {
518                let status = StatusCode::try_from(200).expect("valid status");
519                HttpConnectorFuture::ready(Ok(HttpResponse::new(status, SdkBody::from("hello"))))
520            }
521        }
522    }
523
524    fn connector(pool: &Arc<Pool>) -> PooledConnector {
525        PooledConnector {
526            pool: Arc::clone(pool),
527            settings: HttpConnectorSettings::builder().build(),
528            components: RuntimeComponentsBuilder::for_tests()
529                .build()
530                .expect("valid components"),
531        }
532    }
533
534    #[mz_ore::test(tokio::test)]
535    async fn request_counts_until_body_dropped() {
536        let pool = Pool::new(
537            || SharedHttpClient::new(FakeClient { fail: false }),
538            || MAX_IN_FLIGHT,
539            HttpPoolMetrics::unreported(),
540        );
541        let response = connector(&pool).call(HttpRequest::empty());
542        assert_eq!(pool.metrics.in_flight(), 1);
543        assert_eq!(pool.metrics.requests.get(), 1);
544        let mut response = response.await.expect("fake response");
545        assert_eq!(pool.metrics.in_flight(), 1);
546        assert_eq!(
547            response.body().content_length(),
548            Some(5),
549            "wrapping the body keeps its size hint",
550        );
551        let body = response.take_body();
552        assert_eq!(pool.metrics.in_flight(), 1);
553        let bytes = body.collect().await.expect("fake body").to_bytes();
554        assert_eq!(bytes, "hello");
555        assert_eq!(pool.metrics.in_flight(), 0);
556    }
557
558    #[mz_ore::test(tokio::test)]
559    async fn switch_applies_to_existing_connector() {
560        use std::sync::atomic::{AtomicBool, Ordering};
561
562        let use_pool = Arc::new(AtomicBool::new(false));
563        let client = SwitchedHttpClient::new(
564            SharedHttpClient::new(FakeClient { fail: true }),
565            SharedHttpClient::new(FakeClient { fail: false }),
566            {
567                let use_pool = Arc::clone(&use_pool);
568                move || use_pool.load(Ordering::Relaxed)
569            },
570        );
571        let connector = client.http_connector(
572            &HttpConnectorSettings::builder().build(),
573            &RuntimeComponentsBuilder::for_tests()
574                .build()
575                .expect("valid components"),
576        );
577        connector
578            .call(HttpRequest::empty())
579            .await
580            .expect_err("plain client");
581        use_pool.store(true, Ordering::Relaxed);
582        connector
583            .call(HttpRequest::empty())
584            .await
585            .expect("pooled client");
586    }
587
588    #[mz_ore::test(tokio::test)]
589    async fn failed_or_abandoned_request_stops_counting() {
590        let pool = Pool::new(
591            || SharedHttpClient::new(FakeClient { fail: true }),
592            || MAX_IN_FLIGHT,
593            HttpPoolMetrics::unreported(),
594        );
595        let response = connector(&pool).call(HttpRequest::empty());
596        assert_eq!(pool.metrics.in_flight(), 1);
597        response.await.expect_err("fake failure");
598        assert_eq!(pool.metrics.in_flight(), 0);
599
600        let response = connector(&pool).call(HttpRequest::empty());
601        assert_eq!(pool.metrics.in_flight(), 1);
602        drop(response);
603        assert_eq!(pool.metrics.in_flight(), 0);
604    }
605}
mz_aws_util/http_pool.rs

mz_aws_util/
http_pool.rs
