risingwave_object_store/object/opendal_engine/
rotating_transport.rs

1// Copyright 2026 RisingWave Labs
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 std::sync::Arc;
16use std::time::Duration;
17
18use arc_swap::ArcSwap;
19use http::{Request, Response};
20use opendal::{Buffer, Error, ErrorKind, HttpBody, HttpTransport, Result};
21use opendal_http_transport_reqwest::ReqwestTransport;
22use thiserror_ext::AsReport as _;
23use tokio::sync::oneshot;
24use tokio::time::{Instant, MissedTickBehavior, interval_at};
25use tokio_retry::strategy::jitter;
26
27/// Periodically replace the HTTP client while existing response bodies finish normally.
28pub(super) struct RotatingTransport {
29    current: Arc<ArcSwap<ReqwestTransport>>,
30    // Dropping the last transport handle closes the channel and stops rotation.
31    _shutdown_tx: oneshot::Sender<()>,
32}
33
34impl RotatingTransport {
35    pub(super) fn new(
36        mut factory: impl FnMut() -> Result<ReqwestTransport> + Send + 'static,
37        interval: Duration,
38    ) -> Result<Self> {
39        if interval.is_zero() {
40            return Err(Error::new(
41                ErrorKind::ConfigInvalid,
42                "HTTP client rotation interval must be positive",
43            ));
44        }
45        let current = Arc::new(ArcSwap::from_pointee(factory()?));
46        // Delay the first refresh and spread its phase across operators. Later
47        // refreshes use the configured interval without accumulating build time.
48        let start = Instant::now()
49            .checked_add(interval)
50            .and_then(|start| start.checked_add(jitter(interval)))
51            .ok_or_else(|| {
52                Error::new(
53                    ErrorKind::ConfigInvalid,
54                    "HTTP client rotation deadline overflow",
55                )
56            })?;
57        let mut ticker = interval_at(start, interval);
58        ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
59        let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
60        let _task = tokio::spawn({
61            let current = current.clone();
62            async move {
63                loop {
64                    tokio::select! {
65                        biased;
66                        _ = &mut shutdown_rx => break,
67                        _ = ticker.tick() => {}
68                    }
69                    match factory() {
70                        Ok(client) => current.store(Arc::new(client)),
71                        Err(err) => {
72                            tracing::warn!(error = %err.as_report(), "failed to rotate object store HTTP client")
73                        }
74                    }
75                }
76            }
77        });
78        Ok(Self {
79            current,
80            _shutdown_tx: shutdown_tx,
81        })
82    }
83}
84
85impl HttpTransport for RotatingTransport {
86    async fn fetch(&self, request: Request<Buffer>) -> Result<Response<HttpBody>> {
87        self.current.load_full().fetch(request).await
88    }
89}
90
91#[cfg(all(test, not(madsim)))]
92mod tests {
93    use std::convert::Infallible;
94    use std::net::SocketAddr;
95    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
96    use std::sync::{Mutex, Weak};
97
98    use bytes::Bytes;
99    use futures::stream;
100    use http_body_util::StreamBody;
101    use hyper::body::Frame;
102    use hyper::service::service_fn;
103    use hyper_util::rt::{TokioExecutor, TokioIo};
104    use opendal::raw::oio::ReadStream as _;
105    use reqwest::dns::{Addrs, Name, Resolve, Resolving};
106    use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};
107    use tokio::net::TcpListener;
108    use tokio::sync::{Notify, mpsc};
109    use tokio::task::{JoinHandle, JoinSet};
110
111    use super::*;
112
113    const INTERVAL: Duration = Duration::from_secs(10);
114
115    async fn advance_time(duration: Duration) {
116        // Resume before socket I/O so Tokio cannot auto-advance past idle pool timers.
117        tokio::time::pause();
118        tokio::task::yield_now().await;
119        tokio::time::advance(duration).await;
120        tokio::task::yield_now().await;
121        tokio::time::resume();
122    }
123
124    #[tokio::test]
125    async fn periodic_rotation_keeps_current_client_on_failure() {
126        let fail = Arc::new(AtomicBool::new(false));
127        let builds = Arc::new(AtomicUsize::new(0));
128        let transport = RotatingTransport::new(
129            {
130                let fail = fail.clone();
131                let builds = builds.clone();
132                move || {
133                    builds.fetch_add(1, Ordering::Relaxed);
134                    if fail.load(Ordering::Relaxed) {
135                        return Err(Error::new(ErrorKind::Unexpected, "injected build failure"));
136                    }
137                    Ok(ReqwestTransport::new(reqwest::Client::new()))
138                }
139            },
140            INTERVAL,
141        )
142        .unwrap();
143        let original = transport.current.load_full();
144        tokio::task::yield_now().await;
145        assert_eq!(builds.load(Ordering::Relaxed), 1);
146        advance_time(INTERVAL / 2).await;
147        assert_eq!(builds.load(Ordering::Relaxed), 1);
148
149        // Refresh without requests, skipping missed ticks instead of building a burst.
150        fail.store(true, Ordering::Relaxed);
151        advance_time(INTERVAL * 100).await;
152        assert_eq!(builds.load(Ordering::Relaxed), 2);
153        assert!(Arc::ptr_eq(&original, &transport.current.load_full()));
154
155        fail.store(false, Ordering::Relaxed);
156        advance_time(INTERVAL * 2).await;
157        let replacement = transport.current.load_full();
158        assert_eq!(builds.load(Ordering::Relaxed), 3);
159        assert!(!Arc::ptr_eq(&original, &replacement));
160        let original_weak = Arc::downgrade(&original);
161        drop(original);
162        assert!(original_weak.upgrade().is_none());
163    }
164
165    #[tokio::test]
166    async fn streaming_body_survives_client_rotation_and_shutdown() {
167        for http2 in [false, true] {
168            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
169            let address = listener.local_addr().unwrap();
170            let (body_tx, body_rx) = mpsc::unbounded_channel::<Bytes>();
171            let body_rx = Arc::new(Mutex::new(Some(body_rx)));
172            let requests = Arc::new(AtomicUsize::new(0));
173            let server = tokio::spawn({
174                let requests = requests.clone();
175                async move {
176                    let (socket, _) = listener.accept().await.unwrap();
177                    let service = service_fn(move |_request| {
178                        requests.fetch_add(1, Ordering::Relaxed);
179                        let body_rx = body_rx.lock().unwrap().take().unwrap();
180                        async move {
181                            let stream = stream::unfold(body_rx, |mut rx| async move {
182                                rx.recv()
183                                    .await
184                                    .map(|bytes| (Ok::<_, Infallible>(Frame::data(bytes)), rx))
185                            });
186                            Ok::<_, Infallible>(
187                                Response::builder()
188                                    .header("content-length", "6")
189                                    .body(StreamBody::new(stream))
190                                    .unwrap(),
191                            )
192                        }
193                    });
194                    if http2 {
195                        hyper::server::conn::http2::Builder::new(TokioExecutor::new())
196                            .serve_connection(TokioIo::new(socket), service)
197                            .await
198                            .unwrap();
199                    } else {
200                        hyper::server::conn::http1::Builder::new()
201                            .serve_connection(TokioIo::new(socket), service)
202                            .await
203                            .unwrap();
204                    }
205                }
206            });
207            let (factory_tx, factory_rx) = oneshot::channel::<()>();
208            let transport = RotatingTransport::new(
209                move || {
210                    // Dropping the factory closes this channel and confirms task exit.
211                    let _keep_alive = &factory_tx;
212                    let builder = reqwest::Client::builder().no_proxy();
213                    let builder = if http2 {
214                        builder.http2_prior_knowledge()
215                    } else {
216                        builder.http1_only()
217                    };
218                    Ok(ReqwestTransport::new(builder.build().unwrap()))
219                },
220                INTERVAL,
221            )
222            .unwrap();
223            let current = Arc::downgrade(&transport.current);
224            let original = Arc::downgrade(&transport.current.load_full());
225            body_tx.send(Bytes::from_static(b"A")).unwrap();
226            let response = transport
227                .fetch(
228                    Request::builder()
229                        .uri(format!("http://{address}/held"))
230                        .body(Buffer::new())
231                        .unwrap(),
232                )
233                .await
234                .unwrap();
235            assert_eq!(
236                response.version(),
237                if http2 {
238                    http::Version::HTTP_2
239                } else {
240                    http::Version::HTTP_11
241                }
242            );
243            let mut body = response.into_body();
244            assert_eq!(body.read().await.unwrap().to_bytes(), "A");
245
246            advance_time(INTERVAL * 2).await;
247            assert!(original.upgrade().is_none());
248            drop(transport);
249            assert!(
250                tokio::time::timeout(Duration::from_secs(1), factory_rx)
251                    .await
252                    .unwrap()
253                    .is_err()
254            );
255            assert!(current.upgrade().is_none());
256
257            // The remaining bytes are produced only after the old client is gone
258            // and the rotation task has stopped. The native body still completes.
259            body_tx.send(Bytes::from_static(b"tail!")).unwrap();
260            drop(body_tx);
261            assert_eq!(
262                tokio::time::timeout(Duration::from_secs(2), body.to_buffer())
263                    .await
264                    .unwrap()
265                    .unwrap()
266                    .to_bytes(),
267                "tail!"
268            );
269            assert_eq!(requests.load(Ordering::Relaxed), 1);
270            server.abort();
271        }
272    }
273
274    struct Resolver(Mutex<SocketAddr>);
275
276    impl Resolve for Resolver {
277        fn resolve(&self, name: Name) -> Resolving {
278            assert_eq!(name.as_str(), "rotation.test");
279            let address = *self.0.lock().unwrap();
280            Box::pin(async move { Ok(Box::new(std::iter::once(address)) as Addrs) })
281        }
282    }
283
284    struct Server {
285        address: SocketAddr,
286        release_body: Arc<Notify>,
287        connections: Arc<AtomicUsize>,
288        task: JoinHandle<()>,
289    }
290
291    impl Drop for Server {
292        fn drop(&mut self) {
293            self.task.abort();
294        }
295    }
296
297    impl Server {
298        async fn start(name: &'static str) -> Self {
299            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
300            let address = listener.local_addr().unwrap();
301            let release_body = Arc::new(Notify::new());
302            let connections = Arc::new(AtomicUsize::new(0));
303            let task = tokio::spawn({
304                let release_body = release_body.clone();
305                let connections = connections.clone();
306                async move {
307                    let mut requests = JoinSet::new();
308                    loop {
309                        tokio::select! {
310                            accepted = listener.accept() => {
311                                let (socket, _) = accepted.unwrap();
312                                connections.fetch_add(1, Ordering::Relaxed);
313                                let release_body = release_body.clone();
314                                requests.spawn(async move {
315                                    let mut socket = BufReader::new(socket);
316                                    loop {
317                                        let mut request = String::new();
318                                        loop {
319                                            let mut line = String::new();
320                                            match socket.read_line(&mut line).await {
321                                                Ok(0) | Err(_) => return,
322                                                Ok(_) => {}
323                                            }
324                                            request.push_str(&line);
325                                            if line == "\r\n" { break; }
326                                        }
327                                        let headers = request.to_ascii_lowercase();
328                                        assert!(headers.contains("host: rotation.test\r\n"));
329                                        assert!(headers.contains("authorization: test-signature\r\n"));
330                                        assert!(headers.contains("range: bytes=0-5\r\n"));
331                                        let held = request.starts_with("GET /held ");
332                                        let broken = request.starts_with("GET /broken ");
333                                        let length = if held { 6 } else if broken { 100 } else { 1 };
334                                        let response = format!("HTTP/1.1 200 OK\r\nContent-Length: {length}\r\nx-backend: {name}\r\n\r\n");
335                                        socket.write_all(response.as_bytes()).await.unwrap();
336                                        if request.starts_with("HEAD ") { continue; }
337                                        socket.write_all(name.as_bytes()).await.unwrap();
338                                        if broken { return; }
339                                        if held {
340                                            release_body.notified().await;
341                                            socket.write_all(b"tail!").await.unwrap();
342                                        }
343                                    }
344                                });
345                            }
346                            result = requests.join_next(), if !requests.is_empty() => {
347                                result.unwrap().unwrap();
348                            }
349                        }
350                    }
351                }
352            });
353            Self {
354                address,
355                release_body,
356                connections,
357                task,
358            }
359        }
360    }
361
362    fn request(path: &str) -> Request<Buffer> {
363        Request::builder()
364            .uri(format!("http://rotation.test/{path}"))
365            .header("authorization", "test-signature")
366            .header("range", "bytes=0-5")
367            .body(Buffer::new())
368            .unwrap()
369    }
370
371    async fn fetch_backend(transport: &impl HttpTransport) -> String {
372        let response = transport.fetch(request("object")).await.unwrap();
373        assert_eq!(response.status(), 200);
374        let expected = response.headers()["x-backend"].to_str().unwrap().to_owned();
375        let data = response.into_body().to_buffer().await.unwrap();
376        assert_eq!(data.to_bytes(), expected);
377        expected
378    }
379
380    #[tokio::test]
381    async fn dns_rotation_preserves_old_streams_and_releases_clients() {
382        let a = Server::start("A").await;
383        let b = Server::start("B").await;
384        let resolver = Arc::new(Resolver(Mutex::new(a.address)));
385        let fail = Arc::new(AtomicBool::new(false));
386        let transport = RotatingTransport::new(
387            {
388                let resolver = resolver.clone();
389                let fail = fail.clone();
390                move || {
391                    if fail.load(Ordering::Relaxed) {
392                        return Err(Error::new(ErrorKind::Unexpected, "injected build failure"));
393                    }
394                    Ok(ReqwestTransport::new(
395                        reqwest::Client::builder()
396                            .no_proxy()
397                            .dns_resolver(resolver.clone())
398                            .build()
399                            .unwrap(),
400                    ))
401                }
402            },
403            INTERVAL,
404        )
405        .unwrap();
406
407        for _ in 0..3 {
408            assert_eq!(fetch_backend(&transport).await, "A");
409        }
410        assert_eq!(a.connections.load(Ordering::Relaxed), 1);
411        let single_transport = ReqwestTransport::new(
412            reqwest::Client::builder()
413                .no_proxy()
414                .dns_resolver(resolver.clone())
415                .build()
416                .unwrap(),
417        );
418        assert_eq!(fetch_backend(&single_transport).await, "A");
419        *resolver.0.lock().unwrap() = b.address;
420        // DNS alone leaves a continuously reused healthy connection on A.
421        for _ in 0..3 {
422            assert_eq!(fetch_backend(&transport).await, "A");
423        }
424        let original: Weak<ReqwestTransport> = Arc::downgrade(&transport.current.load_full());
425        let mut held = transport.fetch(request("held")).await.unwrap().into_body();
426        assert_eq!(held.read().await.unwrap().to_bytes(), "A");
427
428        fail.store(true, Ordering::Relaxed);
429        advance_time(INTERVAL * 2).await;
430        // The single-client control keeps reusing A while a failed refresh leaves
431        // the rotating transport's current client available to new requests.
432        assert_eq!(fetch_backend(&single_transport).await, "A");
433        assert!(Arc::ptr_eq(
434            &original.upgrade().unwrap(),
435            &transport.current.load_full()
436        ));
437        fetch_backend(&transport).await;
438        fail.store(false, Ordering::Relaxed);
439        advance_time(INTERVAL * 2).await;
440        assert_eq!(fetch_backend(&transport).await, "B");
441        assert!(original.upgrade().is_none());
442        a.release_body.notify_one();
443        assert_eq!(held.to_buffer().await.unwrap().to_bytes(), "tail!");
444        let mut head = request("object");
445        *head.method_mut() = http::Method::HEAD;
446        assert!(
447            transport
448                .fetch(head)
449                .await
450                .unwrap()
451                .into_body()
452                .to_buffer()
453                .await
454                .unwrap()
455                .is_empty()
456        );
457
458        let mut broken = transport
459            .fetch(request("broken"))
460            .await
461            .unwrap()
462            .into_body();
463        let error = broken.to_buffer().await.unwrap_err();
464        assert!(error.is_temporary());
465    }
466}
risingwave_object_store/object/opendal_engine/rotating_transport.rs

risingwave_object_store/object/opendal_engine/
rotating_transport.rs
