diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index caede75b90f..94af0414d8f 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -8747,6 +8747,29 @@ dependencies = [ "stateright", ] +[[package]] +name = "quickwit-http-client" +version = "0.9.0" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "futures", + "http 1.4.2", + "http-body 1.0.1", + "http-body-util", + "httparse", + "rustls 0.23.41", + "rustls-native-certs", + "rustls-pemfile", + "socket2 0.6.4", + "thiserror 2.0.18", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", + "tracing", +] + [[package]] name = "quickwit-index-management" version = "0.9.0" @@ -9403,6 +9426,7 @@ dependencies = [ "aws-credential-types", "aws-sdk-s3", "aws-smithy-runtime", + "aws-smithy-runtime-api", "aws-smithy-types", "azure_core", "azure_identity", @@ -9428,6 +9452,7 @@ dependencies = [ "quickwit-aws", "quickwit-common", "quickwit-config", + "quickwit-http-client", "quickwit-metrics", "quickwit-proto", "regex", @@ -12253,6 +12278,7 @@ dependencies = [ "futures-core", "futures-io", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index 8904ed1dd3c..2312a0552d3 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -17,6 +17,7 @@ members = [ "quickwit-directories", "quickwit-doc-mapper", "quickwit-dst", + "quickwit-http-client", "quickwit-index-management", "quickwit-indexing", "quickwit-ingest", @@ -64,6 +65,7 @@ default-members = [ "quickwit-datetime", "quickwit-directories", "quickwit-doc-mapper", + "quickwit-http-client", "quickwit-index-management", "quickwit-indexing", "quickwit-ingest", @@ -154,6 +156,8 @@ http = "1.4" http-body = "1.0" http-body-util = "0.1" http-serde = "2.1" +httparse = "1.10" +socket2 = { version = "0.6", features = ["all"] } humantime = "2.3" hyper = { version = "1.9", features = ["client", "http1", "http2", "server"] } hyper-rustls = "0.27" @@ -253,6 +257,7 @@ reqwest-retry = "0.8" rust-embed = "8.11" rustc-hash = "2.1" rustls = "0.23" +rustls-native-certs = "0.8" rustls-pemfile = "2.2" sea-query = { version = "0.32" } sea-query-binder = { version = "0.7", features = [ @@ -296,6 +301,7 @@ tokio-metrics = { version = "0.5", features = ["rt"] } tokio-rustls = { version = "0.26", default-features = false } tokio-stream = { version = "0.1", features = ["sync"] } tokio-util = { version = "0.7", default-features = false, features = [ + "codec", "compat", "io-util", ] } @@ -403,6 +409,7 @@ quickwit-df-core = { path = "quickwit-df-core" } quickwit-directories = { path = "quickwit-directories" } quickwit-doc-mapper = { path = "quickwit-doc-mapper" } quickwit-dst = { path = "quickwit-dst" } +quickwit-http-client = { path = "quickwit-http-client" } quickwit-index-management = { path = "quickwit-index-management" } quickwit-indexing = { path = "quickwit-indexing" } quickwit-ingest = { path = "quickwit-ingest" } diff --git a/quickwit/quickwit-http-client/Cargo.toml b/quickwit/quickwit-http-client/Cargo.toml new file mode 100644 index 00000000000..0a2360e246f --- /dev/null +++ b/quickwit/quickwit-http-client/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "quickwit-http-client" +description = "A fast HTTP/1.1 client for Quickwit object storage" +version.workspace = true +edition.workspace = true +homepage.workspace = true +documentation.workspace = true +repository.workspace = true +authors.workspace = true +license.workspace = true + +[features] +# enable s3 adaptor +s3 = ["dep:aws-smithy-runtime-api", "dep:aws-smithy-types"] + +[dependencies] +bytes = { workspace = true } +futures = { workspace = true } +http = { workspace = true } +http-body = { workspace = true } +http-body-util = { workspace = true } +httparse = { workspace = true } +rustls = { workspace = true } +rustls-native-certs = { workspace = true } +socket2 = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true } +tokio-rustls = { workspace = true, features = ["aws_lc_rs", "tls12"] } +tokio-util = { workspace = true, features = ["rt"] } +tracing = { workspace = true } + +aws-smithy-runtime-api = { workspace = true, features = ["client", "http-1x"], optional = true } +aws-smithy-types = { workspace = true, optional = true } + +[dev-dependencies] +rustls-pemfile = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } diff --git a/quickwit/quickwit-http-client/src/body.rs b/quickwit/quickwit-http-client/src/body.rs new file mode 100644 index 00000000000..0bb38d6dffc --- /dev/null +++ b/quickwit/quickwit-http-client/src/body.rs @@ -0,0 +1,603 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::io::{Cursor, ErrorKind}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; + +use bytes::Bytes; +use futures::{Future, Stream}; +use http_body::{Body, Frame, SizeHint}; +use tokio::io::{AsyncRead, AsyncReadExt, Chain, ReadBuf}; +use tokio::time::Sleep; +use tokio_util::codec::FramedRead; + +use crate::error::HttpError; +use crate::response::BodyStrategy; + +mod decoder; +use decoder::{DecodedItem, HttpBodyDecoder}; + +/// Frame coalescing target. Read-ahead can make frames larger than `target`, +/// while a known body no larger than `target` is yielded in one frame. +#[derive(Clone, Copy, Debug)] +pub struct BufferHint { + pub target: usize, +} + +impl BufferHint { + /// Default block size balancing TTFB and allocation overhead. + pub const DEFAULT: BufferHint = BufferHint { target: 256 * 1024 }; +} + +/// Default per-read idle timeout. +pub const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); + +pub(crate) const MIN_TARGET: usize = 8 * 1024; + +type PrefixedReader = Chain, R>; +type BodyFramedReader = FramedRead>, HttpBodyDecoder>; + +/// Adds a progress-based idle timeout to an [`AsyncRead`]. +struct IdleTimeoutReader { + inner: R, + timeout: Duration, + sleep: Pin>, + sleep_armed: bool, +} + +impl IdleTimeoutReader { + fn new(inner: R, timeout: Duration) -> Self { + Self { + inner, + timeout, + sleep: Box::pin(tokio::time::sleep(Duration::ZERO)), + sleep_armed: false, + } + } + + fn reset_timeout(&mut self) { + self.sleep + .as_mut() + .reset(tokio::time::Instant::now() + self.timeout); + self.sleep_armed = true; + } + + /// Reclaims the underlying reader. + pub(crate) fn into_inner(self) -> R { + self.inner + } +} + +impl AsyncRead for IdleTimeoutReader { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + read_buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if !this.sleep_armed { + this.reset_timeout(); + } + + let filled_before = read_buf.filled().len(); + match Pin::new(&mut this.inner).poll_read(cx, read_buf) { + Poll::Ready(Ok(())) => { + if read_buf.filled().len() > filled_before { + this.reset_timeout(); + } + Poll::Ready(Ok(())) + } + Poll::Ready(Err(error)) => Poll::Ready(Err(error)), + Poll::Pending => { + if this.sleep.as_mut().poll(cx).is_ready() { + this.sleep_armed = false; + Poll::Ready(Err(std::io::Error::new( + ErrorKind::TimedOut, + "response read timed out", + ))) + } else { + Poll::Pending + } + } + } + } +} + +struct KnownLen { + /// Body bytes read alongside headers, to prepend to the output until empty + leftover: Bytes, + /// Remaining body bytes to read. + remaining: usize, +} + +struct UntilClose { + /// Body bytes read alongside headers, to prepend to the output until empty + leftover: Bytes, +} + +enum FastState { + Known(KnownLen), + UntilClose(UntilClose), +} + +/// The outcome of one frame read on the fast path. The reader is always +/// returned so it can be stored back in `FastBody` (or pooled at EOS). +struct FastReadOutcome { + reader: R, + result: Result<(Bytes, bool), HttpError>, +} + +/// Read one frame of `want` bytes. +/// +/// Can short-read if is_known is false and we reach EOS. +/// May read a bit more than `want` bytes if more is available on the wire +/// depending on how BytesMut decide to actually allocate. +// we might win ever so slightly by having this be a handrolled state machine +// which we can manipulate without boxes, but as is i think it would either cost +// us a zero-initialization of a buffer, or some unsafe to interact with AsyncRead +// efficiently +async fn read_frame_fast( + mut reader: R, + want: usize, + leftover: Bytes, + is_known: bool, + total_remaining: usize, + consumed: usize, + read_timeout: Duration, +) -> FastReadOutcome { + use tokio::io::AsyncReadExt; + + let mut buf = bytes::BytesMut::with_capacity(want); + // Drain the leftover first + if !leftover.is_empty() { + buf.extend_from_slice(&leftover); + } + let mut sleep = Box::pin(tokio::time::sleep(read_timeout)); + while buf.len() < want { + sleep + .as_mut() + .reset(tokio::time::Instant::now() + read_timeout); + tokio::select! { + biased; + // this might read slighly over `want` depending on how bytes decides + // to over-allocated. we don't promise strict bound for this reason + res = reader.read_buf(&mut buf) => match res { + Ok(0) => { + if is_known { + let read = consumed + buf.len(); + return FastReadOutcome { + reader, + result: Err(HttpError::UnexpectedEof { + read, + expected: consumed + total_remaining, + }), + }; + } + return FastReadOutcome { + reader, + result: Ok((buf.freeze(), true)), + }; + } + Ok(_) => {} // try to read more until we have want bytes + Err(err) => { + return FastReadOutcome { + reader, + result: Err(HttpError::Io(err)), + }; + } + }, + _ = &mut sleep => { + return FastReadOutcome { + reader, + result: Err(HttpError::Timeout( + read_timeout, + "response read".to_string(), + )), + }; + } + } + } + FastReadOutcome { + reader, + result: Ok((buf.freeze(), false)), + } +} + +/// Fast path, when body is not chunked +struct FastBody { + // None only when lend it to read_frame_fast to please the borrowchecker + reader: Option, + // In-progress frame read. Only None when reader is Some, and Some when reader is None. + read_fut: Option> + Send + Sync>>>, + read_timeout: Duration, + target: usize, + // `None` once the body has reached EOS. + state: Option, + pool_hook: Option>, + poolable: bool, + clean_eos: bool, + consumed: usize, +} + +impl FastBody { + fn poll_frame( + &mut self, + cx: &mut Context<'_>, + ) -> Poll, HttpError>>> { + // If a read is in progress, poll it. + if let Some(fut) = self.read_fut.as_mut() { + match fut.as_mut().poll(cx) { + Poll::Ready(outcome) => { + self.read_fut = None; + self.reader = Some(outcome.reader); + return self.finish_frame(outcome.result); + } + Poll::Pending => return Poll::Pending, + } + } + // Start a new frame read. + let Some(state) = self.state.take() else { + self.clean_eos = true; + return Poll::Ready(None); + }; + let (want, leftover, is_known, total_remaining) = match &state { + FastState::Known(k) => { + if k.leftover.len() > k.remaining { + self.state = None; + return Poll::Ready(Some(Err(HttpError::InvalidLength( + "body longer than Content-Length".to_string(), + )))); + } + if k.remaining == 0 { + self.clean_eos = true; + return Poll::Ready(None); + } + ( + k.remaining.min(self.target), + k.leftover.clone(), + true, + k.remaining, + ) + } + FastState::UntilClose(u) => (self.target, u.leftover.clone(), false, 0), + }; + let reader = self + .reader + .take() + .expect("reader is Some when state is Some"); + let consumed = self.consumed; + let read_timeout = self.read_timeout; + let fut = Box::pin(read_frame_fast( + reader, + want, + leftover, + is_known, + total_remaining, + consumed, + read_timeout, + )); + self.read_fut = Some(fut); + self.state = Some(match state { + FastState::Known(k) => FastState::Known(KnownLen { + leftover: Bytes::new(), + remaining: k.remaining, + }), + FastState::UntilClose(_) => FastState::UntilClose(UntilClose { + leftover: Bytes::new(), + }), + }); + // Poll the future we just created. This will recurse at most once: read_fut is now set + self.poll_frame(cx) + } + + /// Processes the result of a frame read, updating state and pooling on EOS. + fn finish_frame( + &mut self, + result: Result<(Bytes, bool), HttpError>, + ) -> Poll, HttpError>>> { + match result { + Ok((bytes, eos)) => { + self.consumed += bytes.len(); + if let Some(FastState::Known(k)) = &mut self.state { + // `read_buf` ran read past `want`, it's fine as long as it doesn't go further + // than the actual body content (we don't do pipelining so server shouldn't + // send anything more) + if bytes.len() > k.remaining { + self.state = None; + return Poll::Ready(Some(Err(HttpError::InvalidLength( + "body longer than Content-Length".to_string(), + )))); + } + k.remaining -= bytes.len(); + } + if eos || matches!(&self.state, Some(FastState::Known(k)) if k.remaining == 0) { + self.clean_eos = true; + self.state = None; + self.release_to_pool(); + } + Poll::Ready(Some(Ok(Frame::data(bytes)))) + } + Err(err) => { + self.state = None; + Poll::Ready(Some(Err(err))) + } + } + } +} + +impl FastBody { + fn release_to_pool(&mut self) { + if self.clean_eos + && self.poolable + && let (Some(hook), Some(reader)) = (self.pool_hook.take(), self.reader.take()) + { + hook(reader); + } + } +} + +struct ChunkedBody { + framed: Option>, + pool_hook: Option>, + poolable: bool, + clean_eos: bool, +} + +impl ChunkedBody { + fn release_to_pool(&mut self) { + if self.poolable + && let (Some(hook), Some(framed)) = (self.pool_hook.take(), self.framed.take()) + { + let idle_reader = framed.into_inner(); + let chain = idle_reader.into_inner(); + let (_leftover, reader) = chain.into_inner(); + hook(reader); + } + } +} + +pub struct ResponseBody { + kind: BodyKind, + read_timeout: Duration, + consumed: usize, + done: bool, +} + +enum BodyKind { + /// Fast pass for unframed body (content-lenght or until EOS) + Fast(FastBody), + /// Transfer-Encoding: chunked path + Chunked(ChunkedBody), + /// Empty body / zero-length body + Complete { + reader: Option, + pool_hook: Option>, + poolable: bool, + }, +} + +impl std::fmt::Debug for ResponseBody { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ResponseBody") + .field("read_timeout", &self.read_timeout) + .field("consumed", &self.consumed) + .field("done", &self.done) + .finish_non_exhaustive() + } +} + +impl ResponseBody { + pub(crate) fn new( + reader: R, + strategy: BodyStrategy, + leftover: Bytes, + buffer_hint: BufferHint, + read_timeout: Duration, + pool_hook: Option>, + keep_alive: bool, + ) -> Self { + let target = buffer_hint.target.max(MIN_TARGET); + let leftover_len = leftover.len(); + // `UntilClose` ends on a peer EOF, connection isn't reusable. + let poolable = keep_alive && !matches!(strategy, BodyStrategy::UntilClose); + let kind = match strategy { + BodyStrategy::Empty | BodyStrategy::Known(0) => { + // An empty body with non-empty leftover means protocol desync, + // mark it non poolable so we close the connection + BodyKind::Complete { + reader: Some(reader), + pool_hook, + poolable: poolable && leftover_len == 0, + } + } + BodyStrategy::Known(len) => BodyKind::Fast(FastBody { + reader: Some(reader), + read_fut: None, + read_timeout, + target, + state: Some(FastState::Known(KnownLen { + leftover, + remaining: len, + })), + pool_hook, + poolable, + clean_eos: false, + consumed: 0, + }), + BodyStrategy::UntilClose => BodyKind::Fast(FastBody { + reader: Some(reader), + read_fut: None, + read_timeout, + target, + state: Some(FastState::UntilClose(UntilClose { leftover })), + pool_hook, + poolable, + clean_eos: false, + consumed: 0, + }), + BodyStrategy::Chunked => { + let decoder = HttpBodyDecoder::new(target); + let initial_capacity = decoder.initial_capacity(); + let prefixed_reader = AsyncReadExt::chain(Cursor::new(leftover), reader); + let timeout_reader = IdleTimeoutReader::new(prefixed_reader, read_timeout); + let framed = FramedRead::with_capacity(timeout_reader, decoder, initial_capacity); + BodyKind::Chunked(ChunkedBody { + framed: Some(framed), + pool_hook, + poolable, + clean_eos: false, + }) + } + }; + let done = matches!(kind, BodyKind::Complete { .. }); + Self { + kind, + read_timeout, + consumed: 0, + done, + } + } + + /// Number of body bytes yielded so far. + pub fn consumed(&self) -> usize { + self.consumed + } + + fn normalize_error(&self, error: HttpError) -> HttpError { + match error { + HttpError::Io(io_error) if io_error.kind() == ErrorKind::TimedOut => { + HttpError::Timeout(self.read_timeout, "response read".to_string()) + } + other => other, + } + } +} + +impl Body for ResponseBody { + type Data = Bytes; + type Error = HttpError; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + if this.done { + return Poll::Ready(None); + } + match &mut this.kind { + BodyKind::Complete { .. } => { + this.done = true; + Poll::Ready(None) + } + BodyKind::Fast(fast) => match fast.poll_frame(cx) { + Poll::Ready(Some(Ok(frame))) => { + this.consumed = fast.consumed; + if fast.state.is_none() { + this.done = true; + } + Poll::Ready(Some(Ok(frame))) + } + Poll::Ready(None) => { + this.done = true; + Poll::Ready(None) + } + Poll::Ready(Some(Err(err))) => { + this.done = true; + Poll::Ready(Some(Err(this.normalize_error(err)))) + } + Poll::Pending => Poll::Pending, + }, + BodyKind::Chunked(chunked) => { + let Some(framed) = chunked.framed.as_mut() else { + this.done = true; + return Poll::Ready(None); + }; + match Pin::new(framed).poll_next(cx) { + Poll::Ready(Some(Ok(DecodedItem::Data { bytes, end_stream }))) => { + this.consumed += bytes.len(); + if end_stream { + chunked.clean_eos = true; + this.done = true; + chunked.release_to_pool(); + } + Poll::Ready(Some(Ok(Frame::data(bytes)))) + } + Poll::Ready(Some(Ok(DecodedItem::End))) | Poll::Ready(None) => { + chunked.clean_eos = true; + this.done = true; + chunked.release_to_pool(); + Poll::Ready(None) + } + Poll::Ready(Some(Err(error))) => { + this.done = true; + Poll::Ready(Some(Err(this.normalize_error(error)))) + } + Poll::Pending => Poll::Pending, + } + } + } + } + + fn is_end_stream(&self) -> bool { + self.done + } + + fn size_hint(&self) -> SizeHint { + if self.done { + return SizeHint::with_exact(0); + } + match &self.kind { + BodyKind::Fast(fast) => match &fast.state { + Some(FastState::Known(k)) => SizeHint::with_exact(k.remaining as u64), + _ => SizeHint::default(), + }, + // Chunked: we don't know the total length until the terminal + // 0 chunk, so no exact hint. + BodyKind::Chunked(_) => SizeHint::default(), + BodyKind::Complete { .. } => SizeHint::with_exact(0), + } + } +} + +impl Drop for ResponseBody { + fn drop(&mut self) { + match &mut self.kind { + BodyKind::Fast(fast) => { + // this is probably dead code, we already release earlier on most paths + fast.release_to_pool(); + } + BodyKind::Chunked(chunked) => { + if chunked.clean_eos && chunked.poolable { + chunked.release_to_pool(); + } + } + BodyKind::Complete { + reader, + pool_hook, + poolable, + } => { + if *poolable && let (Some(hook), Some(reader)) = (pool_hook.take(), reader.take()) { + hook(reader); + } + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/quickwit/quickwit-http-client/src/body/decoder.rs b/quickwit/quickwit-http-client/src/body/decoder.rs new file mode 100644 index 00000000000..ebc051a47c8 --- /dev/null +++ b/quickwit/quickwit-http-client/src/body/decoder.rs @@ -0,0 +1,262 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use bytes::{Buf, Bytes, BytesMut}; +use tokio_util::codec::Decoder; + +use crate::error::HttpError; + +pub(super) const MAX_CHUNK_SIZE_LINE_SIZE: usize = 8 * 1024; +pub(super) const MAX_TRAILER_SECTION_SIZE: usize = 64 * 1024; + +#[derive(Debug)] +enum ChunkedState { + /// Reading the next chunk-size line. + Size, + /// Reading chunk data; `remaining` bytes left in the current chunk. + Data { remaining: usize }, + /// Expecting the CRLF after a chunk's data. + AfterCrlf, + /// Reading the trailer section (after the `0`-sized chunk). + Trailers, +} + +#[derive(Debug)] +enum State { + /// Actively decoding chunks. + Chunked(ChunkedState), + /// The terminal `0` chunk's trailers have been fully consumed. + Done, +} + +pub(super) enum DecodedItem { + Data { bytes: Bytes, end_stream: bool }, + End, +} + +#[derive(Debug)] +pub(super) struct HttpBodyDecoder { + state: State, + target: usize, + chunk_buf: BytesMut, + chunk_decoded: usize, + trailer_bytes: usize, +} + +impl HttpBodyDecoder { + /// Creates a decoder for a `Transfer-Encoding: chunked` body. `target` + /// is the frame coalescing size from [`super::BufferHint`]. + pub(super) fn new(target: usize) -> Self { + Self { + state: State::Chunked(ChunkedState::Size), + target, + chunk_buf: BytesMut::new(), + chunk_decoded: 0, + trailer_bytes: 0, + } + } + + pub(super) fn initial_capacity(&self) -> usize { + match self.state { + State::Chunked(_) => MAX_CHUNK_SIZE_LINE_SIZE, + State::Done => 1, + } + } + + fn take_chunk_frame(&mut self, end_stream: bool) -> DecodedItem { + let bytes = self.chunk_buf.split().freeze(); + self.chunk_decoded += bytes.len(); + DecodedItem::Data { bytes, end_stream } + } + + fn decode_chunked(&mut self, src: &mut BytesMut) -> Result, HttpError> { + loop { + match &mut self.state { + State::Chunked(ChunkedState::Size) => match httparse::parse_chunk_size(src) { + Ok(httparse::Status::Complete((line_len, chunk_size))) => { + if line_len > MAX_CHUNK_SIZE_LINE_SIZE { + return Err(chunk_size_line_too_large()); + } + let chunk_size = usize::try_from(chunk_size).map_err(|_| { + HttpError::InvalidLength("chunk size does not fit in usize".to_string()) + })?; + src.advance(line_len); + self.state = if chunk_size == 0 { + State::Chunked(ChunkedState::Trailers) + } else { + State::Chunked(ChunkedState::Data { + remaining: chunk_size, + }) + }; + } + Ok(httparse::Status::Partial) => { + if src.len() > MAX_CHUNK_SIZE_LINE_SIZE { + return Err(chunk_size_line_too_large()); + } + return Ok(None); + } + Err(error) => { + return Err(HttpError::InvalidLength(format!( + "invalid chunk size: {error}" + ))); + } + }, + State::Chunked(ChunkedState::Data { remaining }) => { + let output_space = self.target - self.chunk_buf.len(); + let take = src.len().min(*remaining).min(output_space); + // TODO we could try to avoid this copy, it's not on the main path though + self.chunk_buf.extend_from_slice(&src[..take]); + src.advance(take); + *remaining -= take; + + if self.chunk_buf.len() == self.target { + return Ok(Some(self.take_chunk_frame(false))); + } + if *remaining == 0 { + self.state = State::Chunked(ChunkedState::AfterCrlf); + continue; + } + return Ok(None); + } + State::Chunked(ChunkedState::AfterCrlf) => { + if src.len() < 2 { + return Ok(None); + } + if &src[..2] != b"\r\n" { + return Err(HttpError::InvalidLength(format!( + "expected CRLF after chunk, got {:?}", + &src[..2] + ))); + } + src.advance(2); + self.state = State::Chunked(ChunkedState::Size); + } + State::Chunked(ChunkedState::Trailers) => { + let Some(line_len) = find_crlf(src) else { + if src.len() > MAX_TRAILER_SECTION_SIZE - self.trailer_bytes { + return Err(trailer_section_too_large()); + } + return Ok(None); + }; + let encoded_line_len = line_len + 2; + if encoded_line_len > MAX_TRAILER_SECTION_SIZE - self.trailer_bytes { + return Err(trailer_section_too_large()); + } + self.trailer_bytes += encoded_line_len; + src.advance(encoded_line_len); + if line_len != 0 { + continue; + } + self.state = State::Done; + return if self.chunk_buf.is_empty() { + Ok(Some(DecodedItem::End)) + } else { + Ok(Some(self.take_chunk_frame(true))) + }; + } + State::Done => return Ok(Some(DecodedItem::End)), + } + } + } + + fn unexpected_chunked_eof(&self) -> HttpError { + let buffered = self.chunk_decoded + self.chunk_buf.len(); + let expected = match &self.state { + State::Chunked(ChunkedState::Data { remaining }) => buffered + remaining, + _ => 0, + }; + HttpError::UnexpectedEof { + read: buffered, + expected, + } + } +} + +impl Decoder for HttpBodyDecoder { + type Item = DecodedItem; + type Error = HttpError; + + fn decode(&mut self, src: &mut BytesMut) -> Result, Self::Error> { + self.decode_chunked(src) + } + + fn decode_eof(&mut self, src: &mut BytesMut) -> Result, Self::Error> { + if let Some(item) = self.decode(src)? { + return Ok(Some(item)); + } + match self.state { + State::Chunked(_) => Err(self.unexpected_chunked_eof()), + State::Done => Ok(Some(DecodedItem::End)), + } + } +} + +fn find_crlf(src: &[u8]) -> Option { + src.windows(2).position(|window| window == b"\r\n") +} + +fn chunk_size_line_too_large() -> HttpError { + HttpError::InvalidLength(format!( + "chunk-size line exceeded {MAX_CHUNK_SIZE_LINE_SIZE} bytes" + )) +} + +fn trailer_section_too_large() -> HttpError { + HttpError::InvalidLength(format!( + "trailer section exceeded {MAX_TRAILER_SECTION_SIZE} bytes" + )) +} + +#[cfg(test)] +mod tests { + use bytes::BytesMut; + use tokio_util::codec::Decoder; + + use super::*; + + #[test] + fn chunked_decodes_single_chunk() { + let mut decoder = HttpBodyDecoder::new(4); + let mut src = BytesMut::new(); + // One chunk of 4 bytes, then a terminating 0 chunk. + src.extend_from_slice(b"4\r\nwiki\r\n0\r\n\r\n"); + let item = decoder.decode(&mut src).unwrap().unwrap(); + let DecodedItem::Data { bytes, end_stream } = item else { + panic!("expected data"); + }; + assert_eq!(&*bytes, b"wiki"); + assert!(!end_stream); + let item = decoder.decode(&mut src).unwrap().unwrap(); + assert!(matches!(item, DecodedItem::End)); + } + + #[test] + fn chunked_coalesces_across_chunks() { + let mut decoder = HttpBodyDecoder::new(8 * 1024); + let mut src = BytesMut::new(); + // 16 bytes in 4-byte chunks, should coalesce into one 8K frame. + for _ in 0..4 { + src.extend_from_slice(b"4\r\nwiki\r\n"); + } + src.extend_from_slice(b"0\r\n\r\n"); + let item = decoder.decode(&mut src).unwrap().unwrap(); + let DecodedItem::Data { bytes, end_stream } = item else { + panic!("expected data"); + }; + assert_eq!(bytes.len(), 16); + assert!(end_stream); + let item = decoder.decode(&mut src).unwrap().unwrap(); + assert!(matches!(item, DecodedItem::End)); + } +} diff --git a/quickwit/quickwit-http-client/src/body/tests.rs b/quickwit/quickwit-http-client/src/body/tests.rs new file mode 100644 index 00000000000..40f25afd6da --- /dev/null +++ b/quickwit/quickwit-http-client/src/body/tests.rs @@ -0,0 +1,333 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::io::Cursor; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::task::{Context, Poll}; +use std::time::Duration; + +use http_body_util::BodyExt; +use tokio::io::{AsyncRead, AsyncWriteExt, ReadBuf}; + +use super::decoder::{MAX_CHUNK_SIZE_LINE_SIZE, MAX_TRAILER_SECTION_SIZE}; +use super::*; + +struct FragmentedReader { + inner: Cursor>, + max_read: usize, +} + +impl FragmentedReader { + fn new(bytes: impl Into>, max_read: usize) -> Self { + Self { + inner: Cursor::new(bytes.into()), + max_read, + } + } +} + +impl AsyncRead for FragmentedReader { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + read_buf: &mut ReadBuf<'_>, + ) -> Poll> { + let position = self.inner.position() as usize; + let available = &self.inner.get_ref()[position..]; + let read_len = available.len().min(self.max_read).min(read_buf.remaining()); + read_buf.put_slice(&available[..read_len]); + self.inner.set_position((position + read_len) as u64); + Poll::Ready(Ok(())) + } +} + +async fn collect_body( + body: &mut ResponseBody, +) -> Result, HttpError> { + let mut frames = Vec::new(); + while let Some(frame) = body.frame().await { + if let Ok(data) = frame?.into_data() { + frames.push(data); + } + } + Ok(frames) +} + +fn concat(frames: &[Bytes]) -> Vec { + let mut bytes = Vec::new(); + for frame in frames { + bytes.extend_from_slice(frame); + } + bytes +} + +fn body( + reader: R, + strategy: BodyStrategy, + leftover: impl Into, + target: usize, +) -> ResponseBody { + ResponseBody::new( + reader, + strategy, + leftover.into(), + BufferHint { target }, + Duration::from_secs(5), + None, + false, + ) +} + +#[tokio::test] +async fn known_length_single_frame_from_fragmented_reads() { + let expected = vec![b'x'; 20 * 1024]; + let reader = FragmentedReader::new(expected.clone(), 317); + let mut body = body( + reader, + BodyStrategy::Known(expected.len()), + Bytes::new(), + 32 * 1024, + ); + + assert_eq!(body.size_hint().exact(), Some(expected.len() as u64)); + let frames = collect_body(&mut body).await.unwrap(); + assert_eq!(frames.len(), 1); + assert_eq!(concat(&frames), expected); + assert_eq!(body.consumed(), expected.len()); + assert!(body.is_end_stream()); +} + +#[tokio::test] +async fn known_length_coalesces_to_at_least_target_sized_frames() { + let target = 8 * 1024; + let expected = vec![b'x'; 20 * 1024]; + let reader = FragmentedReader::new(expected.clone(), 509); + let mut body = body( + reader, + BodyStrategy::Known(expected.len()), + Bytes::new(), + target, + ); + + let frames = collect_body(&mut body).await.unwrap(); + assert!( + frames[..frames.len() - 1] + .iter() + .all(|frame| frame.len() >= target) + ); + assert_eq!(concat(&frames), expected); +} + +#[tokio::test] +async fn chunked_coalesces_across_wire_chunks_and_fragmented_reads() { + let mut encoded = Vec::new(); + let mut expected = Vec::new(); + for byte in 0u8..16 { + let chunk = vec![b'a' + byte; 1024]; + encoded.extend_from_slice(b"400\r\n"); + encoded.extend_from_slice(&chunk); + encoded.extend_from_slice(b"\r\n"); + expected.extend_from_slice(&chunk); + } + encoded.extend_from_slice(b"0\r\n\r\n"); + let reader = FragmentedReader::new(encoded, 113); + let mut body = body(reader, BodyStrategy::Chunked, Bytes::new(), 8 * 1024); + + let frames = collect_body(&mut body).await.unwrap(); + assert_eq!( + frames.iter().map(Bytes::len).collect::>(), + [8 * 1024, 8 * 1024] + ); + assert_eq!(concat(&frames), expected); +} + +#[tokio::test] +async fn chunk_extensions_and_trailers_can_span_reads() { + let encoded = b"b;extension-name=extension-value\r\nhello world\r\n\ + 0\r\nx-checksum: abcdefghijklmnopqrstuvwxyz\r\n\r\n"; + let reader = FragmentedReader::new(encoded, 3); + let mut body = body(reader, BodyStrategy::Chunked, Bytes::new(), 8 * 1024); + + let frames = collect_body(&mut body).await.unwrap(); + assert_eq!(frames.len(), 1); + assert_eq!(&concat(&frames), b"hello world"); +} + +#[tokio::test] +async fn chunk_size_line_is_limited() { + let mut encoded = b"1;".to_vec(); + encoded.extend(std::iter::repeat_n(b'x', MAX_CHUNK_SIZE_LINE_SIZE)); + encoded.extend_from_slice(b"\r\na\r\n0\r\n\r\n"); + let reader = FragmentedReader::new(encoded, 257); + let mut body = body(reader, BodyStrategy::Chunked, Bytes::new(), 8 * 1024); + + let error = body.frame().await.unwrap().unwrap_err(); + let HttpError::InvalidLength(message) = error else { + panic!("expected InvalidLength, got {error:?}"); + }; + assert!(message.contains("chunk-size line exceeded")); +} + +#[tokio::test] +async fn aggregate_trailer_section_is_limited() { + let mut encoded = b"0\r\n".to_vec(); + while encoded.len() <= MAX_TRAILER_SECTION_SIZE { + encoded.extend_from_slice(b"x: "); + encoded.extend(std::iter::repeat_n(b'x', 1024)); + encoded.extend_from_slice(b"\r\n"); + } + encoded.extend_from_slice(b"\r\n"); + let reader = FragmentedReader::new(encoded, 257); + let mut body = body(reader, BodyStrategy::Chunked, Bytes::new(), 8 * 1024); + + let error = body.frame().await.unwrap().unwrap_err(); + let HttpError::InvalidLength(message) = error else { + panic!("expected InvalidLength, got {error:?}"); + }; + assert!(message.contains("trailer section exceeded")); +} + +#[tokio::test] +async fn until_close_coalesces_fragmented_reads() { + let expected = vec![b'z'; 20 * 1024]; + let reader = FragmentedReader::new(expected.clone(), 251); + let mut body = body(reader, BodyStrategy::UntilClose, Bytes::new(), 8 * 1024); + + let frames = collect_body(&mut body).await.unwrap(); + assert_eq!( + frames.iter().map(Bytes::len).collect::>(), + [8 * 1024, 8 * 1024, 4 * 1024] + ); + assert_eq!(concat(&frames), expected); +} + +#[tokio::test] +async fn truncated_known_body_errors_once_and_then_ends() { + let reader = FragmentedReader::new(b"abc", 1); + let mut body = body(reader, BodyStrategy::Known(5), Bytes::new(), 8 * 1024); + + let error = body.frame().await.unwrap().unwrap_err(); + assert!(matches!( + error, + HttpError::UnexpectedEof { + read: 3, + expected: 5 + } + )); + assert!(body.frame().await.is_none()); +} + +#[tokio::test] +async fn malformed_chunked_body_errors_once_and_then_ends() { + let reader = FragmentedReader::new(b"3\r\nabcXX", 1); + let mut body = body(reader, BodyStrategy::Chunked, Bytes::new(), 8 * 1024); + + let error = body.frame().await.unwrap().unwrap_err(); + assert!(matches!(error, HttpError::InvalidLength(_))); + assert!(body.frame().await.is_none()); +} + +#[tokio::test] +async fn leftover_longer_than_content_length_is_an_error() { + let reader = Cursor::new(Vec::new()); + let mut body = body(reader, BodyStrategy::Known(5), &b"helloXY"[..], 8 * 1024); + + let error = body.frame().await.unwrap().unwrap_err(); + assert!(matches!(error, HttpError::InvalidLength(_))); + assert!(body.frame().await.is_none()); +} + +#[tokio::test] +async fn per_read_timeout_fires() { + let (client, mut server) = tokio::io::duplex(1024); + let mut body = ResponseBody::new( + client, + BodyStrategy::Known(5), + Bytes::new(), + BufferHint { target: 8 * 1024 }, + Duration::from_millis(20), + None, + false, + ); + + let error = body.frame().await.unwrap().unwrap_err(); + assert!(error.is_timeout(), "expected timeout, got {error:?}"); + assert!(body.frame().await.is_none()); + server.shutdown().await.unwrap(); +} + +#[tokio::test(start_paused = true)] +async fn read_timeout_resets_when_bytes_keep_arriving() { + let (reader, mut writer) = tokio::io::duplex(1024); + let writer_task = tokio::spawn(async move { + for byte in b"abc" { + writer.write_all(&[*byte]).await.unwrap(); + tokio::time::sleep(Duration::from_secs(4)).await; + } + }); + let mut body = ResponseBody::new( + reader, + BodyStrategy::Known(3), + Bytes::new(), + BufferHint { target: 8 * 1024 }, + Duration::from_secs(5), + None, + false, + ); + + let frames = collect_body(&mut body).await.unwrap(); + assert_eq!(&concat(&frames), b"abc"); + writer_task.await.unwrap(); +} + +#[tokio::test] +async fn empty_body_is_already_complete() { + let reader = Cursor::new(Vec::new()); + let mut body = body(reader, BodyStrategy::Empty, Bytes::new(), 8 * 1024); + + assert!(body.is_end_stream()); + assert_eq!(body.size_hint().exact(), Some(0)); + assert!(body.frame().await.is_none()); +} + +#[tokio::test] +async fn until_close_clean_eos_is_not_pooled() { + let reader = Cursor::new(b"abc".to_vec()); + let pooled = Arc::new(AtomicUsize::new(0)); + let pooled_for_hook = pooled.clone(); + let pool_hook: Option> = Some(Box::new(move |_| { + pooled_for_hook.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + })); + let mut body = ResponseBody::new( + reader, + BodyStrategy::UntilClose, + Bytes::new(), + BufferHint { target: 8 * 1024 }, + Duration::from_secs(5), + pool_hook, + true, // keep-alive held in the head, but UntilClose is still not poolable + ); + + let frames = collect_body(&mut body).await.unwrap(); + assert_eq!(&concat(&frames), b"abc"); + assert!(body.is_end_stream(), "reached a clean EOS"); + drop(body); + assert_eq!( + pooled.load(std::sync::atomic::Ordering::SeqCst), + 0, + "UntilClose EOS is a TCP close; the pool hook must not fire" + ); +} diff --git a/quickwit/quickwit-http-client/src/client.rs b/quickwit/quickwit-http-client/src/client.rs new file mode 100644 index 00000000000..a946b0f4499 --- /dev/null +++ b/quickwit/quickwit-http-client/src/client.rs @@ -0,0 +1,327 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::time::Duration; + +use http_body::Body; +use tokio_rustls::TlsConnector; + +use crate::body::{BufferHint, ResponseBody}; +use crate::connection::{ConnStream, connect}; +use crate::dns::DnsResolver; +use crate::endpoint::Endpoint; +use crate::error::HttpError; +use crate::exchange::exchange; +use crate::pool::ConnectionPool; +use crate::request::WriteState; + +/// Default connect timeout. +pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +/// Default per-write timeout for the request side. +pub const DEFAULT_WRITE_TIMEOUT: Duration = Duration::from_secs(10); + +struct HttpClientInner { + pool: ConnectionPool, + dns: Arc, + tls_connector: Option, + tls_config: Option>, + connect_timeout: Duration, + read_timeout: Duration, + write_timeout: Duration, + buffer_hint: BufferHint, +} + +/// A streaming HTTP/1.1 client. +#[derive(Clone)] +pub struct HttpClient { + inner: Arc, +} + +impl std::fmt::Debug for HttpClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("HttpClient") + .field("connect_timeout", &self.inner.connect_timeout) + .field("read_timeout", &self.inner.read_timeout) + .field("write_timeout", &self.inner.write_timeout) + .field("buffer_hint", &self.inner.buffer_hint) + .finish_non_exhaustive() + } +} + +impl HttpClient { + /// Returns a handle to the shared connection pool. + pub fn pool(&self) -> ConnectionPool { + self.inner.pool.clone() + } + + /// Returns the TLS client config, if one was configured. + pub fn tls_config(&self) -> Option> { + self.inner.tls_config.clone() + } + + /// Returns the default `BufferHint` configured on this HttpClient. + pub fn buffer_hint(&self) -> BufferHint { + self.inner.buffer_hint + } + + /// Returns the DNS resolver configured on this `HttpClient`. + pub fn dns_resolver(&self) -> Arc { + self.inner.dns.clone() + } + + /// Performs one request/response exchange and returns the streaming + /// response. + pub async fn execute( + &self, + mut request: http::Request, + ) -> Result>, HttpError> + where + B: Body + Unpin, + B::Error: Into, + { + let endpoint = Endpoint::from_uri(request.uri())?; + derive_host(&mut request); + let method = request.method().clone(); + let buffer_hint = request + .extensions() + .get::() + .copied() + .unwrap_or(self.inner.buffer_hint); + let pool_hook = Some((self.inner.pool.clone(), endpoint.clone())); + + let (conn, was_reused) = self + .inner + .pool + .acquire(&endpoint, self.connect(&endpoint)) + .await?; + + let mut write_state = WriteState::default(); + match exchange( + conn, + &mut request, + buffer_hint, + pool_hook.clone(), + &mut write_state, + self.inner.read_timeout, + self.inner.write_timeout, + ) + .await + { + Ok(response) => Ok(response), + Err(error) => { + // If on a pooled connection and we failed early enought, it might + // just mean the connection was dead: retry the query on a fresh + // connection (but only if doing so is safe) + if was_reused && retry_is_safe(&method, &write_state) { + let conn = self.connect(&endpoint).await?; + let mut retry_state = WriteState::default(); + exchange( + conn, + &mut request, + buffer_hint, + pool_hook, + &mut retry_state, + self.inner.read_timeout, + self.inner.write_timeout, + ) + .await + } else { + Err(error) + } + } + } + } + + async fn connect(&self, endpoint: &Endpoint) -> Result { + connect( + self.inner.dns.as_ref(), + endpoint, + self.inner.tls_connector.as_ref(), + self.inner.connect_timeout, + ) + .await + } +} + +/// Inserts a Host header if none is present. +fn derive_host(request: &mut http::Request) { + use http::header::HOST; + if request.headers().contains_key(HOST) { + return; + } + let Some(authority) = request.uri().authority() else { + return; + }; + // Strip any credentials (`user:pass@`), the Host header is just host[:port] + let host = authority.host(); + let host_value = match authority.port_u16() { + Some(port) => format!("{host}:{port}"), + None => host.to_string(), + }; + if let Ok(value) = host_value.parse::() { + request.headers_mut().insert(HOST, value); + } +} + +/// Whether a dead-connection retry is safe after `exchange` failed, given how +/// far `write_request` got. +/// +/// If body was touched, we cannot replay (we'd me missing part of the body). +/// If not all head is sent, we can replay (the server didn't receive a full request). +/// If all head was sent, it depends on the method +fn retry_is_safe(method: &http::Method, state: &WriteState) -> bool { + !state.body_touched && (!state.head_sent || is_idempotent(method)) +} + +/// Methods for which a full replay is safe even after the first attempt was +/// completely sent (the server may already have processed it). +fn is_idempotent(method: &http::Method) -> bool { + matches!( + *method, + http::Method::GET | http::Method::HEAD | http::Method::OPTIONS + ) +} + +/// Builder for [`HttpClient`]. +pub struct HttpClientBuilder { + dns: Arc, + tls_config: Option>, + connect_timeout: Duration, + read_timeout: Duration, + write_timeout: Duration, + max_idle_per_host: usize, + idle_timeout: Duration, + buffer_hint: BufferHint, + // When `Some`, `build` reuses this pool instead of creating a new one + shared_pool: Option, +} + +impl HttpClientBuilder { + /// Build a default [`HttpClientBuilder`] + pub fn new() -> Self { + Self { + dns: Arc::new(crate::dns::DefaultDnsResolver), + tls_config: None, + connect_timeout: DEFAULT_CONNECT_TIMEOUT, + read_timeout: crate::body::DEFAULT_READ_TIMEOUT, + write_timeout: DEFAULT_WRITE_TIMEOUT, + max_idle_per_host: crate::pool::DEFAULT_MAX_IDLE_PER_HOST, + idle_timeout: crate::pool::DEFAULT_IDLE_TIMEOUT, + buffer_hint: BufferHint::DEFAULT, + shared_pool: None, + } + } + + /// Overrides the DNS resolver. + pub fn dns_resolver(mut self, dns: Arc) -> Self { + self.dns = dns; + self + } + + /// Overrides the TLS client config used for HTTPS endpoints. + /// + /// The caller is responsible for setting ALPN to HTTP/1.1. + pub fn tls_config(mut self, config: Arc) -> Self { + self.tls_config = Some(config); + self + } + + pub fn connect_timeout(mut self, timeout: Duration) -> Self { + self.connect_timeout = timeout; + self + } + + pub fn read_timeout(mut self, timeout: Duration) -> Self { + self.read_timeout = timeout; + self + } + + pub fn write_timeout(mut self, timeout: Duration) -> Self { + self.write_timeout = timeout; + self + } + + /// Per-host idle connection cap. `0` disables pooling. + pub fn max_idle_per_host(mut self, max: usize) -> Self { + self.max_idle_per_host = max; + self + } + + /// Idle connection timeout; idle connections older than this are dropped. + pub fn idle_timeout(mut self, timeout: Duration) -> Self { + self.idle_timeout = timeout; + self + } + + /// Default [`BufferHint`] for requests without one in their extensions. + pub fn buffer_hint(mut self, hint: BufferHint) -> Self { + self.buffer_hint = hint; + self + } + + /// Reuses `pool` instead of creating a new one. + pub fn shared_pool(mut self, pool: ConnectionPool) -> Self { + self.shared_pool = Some(pool); + self + } + + /// Builds the client. + pub fn build(self) -> Result { + let tls_config = match self.tls_config { + Some(config) => config, + None => crate::tls::default_client_config()?, + }; + let tls_connector = Some(TlsConnector::from(tls_config.clone())); + let pool = self + .shared_pool + .unwrap_or_else(|| ConnectionPool::new(self.max_idle_per_host, self.idle_timeout)); + Ok(HttpClient { + inner: Arc::new(HttpClientInner { + pool, + dns: self.dns, + tls_connector, + tls_config: Some(tls_config), + connect_timeout: self.connect_timeout, + read_timeout: self.read_timeout, + write_timeout: self.write_timeout, + buffer_hint: self.buffer_hint, + }), + }) + } + + /// Returns the configured connect timeout. + pub fn configured_connect_timeout(&self) -> Duration { + self.connect_timeout + } + + /// Returns the configured read timeout. + pub fn configured_read_timeout(&self) -> Duration { + self.read_timeout + } + + /// Returns the configured write timeout. + pub fn configured_write_timeout(&self) -> Duration { + self.write_timeout + } +} + +impl Default for HttpClientBuilder { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests; diff --git a/quickwit/quickwit-http-client/src/client/tests.rs b/quickwit/quickwit-http-client/src/client/tests.rs new file mode 100644 index 00000000000..fc590c62b0a --- /dev/null +++ b/quickwit/quickwit-http-client/src/client/tests.rs @@ -0,0 +1,343 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use bytes::Bytes; +use http_body_util::{BodyExt, Empty}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; + +use crate::body::ResponseBody; +use crate::client::{HttpClient, HttpClientBuilder}; +use crate::connection::ConnStream; + +async fn spawn_server( + body: &'static [u8], + keep_alive: bool, +) -> ( + u16, + Arc, + tokio::sync::mpsc::UnboundedReceiver, +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let conn_count = Arc::new(AtomicUsize::new(0)); + let (host_tx, host_rx) = tokio::sync::mpsc::unbounded_channel::(); + let count_for_task = conn_count.clone(); + tokio::spawn(async move { + loop { + let (mut sock, _) = match listener.accept().await { + Ok(s) => s, + Err(_) => break, + }; + let count = count_for_task.clone(); + let host_tx = host_tx.clone(); + tokio::spawn(async move { + count.fetch_add(1, Ordering::SeqCst); + let mut buf: Vec = Vec::new(); + loop { + // Read until the end of the request head. + loop { + if buf.windows(4).any(|w| w == b"\r\n\r\n") { + break; + } + let mut tmp = [0u8; 1024]; + match sock.read(&mut tmp).await { + Ok(0) => return, + Ok(n) => buf.extend_from_slice(&tmp[..n]), + Err(_) => return, + } + } + let head_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4; + if let Some(host) = parse_host(&buf[..head_end]) { + let _ = host_tx.send(host); + } + buf.drain(..head_end); + let resp = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()); + if sock.write_all(resp.as_bytes()).await.is_err() { + return; + } + if sock.write_all(body).await.is_err() { + return; + } + if !keep_alive { + return; + } + } + }); + } + }); + (port, conn_count, host_rx) +} + +/// Extracts the `Host` header value from a raw request head (case-insensitive +/// name match). +fn parse_host(head: &[u8]) -> Option { + let text = std::str::from_utf8(head).ok()?; + for line in text.split("\r\n").skip(1) { + let mut parts = line.splitn(2, ':'); + let name = parts.next()?.trim(); + let value = parts.next()?.trim(); + if name.eq_ignore_ascii_case("host") { + return Some(value.to_string()); + } + } + None +} + +fn get(uri: &str) -> http::Request> { + http::Request::builder() + .method("GET") + .uri(uri) + .body(Empty::new()) + .unwrap() +} + +async fn collect(response: http::Response>) -> Vec { + let (_parts, body) = response.into_parts(); + let bytes = body.collect().await.unwrap().to_bytes(); + bytes.to_vec() +} + +async fn http_client() -> HttpClient { + HttpClientBuilder::new().build().unwrap() +} + +#[tokio::test] +async fn sequential_requests_reuse_one_connection() { + let (port, conn_count, _host_rx) = spawn_server(b"hello", true).await; + let client = http_client().await; + let uri = format!("http://127.0.0.1:{port}/a"); + + let body = collect(client.execute(get(&uri)).await.unwrap()).await; + assert_eq!(&body, b"hello"); + let body = collect(client.execute(get(&uri)).await.unwrap()).await; + assert_eq!(&body, b"hello"); + + assert_eq!( + conn_count.load(Ordering::SeqCst), + 1, + "second request should have reused the pooled connection" + ); +} + +#[tokio::test] +async fn dead_pooled_connection_is_retried_once() { + // One-shot server: it closes after the first response, so the connection + // the client pools is dead on reuse. + let (port, conn_count, _host_rx) = spawn_server(b"world", false).await; + let client = http_client().await; + let uri = format!("http://127.0.0.1:{port}/b"); + + let body = collect(client.execute(get(&uri)).await.unwrap()).await; + assert_eq!(&body, b"world"); + + let body = collect(client.execute(get(&uri)).await.unwrap()).await; + assert_eq!(&body, b"world"); + + // First request opened one connection; the second reused the (dead) + // pooled one, detected the failure, and reconnected -> two connections. + assert_eq!(conn_count.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn concurrent_requests_then_reuse() { + let (port, conn_count, _host_rx) = spawn_server(b"x", true).await; + let client = http_client().await; + let uri = format!("http://127.0.0.1:{port}/c"); + + // empty pool: we should create as much connections as needed + const N: usize = 4; + let mut first = Vec::new(); + for _ in 0..N { + first.push(client.execute(get(&uri))); + } + let first = futures::future::join_all(first).await; + for resp in first { + assert_eq!(&collect(resp.unwrap()).await, b"x"); + } + assert_eq!(conn_count.load(Ordering::SeqCst), N); + + // now we should reuse connection and not open any new one + let mut second = Vec::new(); + for _ in 0..N { + second.push(client.execute(get(&uri))); + } + let second = futures::future::join_all(second).await; + for resp in second { + assert_eq!(&collect(resp.unwrap()).await, b"x"); + } + assert_eq!( + conn_count.load(Ordering::SeqCst), + N, + "second batch should reuse pooled connections" + ); +} + +#[tokio::test] +async fn max_idle_zero_disables_reuse() { + let (port, conn_count, _host_rx) = spawn_server(b"z", true).await; + let client = HttpClientBuilder::new() + .max_idle_per_host(0) + .build() + .unwrap(); + let uri = format!("http://127.0.0.1:{port}/d"); + + let body = collect(client.execute(get(&uri)).await.unwrap()).await; + assert_eq!(&body, b"z"); + let body = collect(client.execute(get(&uri)).await.unwrap()).await; + assert_eq!(&body, b"z"); + + assert_eq!( + conn_count.load(Ordering::SeqCst), + 2, + "pooling disabled: each request opens a new connection" + ); +} + +#[tokio::test] +async fn host_header_is_derived_when_absent() { + let (port, _conn_count, mut host_rx) = spawn_server(b"h", true).await; + let client = http_client().await; + let uri = format!("http://127.0.0.1:{port}/e"); + + // No explicit Host header on the request. + let request = http::Request::builder() + .method("GET") + .uri(&uri) + .body(Empty::::new()) + .unwrap(); + let _ = collect(client.execute(request).await.unwrap()).await; + + let host = host_rx + .try_recv() + .expect("server should have received a Host"); + assert_eq!(host, format!("127.0.0.1:{port}")); +} + +#[tokio::test] +async fn explicit_host_header_is_preserved() { + let (port, _conn_count, mut host_rx) = spawn_server(b"h", true).await; + let client = http_client().await; + let uri = format!("http://127.0.0.1:{port}/f"); + + let request = http::Request::builder() + .method("GET") + .uri(&uri) + .header("host", "example.invalid") + .body(Empty::::new()) + .unwrap(); + let _ = collect(client.execute(request).await.unwrap()).await; + + let host = host_rx + .try_recv() + .expect("server should have received a Host"); + assert_eq!(host, "example.invalid"); +} + +#[tokio::test] +async fn host_header_strips_userinfo() { + let (port, _conn_count, mut host_rx) = spawn_server(b"h", true).await; + let client = http_client().await; + let uri = format!("http://user:pass@127.0.0.1:{port}/g"); + + let request = http::Request::builder() + .method("GET") + .uri(&uri) + .body(Empty::::new()) + .unwrap(); + let _ = collect(client.execute(request).await.unwrap()).await; + + let host = host_rx + .try_recv() + .expect("server should have received a Host"); + assert_eq!( + host, + format!("127.0.0.1:{port}"), + "userinfo must be stripped" + ); +} + +mod retry_safety { + use crate::client::retry_is_safe; + use crate::request::WriteState; + + fn state(head_sent: bool, body_touched: bool) -> WriteState { + WriteState { + head_sent, + body_touched, + } + } + + // --- head not fully sent: safe for any method --- + #[test] + fn head_not_sent_is_safe_for_get() { + assert!(retry_is_safe(&http::Method::GET, &state(false, false))); + } + + #[test] + fn head_not_sent_is_safe_for_post() { + assert!(retry_is_safe(&http::Method::POST, &state(false, false))); + } + + #[test] + fn head_not_sent_is_safe_for_put() { + assert!(retry_is_safe(&http::Method::PUT, &state(false, false))); + } + + // --- head fully sent, bodyless request, side-effect free --- + #[test] + fn head_sent_bodyless_get_is_safe() { + assert!(retry_is_safe(&http::Method::GET, &state(true, false))); + } + + #[test] + fn head_sent_bodyless_head_is_safe() { + assert!(retry_is_safe(&http::Method::HEAD, &state(true, false))); + } + + #[test] + fn head_sent_bodyless_options_is_safe() { + assert!(retry_is_safe(&http::Method::OPTIONS, &state(true, false))); + } + + // --- head fully sent, might have had an empty body, not side-effect free --- + #[test] + fn head_sent_bodyless_post_is_unsafe() { + assert!(!retry_is_safe(&http::Method::POST, &state(true, false))); + } + + #[test] + fn head_sent_bodyless_put_is_unsafe() { + assert!(!retry_is_safe(&http::Method::PUT, &state(true, false))); + } + + #[test] + fn head_sent_bodyless_delete_is_unsafe() { + assert!(!retry_is_safe(&http::Method::DELETE, &state(true, false))); + } + + // --- body started: never safe --- + #[test] + fn body_touched_get_is_unsafe() { + assert!(!retry_is_safe(&http::Method::GET, &state(true, true))); + } + + #[test] + fn body_touched_post_is_unsafe() { + assert!(!retry_is_safe(&http::Method::POST, &state(true, true))); + } +} diff --git a/quickwit/quickwit-http-client/src/connection.rs b/quickwit/quickwit-http-client/src/connection.rs new file mode 100644 index 00000000000..d5b3d865ea8 --- /dev/null +++ b/quickwit/quickwit-http-client/src/connection.rs @@ -0,0 +1,710 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::io; +use std::net::{IpAddr, SocketAddr}; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use socket2::TcpKeepalive; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio::net::TcpStream; +use tokio_rustls::TlsConnector; +use tokio_rustls::client::TlsStream; + +use crate::dns::DnsResolver; +use crate::endpoint::Endpoint; +use crate::error::HttpError; + +type TlsConnStream = TlsStream; + +/// Either a plain TCP connection, or a TLS connection. +#[derive(Debug)] +pub enum ConnStream { + Plain(TcpStream), + // TlsConnStream is rather big, box it to keep the enum small + Tls(Box), +} + +impl AsyncRead for ConnStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + match self.get_mut() { + ConnStream::Plain(s) => Pin::new(s).poll_read(cx, buf), + ConnStream::Tls(s) => Pin::new(&mut **s).poll_read(cx, buf), + } + } +} + +impl AsyncWrite for ConnStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + match self.get_mut() { + ConnStream::Plain(s) => Pin::new(s).poll_write(cx, buf), + ConnStream::Tls(s) => Pin::new(&mut **s).poll_write(cx, buf), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + ConnStream::Plain(s) => Pin::new(s).poll_flush(cx), + ConnStream::Tls(s) => Pin::new(&mut **s).poll_flush(cx), + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + ConnStream::Plain(s) => Pin::new(s).poll_shutdown(cx), + ConnStream::Tls(s) => Pin::new(&mut **s).poll_shutdown(cx), + } + } +} + +/// Default TCP keepalive idle time (matches `reqwest`'s default). +const KEEPALIVE_TIME: Duration = Duration::from_secs(15); +/// Happy-eyeballs fallback delay before starting the second address family +const HAPPY_EYEBALLS_DELAY: Duration = Duration::from_millis(300); + +/// Process-wide counter used to round-robin the starting index into the +/// resolved IP list across concurrent fresh connects, so they spread across +/// the fleet instead of all converging on the first reachable IP. +// i wonder if this should be an internal of the DnsResolver instead +static CONNECT_ROTATION: AtomicUsize = AtomicUsize::new(0); + +/// Open a connection to an Endpoint. +/// +/// Perform DNS resolution, TCP and optionally TLS handshake. +/// Set a few socket options for keepalive and no-delay. +pub async fn connect( + resolver: &dyn DnsResolver, + endpoint: &Endpoint, + tls_connector: Option<&TlsConnector>, + connect_timeout: Duration, +) -> Result { + let connect = async { + let ips = resolver.resolve(&endpoint.host).await?; + let tls_connector = match (endpoint.tls, tls_connector) { + (true, Some(connector)) => Some(connector), + (true, None) => { + return Err(HttpError::Tls( + "https endpoint requested but no TLS connector was provided".to_string(), + )); + } + (false, _) => None, + }; + let server_name = match tls_connector { + Some(_) => Some(endpoint.server_name()?), + None => None, + }; + connect_addresses( + &ips, + endpoint.port, + connect_timeout, + tls_connector, + server_name.as_ref(), + ) + .await + }; + + tokio::time::timeout(connect_timeout, connect) + .await + .map_err(|_| HttpError::Timeout(connect_timeout, "connect".to_string()))? +} + +/// Connects to one of the requested IPs +/// +/// Uses happy-eyeball for dual-stacked endpoints +async fn connect_addresses( + ips: &[IpAddr], + port: u16, + connect_timeout: Duration, + tls_connector: Option<&TlsConnector>, + server_name: Option<&rustls::pki_types::ServerName<'static>>, +) -> Result { + if ips.is_empty() { + return Err(HttpError::Dns { + host: String::new(), + message: "no addresses resolved".to_string(), + }); + } + let (v4, v6) = split_by_family(ips); + // clamp so many ips don't cause overly short timeout + let divisor = v4.len().max(v6.len()).clamp(1, 4) as u32; + let per_addr = connect_timeout / divisor; + + let v4 = rotate(v4); + let v6 = rotate(v6); + + if v6.is_empty() { + return connect_family(v4, port, per_addr, tls_connector, server_name).await; + } + if v4.is_empty() { + return connect_family(v6, port, per_addr, tls_connector, server_name).await; + } + + let v4_fut = connect_family(v4, port, per_addr, tls_connector, server_name); + let v6_fut = async { + tokio::time::sleep(HAPPY_EYEBALLS_DELAY).await; + connect_family(v6, port, per_addr, tls_connector, server_name).await + }; + race_first_success(v4_fut, v6_fut).await +} + +async fn connect_family( + ips: Vec, + port: u16, + per_addr: Duration, + tls_connector: Option<&TlsConnector>, + server_name: Option<&rustls::pki_types::ServerName<'static>>, +) -> Result { + let mut last_err: Option = None; + for ip in ips { + let addr = SocketAddr::new(ip, port); + match tokio::time::timeout(per_addr, TcpStream::connect(addr)).await { + Ok(Ok(tcp)) => { + set_socket_opts(&tcp); + if let Some(connector) = tls_connector { + match tokio::time::timeout( + per_addr, + connector.connect(server_name.cloned().unwrap(), tcp), + ) + .await + { + Ok(Ok(tls)) => return Ok(ConnStream::Tls(Box::new(tls))), + Ok(Err(err)) => { + last_err = Some(HttpError::Tls(format!("tls handshake failed: {err}"))); + continue; + } + Err(_) => { + last_err = + Some(HttpError::Timeout(per_addr, "tls handshake".to_string())); + continue; + } + } + } + return Ok(ConnStream::Plain(tcp)); + } + Ok(Err(err)) => last_err = Some(HttpError::Io(err)), + Err(_) => { + last_err = Some(HttpError::Timeout(per_addr, "tcp connect".to_string())); + } + } + } + Err(last_err.unwrap_or_else(|| HttpError::Dns { + host: String::new(), + message: "no addresses resolved".to_string(), + })) +} + +async fn race_first_success(a: A, b: B) -> Result +where + A: std::future::Future> + Send, + B: std::future::Future> + Send, +{ + let mut a = Box::pin(a); + let mut b = Box::pin(b); + tokio::select! { + result = &mut a => match result { + Ok(conn) => Ok(conn), + Err(_) => b.await, + }, + result = &mut b => match result { + Ok(conn) => Ok(conn), + Err(_) => a.await, + }, + } +} + +fn set_socket_opts(tcp: &TcpStream) { + let _ = tcp.set_nodelay(true); + let keepalive = TcpKeepalive::new() + .with_time(KEEPALIVE_TIME) + .with_interval(KEEPALIVE_TIME) + .with_retries(3); + if let Err(err) = socket2::SockRef::from(tcp).set_tcp_keepalive(&keepalive) { + tracing::debug!("failed to set TCP keepalive: {err}"); + } +} + +fn split_by_family(ips: &[IpAddr]) -> (Vec, Vec) { + let mut v4 = Vec::new(); + let mut v6 = Vec::new(); + for ip in ips { + match ip { + IpAddr::V4(_) => v4.push(*ip), + IpAddr::V6(_) => v6.push(*ip), + } + } + (v4, v6) +} + +/// Rotates the list so concurrent fresh connects start at different IPs, +/// spreading load across the resolved fleet. Uses a process-wide counter +/// so concurrent connects from multiple threads get distinct offsets. +fn rotate(mut ips: Vec) -> Vec { + if ips.len() <= 1 { + return ips; + } + let offset = CONNECT_ROTATION.fetch_add(1, Ordering::Relaxed) % ips.len(); + ips.rotate_left(offset); + ips +} + +#[cfg(test)] +mod tests { + use std::net::IpAddr; + use std::sync::Arc; + use std::time::Duration; + + use rustls::pki_types::{CertificateDer, PrivateKeyDer}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use tokio_rustls::{TlsAcceptor, TlsConnector}; + + use super::connect; + use crate::dns::{DefaultDnsResolver, DnsResolver, ResolveFuture}; + use crate::endpoint::Endpoint; + use crate::error::HttpError; + + // The test certificates live in the shared test-resources directory, one + // level up from this crate. The server cert's SAN includes `127.0.0.1`, + // so we connect by IP literal and verify against our own CA rather than + // the native root store. + const CA_CERT_PATH: &str = + concat!(env!("CARGO_MANIFEST_DIR"), "/../resources/tests/tls/ca.crt"); + const SERVER_CERT_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../resources/tests/tls/server.crt" + ); + const SERVER_KEY_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../resources/tests/tls/server.key" + ); + + struct HangingResolver; + impl DnsResolver for HangingResolver { + fn resolve<'a>(&'a self, _host: &'a str) -> ResolveFuture<'a> { + Box::pin(std::future::pending::, HttpError>>()) + } + } + + struct EmptyResolver; + impl DnsResolver for EmptyResolver { + fn resolve<'a>(&'a self, host: &'a str) -> ResolveFuture<'a> { + let host = host.to_string(); + Box::pin(async move { + Err(HttpError::Dns { + host, + message: "no addresses resolved".to_string(), + }) + }) + } + } + + struct FixedResolver(Vec); + impl DnsResolver for FixedResolver { + fn resolve<'a>(&'a self, _host: &'a str) -> ResolveFuture<'a> { + let ips = self.0.clone(); + Box::pin(async move { Ok(ips) }) + } + } + + fn load_certs(path: &str) -> Vec> { + let mut reader = std::io::BufReader::new(std::fs::File::open(path).unwrap()); + rustls_pemfile::certs(&mut reader) + .collect::, _>>() + .unwrap() + } + + fn load_key(path: &str) -> PrivateKeyDer<'static> { + let mut reader = std::io::BufReader::new(std::fs::File::open(path).unwrap()); + rustls_pemfile::private_key(&mut reader).unwrap().unwrap() + } + + fn server_config() -> Arc { + let certs = load_certs(SERVER_CERT_PATH); + let key = load_key(SERVER_KEY_PATH); + let config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certs, key) + .unwrap(); + Arc::new(config) + } + + fn client_config() -> Arc { + let mut roots = rustls::RootCertStore::empty(); + for cert in load_certs(CA_CERT_PATH) { + roots.add(cert).unwrap(); + } + let config = rustls::ClientConfig::builder() + .with_root_certificates(roots) + .with_no_client_auth(); + Arc::new(config) + } + + #[tokio::test] + async fn plain_tcp_echo() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut sock, _) = listener.accept().await.unwrap(); + let mut buf = [0u8; 1024]; + loop { + let n = match sock.read(&mut buf).await { + Ok(0) => break, + Ok(n) => n, + Err(_) => break, + }; + sock.write_all(&buf[..n]).await.unwrap(); + } + }); + + let endpoint = Endpoint { + tls: false, + host: "127.0.0.1".to_string(), + port: addr.port(), + }; + let mut conn = connect(&DefaultDnsResolver, &endpoint, None, Duration::from_secs(5)) + .await + .unwrap(); + + let payload = b"hello plain"; + conn.write_all(payload).await.unwrap(); + let mut got = vec![0u8; payload.len()]; + conn.read_exact(&mut got).await.unwrap(); + assert_eq!(&got, payload); + drop(conn); + server.await.unwrap(); + } + + #[tokio::test] + async fn tls_handshake_echo() { + // try to install a crypto provider, might fail if another test already ran in the same + // process, ignore the failure + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let acceptor = TlsAcceptor::from(server_config()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let server = tokio::spawn(async move { + let (sock, _) = listener.accept().await.unwrap(); + let mut tls = acceptor.accept(sock).await.unwrap(); + let mut buf = [0u8; 1024]; + loop { + let n = match tls.read(&mut buf).await { + Ok(0) => break, + Ok(n) => n, + Err(_) => break, + }; + tls.write_all(&buf[..n]).await.unwrap(); + } + }); + + let connector = TlsConnector::from(client_config()); + let endpoint = Endpoint { + tls: true, + host: "127.0.0.1".to_string(), + port: addr.port(), + }; + let mut conn = connect( + &DefaultDnsResolver, + &endpoint, + Some(&connector), + Duration::from_secs(5), + ) + .await + .unwrap(); + + let payload = b"hello tls"; + conn.write_all(payload).await.unwrap(); + let mut got = vec![0u8; payload.len()]; + conn.read_exact(&mut got).await.unwrap(); + assert_eq!(&got, payload); + drop(conn); + server.await.unwrap(); + } + + // this test need 127.0.0.2 to be a thing, which it isn't on many other OSes + #[cfg(target_os = "linux")] + #[tokio::test] + async fn tls_handshake_failure_tries_next_ip() { + // try to install a crypto provider, might fail if another test already ran in the same + // process, ignore the failure + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let acceptor = TlsAcceptor::from(server_config()); + let good_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = good_listener.local_addr().unwrap().port(); + + let bad_listener = TcpListener::bind(("127.0.0.2", port)).await.unwrap(); + + let good_acceptor = acceptor.clone(); + let good_server = tokio::spawn(async move { + let (sock, _) = good_listener.accept().await.unwrap(); + let mut tls = good_acceptor.accept(sock).await.unwrap(); + let mut buf = [0u8; 1024]; + loop { + let n = match tls.read(&mut buf).await { + Ok(0) => break, + Ok(n) => n, + Err(_) => break, + }; + tls.write_all(&buf[..n]).await.unwrap(); + } + }); + // The bad server just accepts and immediately closes, so the rustls + // handshake reads EOF and errors. + let bad_server = tokio::spawn(async move { + let (_sock, _) = bad_listener.accept().await.unwrap(); + }); + + let connector = TlsConnector::from(client_config()); + let endpoint = Endpoint { + tls: true, + // this will get resolved to both ip, but must match the SAN in our certificate + host: "127.0.0.1".to_string(), + port, + }; + let resolver = FixedResolver(vec![ + "127.0.0.2".parse().unwrap(), + "127.0.0.1".parse().unwrap(), + ]); + let mut conn = connect( + &resolver, + &endpoint, + Some(&connector), + Duration::from_secs(5), + ) + .await + .expect("should fall back to the second IP"); + + let payload = b"hello after fallback"; + conn.write_all(payload).await.unwrap(); + let mut got = vec![0u8; payload.len()]; + conn.read_exact(&mut got).await.unwrap(); + assert_eq!(&got, payload); + drop(conn); + good_server.await.unwrap(); + bad_server.await.unwrap(); + } + + #[tokio::test] + async fn connect_timeout_fires() { + let resolver = HangingResolver; + let endpoint = Endpoint { + tls: false, + host: "example.com".to_string(), + port: 80, + }; + let err = connect(&resolver, &endpoint, None, Duration::from_millis(100)) + .await + .unwrap_err(); + assert!(err.is_timeout(), "expected a timeout, got {err:?}"); + } + + #[tokio::test] + async fn no_address_resolution_fails() { + let resolver = EmptyResolver; + let endpoint = Endpoint { + tls: false, + host: "nope.example".to_string(), + port: 80, + }; + let err = connect(&resolver, &endpoint, None, Duration::from_secs(5)) + .await + .unwrap_err(); + assert!( + matches!(err, HttpError::Dns { .. }), + "expected a dns error, got {err:?}" + ); + } + + #[tokio::test] + async fn connect_refused_surfaces_io_error() { + // Bind to grab a free port, then drop the listener so the port has no + // listener and the connect is refused (ECONNREFUSED on loopback). + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + drop(listener); + + let endpoint = Endpoint { + tls: false, + host: "127.0.0.1".to_string(), + port, + }; + let err = connect(&DefaultDnsResolver, &endpoint, None, Duration::from_secs(5)) + .await + .unwrap_err(); + assert!( + err.is_io(), + "expected an io error for a refused connect, got {err:?}" + ); + } + + #[tokio::test] + async fn https_without_connector_is_an_error() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + // Keep the listener alive so the TCP connect succeeds and we reach the + // TLS-connector check rather than a connect error. + let _server = tokio::spawn(async move { + let (_sock, _) = listener.accept().await.unwrap(); + }); + let endpoint = Endpoint { + tls: true, + host: "127.0.0.1".to_string(), + port: addr.port(), + }; + let err = connect(&DefaultDnsResolver, &endpoint, None, Duration::from_secs(5)) + .await + .unwrap_err(); + assert!( + matches!(err, HttpError::Tls(_)), + "expected a tls error for a missing connector, got {err:?}" + ); + } + + async fn happy_eyeballs_helper(broken: IpAddr, working: IpAddr, working_listener: TcpListener) { + let port = working_listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + let (mut sock, _) = working_listener.accept().await.unwrap(); + let mut buf = [0u8; 1024]; + loop { + let n = match sock.read(&mut buf).await { + Ok(0) => break, + Ok(n) => n, + Err(_) => break, + }; + sock.write_all(&buf[..n]).await.unwrap(); + } + }); + + let endpoint = Endpoint { + tls: false, + host: "127.0.0.1".to_string(), + port, + }; + let resolver = FixedResolver(vec![broken, working]); + + let start = std::time::Instant::now(); + let mut conn = connect(&resolver, &endpoint, None, Duration::from_secs(2)) + .await + .expect("should connect via the working family"); + assert!( + start.elapsed() < Duration::from_secs(2), + "happy-eyeballs should not wait for the broken family" + ); + + let payload = b"hello happy eyeballs"; + conn.write_all(payload).await.unwrap(); + let mut got = vec![0u8; payload.len()]; + conn.read_exact(&mut got).await.unwrap(); + assert_eq!(&got, payload); + drop(conn); + server.await.unwrap(); + } + + #[tokio::test] + async fn happy_eyeballs_prefers_v4_when_v6_broken() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + happy_eyeballs_helper( + IpAddr::V6(std::net::Ipv6Addr::LOCALHOST), + "127.0.0.1".parse().unwrap(), + listener, + ) + .await; + } + + #[tokio::test] + async fn happy_eyeballs_prefers_v6_when_v4_broken() { + // Bind an IPv6-only listener on ::1. + let socket = + socket2::Socket::new(socket2::Domain::IPV6, socket2::Type::STREAM, None).unwrap(); + // Ensure it is not dual-stack: only IPv6 connections are accepted. + socket.set_only_v6(true).unwrap(); + socket.set_nonblocking(true).unwrap(); + socket + .bind(&socket2::SockAddr::from(std::net::SocketAddr::new( + std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST), + 0u16, + ))) + .unwrap(); + socket.listen(1024).unwrap(); + let listener = TcpListener::from_std(socket.into()).unwrap(); + + happy_eyeballs_helper( + "127.0.0.1".parse().unwrap(), + IpAddr::V6(std::net::Ipv6Addr::LOCALHOST), + listener, + ) + .await; + } + + #[tokio::test(start_paused = true)] + async fn connect_timeout_divided_across_addresses() { + // Grab 4 free ports, then drop the listeners so all 4 connects refuse. + let ports: Vec = (0..4) + .map(|_| { + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + l.local_addr().unwrap().port() + }) + .collect(); + + let ips: Vec = (0..4).map(|_| "127.0.0.1".parse().unwrap()).collect(); + let resolver = FixedResolver(ips); + let endpoint = Endpoint { + tls: false, + host: "127.0.0.1".to_string(), + port: ports[0], + }; + + tokio::select! { + result = connect(&resolver, &endpoint, None, Duration::from_secs(4)) => { + let err = result.unwrap_err(); + assert!( + err.is_io() || err.is_timeout(), + "expected an io or timeout error, got {err:?}" + ); + } + _ = tokio::time::sleep(Duration::from_secs(5)) => { + panic!("connect should win the race against sleep(5s)"); + } + } + } + + #[test] + fn rotate_changes_starting_index() { + use super::rotate; + let ips: Vec = (0..4) + .map(|i| format!("127.0.0.{i}").parse().unwrap()) + .collect(); + // Two calls should produce different rotations (the counter increments). + let r1 = rotate(ips.clone()); + let r2 = rotate(ips.clone()); + assert_ne!(r1, r2, "consecutive rotates should start at different IPs"); + // Each rotation is a valid permutation of the input. + let mut sorted = r1.clone(); + sorted.sort(); + assert_eq!(sorted, ips, "rotate preserves the set of IPs"); + } +} diff --git a/quickwit/quickwit-http-client/src/connector.rs b/quickwit/quickwit-http-client/src/connector.rs new file mode 100644 index 00000000000..67897a04fce --- /dev/null +++ b/quickwit/quickwit-http-client/src/connector.rs @@ -0,0 +1,246 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::time::Duration; + +use aws_smithy_runtime_api::client::http::{ + HttpClient as SdkHttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, + SharedHttpClient, SharedHttpConnector, +}; +use aws_smithy_runtime_api::client::result::ConnectorError; +use aws_smithy_runtime_api::http::{Request as SdkRequest, Response as SdkResponse, StatusCode}; +use aws_smithy_types::body::SdkBody; +use tokio_util::task::AbortOnDropHandle; + +use crate::body::{BufferHint, ResponseBody}; +use crate::client::{HttpClient, HttpClientBuilder}; +use crate::connection::ConnStream; +use crate::error::HttpError; + +#[derive(Clone)] +pub struct SingleBufferHttp1HttpClient { + // A template client whose pool is shared across per-call connectors. + template: HttpClient, + tls_config: Arc, + default_connect_timeout: Duration, + default_read_timeout: Duration, + default_write_timeout: Duration, +} + +impl std::fmt::Debug for SingleBufferHttp1HttpClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SingleBufferHttp1HttpClient") + .field("default_connect_timeout", &self.default_connect_timeout) + .field("default_read_timeout", &self.default_read_timeout) + .field("default_write_timeout", &self.default_write_timeout) + .finish_non_exhaustive() + } +} + +/// Builder for [`SingleBufferHttp1HttpClient`]. +pub struct SingleBufferHttp1HttpClientBuilder { + inner: HttpClientBuilder, +} + +impl SingleBufferHttp1HttpClientBuilder { + pub fn connect_timeout(mut self, d: Duration) -> Self { + self.inner = self.inner.connect_timeout(d); + self + } + pub fn read_timeout(mut self, d: Duration) -> Self { + self.inner = self.inner.read_timeout(d); + self + } + pub fn write_timeout(mut self, d: Duration) -> Self { + self.inner = self.inner.write_timeout(d); + self + } + pub fn max_idle_per_host(mut self, n: usize) -> Self { + self.inner = self.inner.max_idle_per_host(n); + self + } + pub fn idle_timeout(mut self, d: Duration) -> Self { + self.inner = self.inner.idle_timeout(d); + self + } + /// Overrides the DNS resolver. + pub fn dns_resolver(mut self, resolver: Arc) -> Self { + self.inner = self.inner.dns_resolver(resolver); + self + } + + /// Overrides the TLS client config used for HTTPS endpoints. + /// + /// The caller is responsible for setting ALPN to HTTP/1.1. + pub fn tls_config(mut self, config: Arc) -> Self { + self.inner = self.inner.tls_config(config); + self + } + /// Overrides the default [`BufferHint`] for request without one in their extensions. + pub fn buffer_hint(mut self, hint: BufferHint) -> Self { + self.inner = self.inner.buffer_hint(hint); + self + } + pub fn build(self) -> Result { + let default_connect_timeout = self.inner.configured_connect_timeout(); + let default_read_timeout = self.inner.configured_read_timeout(); + let default_write_timeout = self.inner.configured_write_timeout(); + let template = self.inner.build()?; + let tls_config = template + .tls_config() + .expect("build always sets a TLS config"); + Ok(SingleBufferHttp1HttpClient { + default_connect_timeout, + default_read_timeout, + default_write_timeout, + template, + tls_config, + }) + } +} + +impl SingleBufferHttp1HttpClient { + /// Creates a new SDK HTTP client with the OS native root store and the + /// aws-lc-rs crypto provider. + pub fn new() -> Result { + Self::builder().build() + } + + /// Returns a builder allowing the to configure the client. + pub fn builder() -> SingleBufferHttp1HttpClientBuilder { + SingleBufferHttp1HttpClientBuilder { + inner: HttpClientBuilder::new(), + } + } + + /// Builds a per-call [`HttpClient`] honoring `settings`, sharing this + /// selector's pool. + fn client_for_settings(&self, settings: &HttpConnectorSettings) -> HttpClient { + let connect_timeout = settings + .connect_timeout() + .unwrap_or(self.default_connect_timeout); + let read_timeout = settings.read_timeout().unwrap_or(self.default_read_timeout); + HttpClientBuilder::new() + .connect_timeout(connect_timeout) + .read_timeout(read_timeout) + .write_timeout(self.default_write_timeout) + .tls_config(self.tls_config.clone()) + .buffer_hint(self.template.buffer_hint()) + .dns_resolver(self.template.dns_resolver()) + .shared_pool(self.template.pool()) + .build() + .expect("tls config was provided explicitly") + } +} + +impl Default for SingleBufferHttp1HttpClient { + fn default() -> Self { + Self::new().expect("native root TLS config loads") + } +} + +impl SdkHttpClient for SingleBufferHttp1HttpClient { + fn http_connector( + &self, + settings: &HttpConnectorSettings, + _components: &aws_smithy_runtime_api::client::runtime_components::RuntimeComponents, + ) -> SharedHttpConnector { + let client = self.client_for_settings(settings); + SharedHttpConnector::new(SingleBufferHttp1Connector { client }) + } +} + +#[derive(Clone, Debug)] +pub struct SingleBufferHttp1Connector { + client: HttpClient, +} + +impl SingleBufferHttp1Connector { + pub fn new(client: HttpClient) -> Self { + Self { client } + } +} + +impl HttpConnector for SingleBufferHttp1Connector { + fn call(&self, request: SdkRequest) -> HttpConnectorFuture { + let client = self.client.clone(); + HttpConnectorFuture::new(async move { + // Run the client in its own task so the task handling it has a much + // shorter stack/state machine depth, and the sdk only get waked up + // when a buffer is received + let driver = AbortOnDropHandle::new(tokio::spawn(async move { + let http_request: http::Request = request + .try_into_http1x() + .map_err(|err| ConnectorError::other(err.into(), None))?; + let response = client + .execute(http_request) + .await + .map_err(to_connector_error)?; + convert_response(response) + })); + driver + .await + .map_err(|err| ConnectorError::other(err.into(), None))? + }) + } +} + +/// Converts the core client's streaming `http::Response` into +/// the SDK's `Response`. +fn convert_response( + response: http::Response>, +) -> Result, ConnectorError> { + let (parts, body) = response.into_parts(); + let status = StatusCode::try_from(parts.status.as_u16()) + .map_err(|err| ConnectorError::other(err.into(), None))?; + let mut sdk_response = SdkResponse::new(status, SdkBody::from_body_1_x(body)); + for (name, value) in &parts.headers { + sdk_response + .headers_mut() + .try_append( + name.as_str().to_string(), + std::str::from_utf8(value.as_bytes()) + .map_err(|err| ConnectorError::other(err.into(), None))? + .to_string(), + ) + .map_err(|err| ConnectorError::other(err.into(), None))?; + } + Ok(sdk_response) +} + +/// Maps a core [`HttpError`] to the SDK's [`ConnectorError`]. +fn to_connector_error(err: HttpError) -> ConnectorError { + let is_io = err.is_io() + || matches!( + err, + HttpError::UnexpectedEof { .. } | HttpError::InvalidLength(_) | HttpError::Dns { .. } + ); + if err.is_timeout() { + ConnectorError::timeout(err.into()) + } else if is_io { + ConnectorError::io(err.into()) + } else { + ConnectorError::other(err.into(), None) + } +} + +/// Convenience: wrap a [`SingleBufferHttp1HttpClient`] into the SDK's +/// [`SharedHttpClient`] for `s3_config.set_http_client(...)`. +pub fn shared_http_client(client: SingleBufferHttp1HttpClient) -> SharedHttpClient { + SharedHttpClient::new(client) +} + +#[cfg(test)] +mod tests; diff --git a/quickwit/quickwit-http-client/src/connector/tests.rs b/quickwit/quickwit-http-client/src/connector/tests.rs new file mode 100644 index 00000000000..a0ea2a329df --- /dev/null +++ b/quickwit/quickwit-http-client/src/connector/tests.rs @@ -0,0 +1,155 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Duration; + +use aws_smithy_runtime_api::client::http::HttpConnector; +use aws_smithy_runtime_api::http::Request as SdkRequest; +use aws_smithy_types::body::SdkBody; +use http_body_util::BodyExt; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; + +use super::SingleBufferHttp1Connector; +use crate::client::HttpClientBuilder; + +async fn read_request_head(sock: &mut tokio::net::TcpStream) -> bool { + let mut acc: Vec = Vec::new(); + let mut tmp = [0u8; 1024]; + loop { + if acc.windows(4).any(|w| w == b"\r\n\r\n") { + return true; + } + let n = match tokio::time::timeout(Duration::from_secs(2), sock.read(&mut tmp)).await { + Ok(Ok(n)) => n, + _ => return false, + }; + if n == 0 { + return false; + } + acc.extend_from_slice(&tmp[..n]); + } +} + +fn sdk_get_request(uri: &str, host: &str) -> SdkRequest { + let mut request = SdkRequest::new(SdkBody::empty()); + request.set_uri(uri).unwrap(); + request.set_method("GET").unwrap(); + request + .headers_mut() + .try_insert("host", host.to_string()) + .unwrap(); + request +} + +#[tokio::test] +async fn s3_adapter_preserves_body_content() { + // A few MB so a streaming client would necessarily receive it as many + // chunks; a recognizable pattern verifies integrity over the whole body. + let body: Vec = { + let mut buf = Vec::with_capacity(4 * 1024 * 1024); + let mut counter: u32 = 0; + while buf.len() < 4 * 1024 * 1024 { + counter = counter.wrapping_add(1); + buf.extend_from_slice(&counter.to_le_bytes()); + } + buf + }; + let response = + format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).into_bytes(); + let mut full = response; + full.extend_from_slice(&body); + let response_bytes = bytes::Bytes::from(full); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + let (mut sock, _) = listener.accept().await.unwrap(); + loop { + if !read_request_head(&mut sock).await { + return; + } + sock.write_all(&response_bytes).await.unwrap(); + sock.flush().await.unwrap(); + } + }); + + let client = HttpClientBuilder::new().build().unwrap(); + let connector = SingleBufferHttp1Connector::new(client); + let request = sdk_get_request( + &format!("http://127.0.0.1:{port}/bucket/key"), + &format!("127.0.0.1:{port}"), + ); + + let response = connector.call(request).await.unwrap(); + assert_eq!(response.status().as_u16(), 200); + + let body_bytes = response + .into_body() + .collect() + .await + .expect("body collect") + .to_bytes(); + assert_eq!(body_bytes.len(), body.len()); + assert_eq!(body_bytes, body.as_slice()); + + server.abort(); +} + +#[tokio::test] +async fn s3_adapter_maps_connect_refused_to_connector_error() { + let client = HttpClientBuilder::new() + .connect_timeout(Duration::from_secs(2)) + .build() + .unwrap(); + let connector = SingleBufferHttp1Connector::new(client); + // Grab a free port and drop the listener so the connect is refused. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + drop(listener); + + let request = sdk_get_request( + &format!("http://127.0.0.1:{port}/bucket/key"), + &format!("127.0.0.1:{port}"), + ); + let err = connector.call(request).await.unwrap_err(); + // A refused connect is an I/O-class connector error. + assert!(err.is_io(), "expected an io connector error, got {err:?}"); +} + +#[tokio::test] +async fn s3_adapter_maps_read_timeout_to_connector_error() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let _server = tokio::spawn(async move { + let (_sock, _) = listener.accept().await.unwrap(); + // Hold the connection open without responding + std::future::pending::<()>().await; + }); + + let client = HttpClientBuilder::new() + .read_timeout(Duration::from_millis(50)) + .build() + .unwrap(); + let connector = SingleBufferHttp1Connector::new(client); + let request = sdk_get_request( + &format!("http://127.0.0.1:{port}/bucket/key"), + &format!("127.0.0.1:{port}"), + ); + let err = connector.call(request).await.unwrap_err(); + assert!( + err.is_timeout(), + "expected a timeout connector error, got {err:?}" + ); +} diff --git a/quickwit/quickwit-http-client/src/dns.rs b/quickwit/quickwit-http-client/src/dns.rs new file mode 100644 index 00000000000..422226ec18a --- /dev/null +++ b/quickwit/quickwit-http-client/src/dns.rs @@ -0,0 +1,57 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::future::Future; +use std::net::IpAddr; +use std::pin::Pin; + +use crate::error::HttpError; + +/// The future returned by [`DnsResolver::resolve`]. +pub type ResolveFuture<'a> = + Pin, HttpError>> + Send + 'a>>; + +/// Resolves a hostname to a list of IP addresses. +pub trait DnsResolver: Send + Sync { + fn resolve<'a>(&'a self, host: &'a str) -> ResolveFuture<'a>; +} + +/// The default resolver, backed by `tokio::net::lookup_host`. +#[derive(Clone, Default, Debug)] +pub struct DefaultDnsResolver; + +impl DnsResolver for DefaultDnsResolver { + fn resolve<'a>(&'a self, host: &'a str) -> ResolveFuture<'a> { + Box::pin(async move { + // `lookup_host` needs a `ToSocketAddrs`, so a port is required, + // but it is only stamped onto the results; it does not change the + // resolved IPs. Use port `0` as a placeholder and discard it. + let ips: Vec = tokio::net::lookup_host((host, 0u16)) + .await + .map_err(|err| HttpError::Dns { + host: host.to_string(), + message: err.to_string(), + })? + .map(|addr| addr.ip()) + .collect(); + if ips.is_empty() { + return Err(HttpError::Dns { + host: host.to_string(), + message: "no addresses resolved".to_string(), + }); + } + Ok(ips) + }) + } +} diff --git a/quickwit/quickwit-http-client/src/endpoint.rs b/quickwit/quickwit-http-client/src/endpoint.rs new file mode 100644 index 00000000000..9397fa645f2 --- /dev/null +++ b/quickwit/quickwit-http-client/src/endpoint.rs @@ -0,0 +1,145 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use http::Uri; + +use crate::error::HttpError; + +/// A connection target: protocol (plain/TLS), host, and port. +/// +/// Used to create connections and to key the connection pool. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Endpoint { + pub tls: bool, + pub host: String, + pub port: u16, +} + +impl Endpoint { + /// Builds an `Endpoint` from a parsed request URI + pub fn from_uri(uri: &Uri) -> Result { + let scheme = uri.scheme_str().ok_or_else(|| { + HttpError::InvalidUri("missing scheme (expected http/https)".to_string()) + })?; + let tls = match scheme { + "https" => true, + "http" => false, + other => { + return Err(HttpError::InvalidUri(format!( + "unsupported scheme `{other}`" + ))); + } + }; + + let host = uri + .host() + .ok_or_else(|| HttpError::InvalidUri("missing host (authority)".to_string()))?; + + // `Uri::port_u16` returns `None` both when no port is present and when + // the port is the scheme default; fall back to the scheme default. + let port = uri.port_u16().unwrap_or(if tls { 443 } else { 80 }); + + Ok(Self { + tls, + host: host.to_string(), + port, + }) + } + + /// Parses an absolute URI (`https://host[:port]/...`) into an `Endpoint`. + /// + /// Convenience wrapper around [`Self::from_uri`] for tests and callers that build an endpoint + /// from a string. + pub fn parse(uri: &str) -> Result { + let parsed: Uri = uri + .parse() + .map_err(|err| HttpError::InvalidUri(format!("`{uri}`: {err}")))?; + Self::from_uri(&parsed) + } + + /// The `ServerName` used for TLS validation and SNI. Accepts both DNS + /// names and IP literals + pub fn server_name(&self) -> Result, HttpError> { + rustls::pki_types::ServerName::try_from(self.host.clone()) + .map_err(|err| HttpError::Tls(format!("invalid server name `{}`: {err}", self.host))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_https_default_port() { + let ep = Endpoint::parse("https://bucket.s3.us-east-1.amazonaws.com/key?x=1").unwrap(); + assert!(ep.tls); + assert_eq!(ep.host, "bucket.s3.us-east-1.amazonaws.com"); + assert_eq!(ep.port, 443); + } + + #[test] + fn parse_http_explicit_port() { + let ep = Endpoint::parse("http://localhost:4566/bucket/key").unwrap(); + assert!(!ep.tls); + assert_eq!(ep.host, "localhost"); + assert_eq!(ep.port, 4566); + } + + #[test] + fn parse_discards_path_and_query() { + let ep = Endpoint::parse("https://example.com/some/path?q=1#frag").unwrap(); + assert_eq!(ep.host, "example.com"); + assert_eq!(ep.port, 443); + } + + #[test] + fn from_uri_matches_parse() { + let uri: Uri = "https://example.com:8443/x".parse().unwrap(); + let ep = Endpoint::from_uri(&uri).unwrap(); + assert_eq!(ep, Endpoint::parse("https://example.com:8443/y").unwrap()); + } + + #[test] + fn parse_rejects_missing_scheme() { + assert!(Endpoint::parse("localhost:4566/x").is_err()); + } + + #[test] + fn parse_rejects_unsupported_scheme() { + assert!(Endpoint::parse("ftp://example.com/x").is_err()); + } + + #[test] + fn parse_rejects_missing_host() { + assert!(Endpoint::parse("https:///path").is_err()); + } + + #[test] + fn server_name_accepts_ip_literal() { + let ep = Endpoint::parse("https://127.0.0.1:443/x").unwrap(); + let name = ep.server_name().unwrap(); + assert!(matches!(name, rustls::pki_types::ServerName::IpAddress(_))); + } + + #[test] + fn endpoint_is_pool_key() { + // Same authority, different paths -> same endpoint (same pool key). + let a = Endpoint::parse("https://bucket.s3.amazonaws.com/keyA").unwrap(); + let b = Endpoint::parse("https://bucket.s3.amazonaws.com/keyB?x=1").unwrap(); + assert_eq!(a, b); + // Different port -> different endpoint. + let c = Endpoint::parse("https://bucket.s3.amazonaws.com:8443/keyA").unwrap(); + assert_ne!(a, c); + } +} diff --git a/quickwit/quickwit-http-client/src/error.rs b/quickwit/quickwit-http-client/src/error.rs new file mode 100644 index 00000000000..ad0502320a9 --- /dev/null +++ b/quickwit/quickwit-http-client/src/error.rs @@ -0,0 +1,87 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::io; + +/// Errors encountered while processing an HTTP request +#[derive(Debug, thiserror::Error)] +pub enum HttpError { + /// io error, usually retryable + #[error("io error: {0}")] + Io(#[from] io::Error), + + /// timeout, usually retryable + #[error("timeout after {0:?}: {1}")] + Timeout(std::time::Duration, String), + + /// The request URI could not be parsed or was missing the host/scheme. + #[error("invalid request URI: {0}")] + InvalidUri(String), + + /// The response head did not fit within the configured maximum head size. + /// Not retryable. + #[error("response head exceeded {0} bytes")] + HeadTooLarge(usize), + + /// The response could not be parsed as HTTP/1.1. + #[error("malformed HTTP/1.1 response: {0}")] + Parse(#[from] httparse::Error), + + /// The response announced a body length we could not interpret + /// (e.g. a negative `Content-Range`, a malformed `Content-Length`). + #[error("invalid response length: {0}")] + InvalidLength(String), + + /// The body ended before the expected number of bytes arrived. + #[error("unexpected end of response body: read {read} of {expected} bytes")] + UnexpectedEof { read: usize, expected: usize }, + + /// A body frame produced by the request's `http_body::Body` failed. + #[error("request body error: {0}")] + Body(String), + + /// DNS error before sending the request. It's safe to retry a non idempotent + /// operation after this error. + #[error("dns resolution failed for `{host}`: {message}")] + Dns { host: String, message: String }, + + /// A TLS error + #[error("tls error: {0}")] + Tls(String), +} + +impl From for HttpError { + fn from(err: std::convert::Infallible) -> Self { + match err {} + } +} + +// we need this to accept request body from s3 sdk +impl From> for HttpError { + fn from(err: Box) -> Self { + HttpError::Body(err.to_string()) + } +} + +impl HttpError { + /// Returns `true` if the error represents a timeout. + pub fn is_timeout(&self) -> bool { + matches!(self, HttpError::Timeout(..)) + } + + /// Returns `true` when the error is an I/O failure + pub fn is_io(&self) -> bool { + matches!(self, HttpError::Io(..)) + } +} diff --git a/quickwit/quickwit-http-client/src/exchange.rs b/quickwit/quickwit-http-client/src/exchange.rs new file mode 100644 index 00000000000..af685738c4e --- /dev/null +++ b/quickwit/quickwit-http-client/src/exchange.rs @@ -0,0 +1,82 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Duration; + +use http_body::Body; + +use crate::body::{BufferHint, ResponseBody}; +use crate::connection::ConnStream; +use crate::endpoint::Endpoint; +use crate::error::HttpError; +use crate::pool::ConnectionPool; +use crate::request::{WriteState, write_request}; +use crate::response::read_head; + +/// Performs one request/response exchange over `conn` and returns the +/// streaming response. +/// +/// `write_state` is updated as the request is written so the caller can decide, +/// on error, whether a retry is safe (the request must be both replayable and +/// uncommitted; see `client::retry_is_safe`). +/// +/// Timeouts are per read/write call, not total amounts of time spent on either, +/// i.e. they can be relatively low do detect dead stream without tripping on +/// long upload. If the request allows it and ends cleanly, the connection is +/// passed back to `pool_hook`. +pub(crate) async fn exchange( + mut conn: ConnStream, + request: &mut http::Request, + buffer_hint: BufferHint, + pool_hook: Option<(ConnectionPool, Endpoint)>, + write_state: &mut WriteState, + read_timeout: Duration, + write_timeout: Duration, +) -> Result>, HttpError> +where + B: Body + Unpin, + B::Error: Into, +{ + write_request(&mut conn, request, write_timeout, write_state).await?; + let head = read_head(&mut conn, request.method(), read_timeout).await?; + // Withhold the hook if the query explicitly asked to close the connection. + let request_close = request + .headers() + .get_all(http::header::CONNECTION) + .iter() + .filter_map(|v| v.to_str().ok()) + .flat_map(|v| v.split(',')) + .map(str::trim) + .any(|tok| tok.eq_ignore_ascii_case("close")); + let pool_hook = if request_close { + None + } else { + pool_hook.map(|(pool, endpoint)| { + Box::new(move |conn: ConnStream| { + pool.release(&endpoint, conn); + }) as Box + }) + }; + let body = ResponseBody::new( + conn, + head.body, + head.leftover, + buffer_hint, + read_timeout, + pool_hook, + head.keep_alive, + ); + let response = http::Response::from_parts(head.parts, body); + Ok(response) +} diff --git a/quickwit/quickwit-http-client/src/io.rs b/quickwit/quickwit-http-client/src/io.rs new file mode 100644 index 00000000000..bc4b93dfd47 --- /dev/null +++ b/quickwit/quickwit-http-client/src/io.rs @@ -0,0 +1,49 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Duration; + +use tokio::io::{AsyncWrite, AsyncWriteExt}; + +use crate::error::HttpError; + +/// Writes all of `buf` to `stream`, bounding each individual `write` (and the +/// final flush) by `timeout`. +pub async fn write_all_timeout( + stream: &mut W, + buf: &[u8], + timeout: Duration, +) -> Result<(), HttpError> +where + W: AsyncWrite + Unpin, +{ + let mut written = 0; + while written < buf.len() { + let n = match tokio::time::timeout(timeout, stream.write(&buf[written..])).await { + Ok(res) => res?, + Err(_) => return Err(HttpError::Timeout(timeout, "request write".to_string())), + }; + if n == 0 { + return Err(HttpError::Io(std::io::Error::new( + std::io::ErrorKind::WriteZero, + "wrote zero bytes", + ))); + } + written += n; + } + match tokio::time::timeout(timeout, stream.flush()).await { + Ok(res) => Ok(res?), + Err(_) => Err(HttpError::Timeout(timeout, "request write".to_string())), + } +} diff --git a/quickwit/quickwit-http-client/src/lib.rs b/quickwit/quickwit-http-client/src/lib.rs new file mode 100644 index 00000000000..afdc3d9b2c8 --- /dev/null +++ b/quickwit/quickwit-http-client/src/lib.rs @@ -0,0 +1,64 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! HTTP client for Quickwit object storage. +//! +//! The goal of this crate is to build a fast HTTP client to use to query an +//! object stoage. You may find other uses for it, but it might also not fit +//! your use case. Unless Hyper is somehow too slow for you, you should probably +//! not use this. +//! +//! Some limitations at the moment (which may or may not be improved in the future): +//! - no support for HTTP/2 +//! - no support for connection upgrade (websocket); `101 Switching Protocols` is rejected +//! - focus more on cpu usage and TTLB than TTFB +//! - no support for proxy yet +//! - `write_request` serializes any method and drains/sends body frames, but does not synthesize +//! request framing (`Content-Length` / `Transfer-Encoding: chunked`); the caller must set those +//! headers and do the framing +//! - request trailers are dropped, not serialized (no chunked transfer-encoding is synthesized on +//! requests, so there is no wire slot for them) +//! - the client/pool layer and the single-buffer optimization target GET downloads; other methods +//! are not really exercised yet + +pub mod body; +pub mod client; +pub mod connection; +pub mod dns; +pub mod endpoint; +pub mod error; +pub mod exchange; +pub mod io; +pub mod pool; +pub mod request; +pub mod response; +pub mod tls; + +#[cfg(feature = "s3")] +pub mod connector; + +pub use body::{BufferHint, DEFAULT_READ_TIMEOUT, ResponseBody}; +pub use client::{DEFAULT_CONNECT_TIMEOUT, DEFAULT_WRITE_TIMEOUT, HttpClient, HttpClientBuilder}; +pub use connection::{ConnStream, connect}; +#[cfg(feature = "s3")] +pub use connector::{ + SingleBufferHttp1Connector, SingleBufferHttp1HttpClient, SingleBufferHttp1HttpClientBuilder, + shared_http_client, +}; +pub use dns::{DefaultDnsResolver, DnsResolver, ResolveFuture}; +pub use endpoint::Endpoint; +pub use error::HttpError; +pub use pool::{ConnectionPool, DEFAULT_IDLE_TIMEOUT, DEFAULT_MAX_IDLE_PER_HOST}; +pub use response::{BodyStrategy, ResponseHead, read_head}; +pub use tls::default_client_config; diff --git a/quickwit/quickwit-http-client/src/pool.rs b/quickwit/quickwit-http-client/src/pool.rs new file mode 100644 index 00000000000..af742ed6581 --- /dev/null +++ b/quickwit/quickwit-http-client/src/pool.rs @@ -0,0 +1,734 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, VecDeque}; +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, Mutex, Weak}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use tokio::io::{AsyncRead, ReadBuf}; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; +use tokio::time::Instant; + +use crate::connection::ConnStream; +use crate::endpoint::Endpoint; +use crate::error::HttpError; + +/// Probes a connection with a single non-blocking poll_read. +/// +/// We don't await, Pending means the connection looks healty, +/// anything else means it's not: +/// - Ready(Err): there's an error +/// - Ready(Ok(0)): end of stream (the connection is half closed) +/// - Ready(Ok(n)): protocol desync +/// +/// This costs a non-blocking syscall +fn probe_healthy(conn: &mut ConnStream) -> bool { + let waker = std::task::Waker::noop(); + let mut context = Context::from_waker(waker); + let mut buf = [0u8; 1]; + let mut read_buf = ReadBuf::new(&mut buf); + match Pin::new(conn).poll_read(&mut context, &mut read_buf) { + Poll::Pending => true, + Poll::Ready(Ok(())) => false, + Poll::Ready(Err(_)) => false, + } +} + +/// Default per-host idle connection cap. +pub const DEFAULT_MAX_IDLE_PER_HOST: usize = 32; +/// Default idle timeout: an idle connection older than this is dropped on the +/// next acquire/release (and by the background reaper) rather than reused. +pub const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(90); + +/// An idle connection and the instant it entered the pool. +struct IdleConn { + conn: ConnStream, + idle_since: Instant, +} + +/// Per-host pool state +#[derive(Default)] +struct HostState { + /// Idle connections, ordered oldest (front) to newest (back). + idle: VecDeque, + /// Checkouts waiting for a connection to be returned. + waiters: VecDeque>, +} + +impl HostState { + fn purge_expired(&mut self, idle_timeout: Duration) { + let now = Instant::now(); + if let Some(newest) = self.idle.back() + && now.duration_since(newest.idle_since) >= idle_timeout + { + self.idle.clear(); + } else { + while let Some(oldest) = self.idle.front() { + if now.duration_since(oldest.idle_since) >= idle_timeout { + self.idle.pop_front(); + } else { + break; + } + } + } + while let Some(front) = self.waiters.front() { + if front.is_closed() { + self.waiters.pop_front(); + } else { + break; + } + } + } +} + +/// A per-host entry, shared via `Arc` so the top-level map only needs a brief +/// lookup lock; the deque work happens under this per-host lock. +struct HostEntry { + state: Mutex, +} + +struct PoolInner { + hosts: Mutex>>, + max_idle_per_host: usize, + idle_timeout: Duration, +} + +/// A connection pool keyed by [`Endpoint`]. +#[derive(Clone)] +pub struct ConnectionPool { + inner: Arc, +} + +impl ConnectionPool { + /// Creates a pool with the given per-host idle cap and idle timeout. + /// + /// `max_idle_per_host = 0` disables pooling entirely: [`Self::release`] + /// drops every connection and [`Self::acquire`] always connects fresh. + /// + /// Must be called within a Tokio runtime: this spawns a background reaper + /// task that periodically evicts expired idle connections. The task holds + /// a [`Weak`] handle to this pool and exits when the pool is dropped. + /// Lazy eviction on [`Self::acquire`] and [`Self::release`] keeps the + /// pool correct even if the reaper never runs. + pub fn new(max_idle_per_host: usize, idle_timeout: Duration) -> Self { + Self::new_with_reaper_handle(max_idle_per_host, idle_timeout).0 + } + + /// Like [`Self::new`], but also returns the background reaper task's + /// [`JoinHandle`]. + pub fn new_with_reaper_handle( + max_idle_per_host: usize, + idle_timeout: Duration, + ) -> (Self, JoinHandle<()>) { + let inner = Arc::new(PoolInner { + hosts: Mutex::new(HashMap::new()), + max_idle_per_host, + idle_timeout, + }); + let reaper = spawn_reaper(Arc::downgrade(&inner), idle_timeout); + (Self { inner }, reaper) + } + + /// Creates a pool with [`DEFAULT_MAX_IDLE_PER_HOST`] and + /// [`DEFAULT_IDLE_TIMEOUT`]. + pub fn with_defaults() -> Self { + Self::new(DEFAULT_MAX_IDLE_PER_HOST, DEFAULT_IDLE_TIMEOUT) + } + + fn entry(&self, endpoint: &Endpoint) -> Arc { + let mut hosts = self.inner.hosts.lock().unwrap(); + hosts + .entry(endpoint.clone()) + .or_insert_with(|| { + Arc::new(HostEntry { + state: Mutex::new(HostState::default()), + }) + }) + .clone() + } + + /// Takes an idle connection for `endpoint` out of the pool (MRU: the + /// most recently returned one first), or races a fresh `connect` against a + /// connection being returned by a concurrent [`Self::release`]. + /// + /// Returns the connection along with `was_reused`: `true` when it came + /// from the pool (either an idle entry or a waiter hand-off). Reused connections + /// might have died without it being noticed yet, one failing early should + /// cause a retry rather than a query failure. + pub async fn acquire( + &self, + endpoint: &Endpoint, + connect: F, + ) -> Result<(ConnStream, bool), HttpError> + where + F: Future> + Send, + { + let entry = self.entry(endpoint); + let mut rx = { + let mut state = entry.state.lock().unwrap(); + while let Some(mut conn) = state.idle.pop_back().map(|entry| entry.conn) { + if probe_healthy(&mut conn) { + return Ok((conn, true)); + } else { + drop(conn); + } + } + let (tx, rx) = oneshot::channel(); + state.waiters.push_back(tx); + rx + }; + let mut connect = Box::pin(connect); + tokio::select! { + conn = &mut rx => match conn { + Ok(conn) => Ok((conn, true)), + // The sender was dropped without sending. This shouldn't happen. + // Fall back to connecting. + Err(_) => { + let conn = connect.as_mut().await?; + Ok((conn, false)) + } + }, + conn = connect.as_mut() => { + // leave our oneshot sender alone, next call to release() that tries to send it a + // connection will clean it up + Ok((conn?, false)) + } + } + } + + /// Returns a connection to the pool for later reuse, handing it to the + /// oldest live waiter first if one is waiting. + pub fn release(&self, endpoint: &Endpoint, conn: ConnStream) { + if self.inner.max_idle_per_host == 0 { + return; + } + let entry = self.entry(endpoint); + let mut state = entry.state.lock().unwrap(); + // Hand the connection to the oldest live waiter, or park the connection. + let mut conn = conn; + while let Some(sender) = state.waiters.pop_front() { + match sender.send(conn) { + Ok(()) => return, + Err(returned) => conn = returned, + } + } + state.idle.push_back(IdleConn { + conn, + idle_since: Instant::now(), + }); + if state.idle.len() > self.inner.max_idle_per_host { + state.idle.pop_front(); + } + } + + #[cfg(test)] + pub(crate) fn idle_count(&self, endpoint: &Endpoint) -> usize { + let entry = self.entry(endpoint); + entry.state.lock().unwrap().idle.len() + } + + #[cfg(test)] + pub(crate) fn waiter_count(&self, endpoint: &Endpoint) -> usize { + let entry = self.entry(endpoint); + entry.state.lock().unwrap().waiters.len() + } + + #[cfg(test)] + pub(crate) fn host_count(&self) -> usize { + self.inner.hosts.lock().unwrap().len() + } +} + +/// Background task that periodically evicts expired idle connections and +/// reclaims stale waiters across all hosts, keeping the idle set honest +/// without an acquire happening. Holds a [`Weak`] handle so it exits as soon +/// as the pool is dropped. +fn spawn_reaper(weak: Weak, idle_timeout: Duration) -> tokio::task::JoinHandle<()> { + let tick = (idle_timeout / 2).max(Duration::from_secs(2)); + // enforce this is run only from a runtime so we can actually spawn the task + let handle = tokio::runtime::Handle::try_current() + .expect("ConnectionPool::new must be called within a Tokio runtime"); + handle.spawn(async move { + loop { + tokio::time::sleep(tick).await; + let Some(inner) = weak.upgrade() else { + return; + }; + // Collect handles under the top-level lock, then purge each host + // under its own lock so the map lock is held as little as possible. + let entries: Vec<(Endpoint, Arc)> = inner + .hosts + .lock() + .unwrap() + .iter() + .map(|(ep, entry)| (ep.clone(), entry.clone())) + .collect(); + for (endpoint, entry) in entries { + let mut state = entry.state.lock().unwrap(); + state.purge_expired(inner.idle_timeout); + // remove empty entries so they don't accumulate + // there are races where we end up dropping a state that's being interacted with, + // so that someone adds an idle conn or a waiter, just after we removed it from + // the map. This only means we sometime don't reuse a connection when we could, + // not ideal, not horrible + if state.idle.is_empty() && state.waiters.is_empty() { + drop(state); + let mut hosts = inner.hosts.lock().unwrap(); + if let Some(stale) = hosts.get(&endpoint) { + let stale_state = stale.state.lock().unwrap(); + if stale_state.idle.is_empty() && stale_state.waiters.is_empty() { + drop(stale_state); + hosts.remove(&endpoint); + } + } + } + } + } + }) +} + +#[cfg(test)] +mod tests { + use tokio::net::TcpListener; + + use super::*; + + async fn holding_server_port() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + tokio::spawn(async move { + // Hold every accepted connection so the client side stays open. + let mut held: Vec<_> = Vec::new(); + while let Ok((conn, _)) = listener.accept().await { + held.push(conn); + } + }); + port + } + + fn endpoint(port: u16) -> Endpoint { + Endpoint { + tls: false, + host: "127.0.0.1".to_string(), + port, + } + } + + async fn make_conn(port: u16) -> ConnStream { + let tcp = tokio::net::TcpStream::connect(("127.0.0.1", port)) + .await + .unwrap(); + let _ = tcp.set_nodelay(true); + ConnStream::Plain(tcp) + } + + // Local ephemeral port of a connection, used as a stable identity to + // check that acquire returns the same connection that was released. + fn local_port(conn: &ConnStream) -> u16 { + match conn { + ConnStream::Plain(tcp) => tcp.local_addr().unwrap().port(), + ConnStream::Tls(tls) => tls.get_ref().0.local_addr().unwrap().port(), + } + } + + fn pending_connect() -> std::future::Pending> { + std::future::pending() + } + + async fn fresh_connect(port: u16) -> Result { + Ok(make_conn(port).await) + } + + #[tokio::test] + async fn acquire_on_empty_pool_connects_fresh() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + let (_conn, was_reused) = pool.acquire(&ep, fresh_connect(port)).await.unwrap(); + assert!(!was_reused, "empty pool must connect fresh"); + } + + #[tokio::test] + async fn release_then_acquire_reuses_same_connection() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + let conn = make_conn(port).await; + let expected_port = local_port(&conn); + pool.release(&ep, conn); + let (reused, was_reused) = pool.acquire(&ep, pending_connect()).await.unwrap(); + assert!(was_reused, "should reuse the pooled connection"); + assert_eq!(local_port(&reused), expected_port); + assert_eq!(pool.idle_count(&ep), 0, "acquire drained the pool"); + } + + #[tokio::test] + async fn one_connection_serves_sequential_acquires() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + let mut conn = make_conn(port).await; + let identity = local_port(&conn); + + for _ in 0..3 { + pool.release(&ep, conn); + let (c, was_reused) = pool.acquire(&ep, pending_connect()).await.unwrap(); + assert!(was_reused, "should reuse the pooled connection"); + conn = c; + assert_eq!(local_port(&conn), identity, "reused the same connection"); + } + // After the loop the pool holds no connection (the last acquire took it). + assert_eq!(pool.idle_count(&ep), 0, "acquire drained the pool"); + } + + #[tokio::test] + async fn acquire_is_mru_most_recent_first() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + + let c1 = make_conn(port).await; + let c2 = make_conn(port).await; + let p1 = local_port(&c1); + let p2 = local_port(&c2); + + pool.release(&ep, c1); + pool.release(&ep, c2); + + let (first, _) = pool + .acquire(&ep, pending_connect()) + .await + .expect("first pooled"); + let (second, _) = pool + .acquire(&ep, pending_connect()) + .await + .expect("second pooled"); + assert_eq!(local_port(&first), p2, "most-recently-returned first"); + assert_eq!(local_port(&second), p1, "then the older one"); + assert_eq!(pool.idle_count(&ep), 0, "pool drained"); + } + + #[tokio::test(start_paused = true)] + async fn idle_cap_evicts_oldest_connection() { + let pool = ConnectionPool::new(2, Duration::from_secs(90)); + let port = holding_server_port().await; + let ep = endpoint(port); + + let c1 = make_conn(port).await; + let c2 = make_conn(port).await; + let c3 = make_conn(port).await; + let p2 = local_port(&c2); + let p3 = local_port(&c3); + + pool.release(&ep, c1); + pool.release(&ep, c2); + // Queue is at the cap (2); the third release evicts c1 (oldest). + pool.release(&ep, c3); + + let (got1, _) = pool + .acquire(&ep, pending_connect()) + .await + .expect("first pooled"); + let (got2, _) = pool + .acquire(&ep, pending_connect()) + .await + .expect("second pooled"); + // c1 (the oldest) was evicted when c3 was released. + assert_eq!(pool.idle_count(&ep), 0, "pool drained"); + assert_eq!(local_port(&got1), p3, "MRU: c3 (newest) first"); + assert_eq!(local_port(&got2), p2, "then c2"); + } + + #[tokio::test] + async fn max_idle_zero_disables_pooling() { + let pool = ConnectionPool::new(0, Duration::from_secs(90)); + let port = holding_server_port().await; + let ep = endpoint(port); + let conn = make_conn(port).await; + pool.release(&ep, conn); + let (_, was_reused) = pool.acquire(&ep, fresh_connect(port)).await.unwrap(); + assert!(!was_reused, "pooling disabled: must connect fresh"); + assert_eq!(pool.idle_count(&ep), 0, "nothing parked"); + } + + #[tokio::test] + async fn distinct_endpoints_do_not_share_connections() { + let pool = ConnectionPool::with_defaults(); + let port_a = holding_server_port().await; + let port_b = holding_server_port().await; + let ep_a = endpoint(port_a); + let ep_b = endpoint(port_b); + + let conn_a = make_conn(port_a).await; + let p_a = local_port(&conn_a); + pool.release(&ep_a, conn_a); + let (_, was_reused_b) = pool.acquire(&ep_b, fresh_connect(port_b)).await.unwrap(); + assert!(!was_reused_b, "ep_b must connect fresh"); + let (got_a, was_reused_a) = pool.acquire(&ep_a, pending_connect()).await.unwrap(); + assert!(was_reused_a, "ep_a still has its connection"); + assert_eq!(local_port(&got_a), p_a); + } + + #[tokio::test] + async fn acquire_probes_and_drops_fully_closed_connection() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + + let closing_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let closing_port = closing_listener.local_addr().unwrap().port(); + // The server accepts, and drops it immediately. + let server = tokio::spawn(async move { + let (sock, _) = closing_listener.accept().await.unwrap(); + drop(sock); + }); + let dead_conn = make_conn(closing_port).await; + // Wait for the server to accept and close its end. + server.await.unwrap(); + + pool.release(&ep, dead_conn); + let (_got, was_reused) = pool.acquire(&ep, fresh_connect(port)).await.unwrap(); + assert!( + !was_reused, + "fully-closed conn should be dropped, not reused" + ); + } + + #[tokio::test] + async fn acquire_probes_and_drops_half_closed_connection() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + + let half_close_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let half_close_port = half_close_listener.local_addr().unwrap().port(); + // The server accepts, shuts down its write side, and keeps the + // socket alive to keep the connection half-open. + let server = tokio::spawn(async move { + let (mut sock, _) = half_close_listener.accept().await.unwrap(); + use tokio::io::AsyncWriteExt; + sock.shutdown().await.unwrap(); + std::future::pending::<()>().await; + }); + let half_closed_conn = make_conn(half_close_port).await; + // Give the server time to accept and shut down its write side. + tokio::time::sleep(Duration::from_millis(50)).await; + + pool.release(&ep, half_closed_conn); + let (_got, was_reused) = pool.acquire(&ep, fresh_connect(port)).await.unwrap(); + assert!( + !was_reused, + "half-closed conn should be dropped, not reused" + ); + server.abort(); + } + + #[tokio::test(start_paused = true)] + async fn acquire_races_connect_and_hands_off_to_waiter() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + + let mut acquire_fut = std::pin::pin!(pool.acquire(&ep, pending_connect())); + + // Poll the acquire once: it must register its waiter and then park + // on the race (Pending), since neither the pending connect nor a + // hand-off has resolved. + let waker = futures::task::noop_waker(); + let mut context = std::task::Context::from_waker(&waker); + assert!( + matches!( + acquire_fut.as_mut().poll(&mut context), + std::task::Poll::Pending + ), + "acquire should park waiting for a connection" + ); + assert_eq!(pool.waiter_count(&ep), 1, "acquire should be waiting"); + + // A separate connection is returned; release must hand it to the + // waiter rather than parking it idle. + let donated = make_conn(port).await; + let donated_port = local_port(&donated); + pool.release(&ep, donated); + assert_eq!(pool.idle_count(&ep), 0, "released conn went to the waiter"); + + let (got, was_reused) = acquire_fut.await.expect("acquire ok"); + assert!(was_reused, "should reuse the donated connection"); + assert_eq!(local_port(&got), donated_port); + } + + #[tokio::test(start_paused = true)] + async fn stale_waiter_is_cleaned_up_by_release() { + let pool = ConnectionPool::with_defaults(); + let port = holding_server_port().await; + let ep = endpoint(port); + + let (_fresh, was_reused) = pool + .acquire(&ep, fresh_connect(port)) + .await + .expect("acquire ok"); + assert!(!was_reused); + // The stale sender is still in the waiters queue until release touches it. + assert_eq!(pool.waiter_count(&ep), 1, "stale sender not cleaned yet"); + + let conn = make_conn(port).await; + let identity = local_port(&conn); + pool.release(&ep, conn); + // The stale sender was popped (send failed), and the connection was + // parked idle rather than lost. + assert_eq!(pool.waiter_count(&ep), 0, "stale sender cleaned up"); + assert_eq!(pool.idle_count(&ep), 1, "connection parked idle"); + let (reused, _) = pool + .acquire(&ep, pending_connect()) + .await + .expect("parked connection"); + assert_eq!(local_port(&reused), identity); + } + + #[tokio::test(start_paused = true)] + async fn reaper_evicts_expired_idle_connections_without_an_acquire() { + let idle_timeout = Duration::from_secs(10); + let pool = ConnectionPool::new(8, idle_timeout); + let port = holding_server_port().await; + let ep = endpoint(port); + + let conn = make_conn(port).await; + pool.release(&ep, conn); + assert_eq!(pool.idle_count(&ep), 1, "parked"); + + // Advance past the idle timeout; the connection is now expired. + tokio::time::sleep(idle_timeout + Duration::from_millis(50)).await; + let mut evicted = false; + for _ in 0..20 { + tokio::task::yield_now().await; + if pool.idle_count(&ep) == 0 { + evicted = true; + break; + } + // Nudge time forward past the next tick if needed. + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!(evicted, "reaper should have evicted the expired connection"); + } + + #[tokio::test(start_paused = true)] + async fn reaper_reclaims_stale_waiters_for_unpooled_host() { + let idle_timeout = Duration::from_secs(10); + let port = holding_server_port().await; + let ep = endpoint(port); + + // Establish all connections before starting the reaper to make the test more + // deterministic. Mixing real socket I/O with paused tokio leads to skipped time. + let mut connections = Vec::new(); + for _ in 0..3 { + connections.push(make_conn(port).await); + } + + let pool = ConnectionPool::new(8, idle_timeout); + for conn in connections { + let connect = std::future::ready(Ok::(conn)); + let (_conn, was_reused) = pool.acquire(&ep, connect).await.unwrap(); + assert!(!was_reused); + } + assert_eq!( + pool.waiter_count(&ep), + 3, + "each connect-won acquire leaves a stale waiter" + ); + + tokio::time::sleep(idle_timeout + Duration::from_millis(50)).await; + let mut reclaimed = false; + for _ in 0..20 { + tokio::task::yield_now().await; + if pool.waiter_count(&ep) == 0 { + reclaimed = true; + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!(reclaimed, "reaper should have reclaimed the stale waiters"); + } + + #[tokio::test(start_paused = true)] + async fn reaper_evicts_empty_host_entries() { + // A host whose connections all expire (or whose waiters all get + // reclaimed) should have its `HostEntry` removed from the map, not + // linger forever. Without this, the map accumulates one entry per + // endpoint ever contacted. + let idle_timeout = Duration::from_secs(10); + let (pool, _reaper) = ConnectionPool::new_with_reaper_handle(8, idle_timeout); + let port = holding_server_port().await; + let ep = endpoint(port); + + let conn = make_conn(port).await; + pool.release(&ep, conn); + assert_eq!(pool.host_count(), 1, "host entry created on release"); + + // Advance past the idle timeout; the reaper purges the expired + // connection and then evicts the now-empty host entry. + tokio::time::sleep(idle_timeout + Duration::from_millis(50)).await; + let mut evicted = false; + for _ in 0..20 { + tokio::task::yield_now().await; + if pool.host_count() == 0 { + evicted = true; + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!(evicted, "reaper should have evicted the empty host entry"); + } + + #[tokio::test(start_paused = true)] + async fn reaper_keeps_fresh_connections() { + let idle_timeout = Duration::from_secs(10); + let pool = ConnectionPool::new(8, idle_timeout); + let port = holding_server_port().await; + let ep = endpoint(port); + + let conn = make_conn(port).await; + pool.release(&ep, conn); + // Advance less than the idle timeout, then let the reaper tick: the + // connection is still fresh and must be retained. + tokio::time::sleep(idle_timeout / 2 + Duration::from_millis(50)).await; + tokio::task::yield_now().await; + tokio::task::yield_now().await; + assert_eq!(pool.idle_count(&ep), 1, "fresh connection retained"); + } + + #[tokio::test(start_paused = true)] + async fn reaper_exits_when_pool_is_dropped() { + let idle_timeout = Duration::from_secs(10); + let (pool, reaper) = ConnectionPool::new_with_reaper_handle(8, idle_timeout); + let port = holding_server_port().await; + let ep = endpoint(port); + let conn = make_conn(port).await; + pool.release(&ep, conn); + assert!(!reaper.is_finished(), "reaper runs while the pool lives"); + + drop(pool); + // Advance past one reaper tick so the task wakes + let tick = idle_timeout / 2; + tokio::time::sleep(tick + Duration::from_millis(50)).await; + tokio::task::yield_now().await; + assert!( + reaper.is_finished(), + "reaper should have exited after the pool was dropped" + ); + } +} diff --git a/quickwit/quickwit-http-client/src/request.rs b/quickwit/quickwit-http-client/src/request.rs new file mode 100644 index 00000000000..3c679f72dcc --- /dev/null +++ b/quickwit/quickwit-http-client/src/request.rs @@ -0,0 +1,203 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::pin::Pin; +use std::time::Duration; + +use bytes::Buf; +use tokio::io::AsyncWrite; + +use crate::error::HttpError; +use crate::io::write_all_timeout; + +/// How far [`write_request`] got. +#[derive(Debug, Default)] +pub(crate) struct WriteState { + /// The entire head was sent, this makes body-less non-idempotent queries + /// non re-tryable. + pub(crate) head_sent: bool, + /// We read part of the request (not response) body: we definitely cannot replay + /// the query anymore + pub(crate) body_touched: bool, +} + +pub(crate) async fn write_request( + stream: &mut W, + request: &mut http::Request, + write_timeout: Duration, + state: &mut WriteState, +) -> Result<(), HttpError> +where + W: AsyncWrite + Unpin, + B: http_body::Body + Unpin, + B::Error: Into, +{ + let mut head = Vec::with_capacity(512); + let method = request.method().as_str(); + // TODO: also send proto and autority when talking to an http proxy + let path = request + .uri() + .path_and_query() + .map(|pq| pq.as_str()) + .unwrap_or("/"); + head.extend_from_slice(method.as_bytes()); + head.push(b' '); + head.extend_from_slice(path.as_bytes()); + head.extend_from_slice(b" HTTP/1.1\r\n"); + for (name, value) in request.headers() { + head.extend_from_slice(name.as_str().as_bytes()); + head.extend_from_slice(b": "); + head.extend_from_slice(value.as_bytes()); + head.extend_from_slice(b"\r\n"); + } + head.extend_from_slice(b"\r\n"); + write_all_timeout(stream, &head, write_timeout).await?; + // TODO should we not write the last CRLF, flush, mark head_sent, and then push that CRLF? + state.head_sent = true; + + loop { + let frame = std::future::poll_fn(|ctx| Pin::new(request.body_mut()).poll_frame(ctx)).await; + match frame { + None => break, + Some(Ok(mut frame)) => { + state.body_touched = true; + if let Some(data) = frame.data_mut() { + while data.remaining() > 0 { + let chunk = data.chunk(); + write_all_timeout(stream, chunk, write_timeout).await?; + let n = chunk.len(); + data.advance(n); + } + } + // Trailers are ignored: HTTP/1.1 chunked trailers are not + // emitted by this GET-only client. + } + Some(Err(err)) => return Err(err.into()), + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + use http_body_util::Empty; + + use super::*; + + async fn capture(mut request: http::Request) -> String + where + B: http_body::Body + Unpin, + B::Error: Into, + { + let mut writer = Vec::new(); + write_request( + &mut writer, + &mut request, + Duration::from_secs(5), + &mut WriteState::default(), + ) + .await + .unwrap(); + String::from_utf8(writer).unwrap() + } + + #[tokio::test] + async fn serialize_get_with_empty_body() { + let request = http::Request::builder() + .method("GET") + .uri("https://bucket.s3.amazonaws.com/key?x=1") + .header("host", "bucket.s3.amazonaws.com") + .header("accept", "*/*") + .body(Empty::::new()) + .unwrap(); + let text = capture(request).await; + assert!( + text.starts_with("GET /key?x=1 HTTP/1.1\r\n"), + "got: {text:?}" + ); + assert!(text.contains("host: bucket.s3.amazonaws.com\r\n")); + assert!(text.contains("accept: */*\r\n")); + assert!( + text.ends_with("\r\n\r\n"), + "no body, head ends with blank line" + ); + } + + #[tokio::test] + async fn serialize_defaults_missing_path_to_root() { + let request = http::Request::builder() + .method("GET") + .uri("https://example.com") + .header("host", "example.com") + .body(Empty::::new()) + .unwrap(); + let text = capture(request).await; + assert!(text.starts_with("GET / HTTP/1.1\r\n"), "got: {text:?}"); + } + + #[tokio::test] + async fn serialize_drains_body_frames() { + let body = http_body_util::combinators::BoxBody::new(http_body_util::StreamBody::new( + futures::stream::iter([ + Ok::<_, std::convert::Infallible>(http_body::Frame::data(Bytes::from_static( + b"hello ", + ))), + Ok(http_body::Frame::data(Bytes::from_static(b"world"))), + ]), + )); + let request = http::Request::builder() + .method("POST") + .uri("/ingest") + .header("host", "example.com") + .header("content-length", "11") + .body(body) + .unwrap(); + let text = capture(request).await; + assert!(text.starts_with("POST /ingest HTTP/1.1\r\n")); + assert!( + text.ends_with("\r\n\r\nhello world"), + "body not drained: {text:?}" + ); + } + + #[tokio::test] + async fn write_times_out_when_peer_stops_draining() { + use tokio::io::AsyncReadExt; + // A duplex with a tiny buffer and no reader: the write blocks once the + // buffer fills, and the per-write idle timeout (20 ms) fires. + let (mut client, mut server) = tokio::io::duplex(16); + let mut request = http::Request::builder() + .method("POST") + .uri("/big") + .header("host", "example.com") + .body(http_body_util::Full::new(bytes::Bytes::from(vec![ + b'x'; + 1024 + ]))) + .unwrap(); + let err = write_request( + &mut client, + &mut request, + Duration::from_millis(20), + &mut WriteState::default(), + ) + .await + .unwrap_err(); + assert!(err.is_timeout(), "expected a timeout, got {err:?}"); + // Drain to avoid a broken-pipe panic on drop. + let mut buf = vec![0u8; 1024]; + let _ = server.read(&mut buf).await; + } +} diff --git a/quickwit/quickwit-http-client/src/response.rs b/quickwit/quickwit-http-client/src/response.rs new file mode 100644 index 00000000000..5a287da338d --- /dev/null +++ b/quickwit/quickwit-http-client/src/response.rs @@ -0,0 +1,454 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Duration; + +use bytes::{Buf, Bytes, BytesMut}; +use tokio::io::{AsyncRead, AsyncReadExt}; + +use crate::error::HttpError; + +/// How to read the response body. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BodyStrategy { + /// A fixed number of bytes are expected (from `Content-Length` or + /// `Content-Range`). + Known(usize), + /// Size unknown (`Transfer-Encoding: chunked`) + Chunked, + /// Size unknown, read until end of connection + UntilClose, + /// No body to read + Empty, +} + +/// A parsed response head plus the strategy for reading its body. +#[derive(Debug)] +pub struct ResponseHead { + /// The response Head + pub parts: http::response::Parts, + /// Whether the connection may be returned to the pool after the body is + /// fully consumed. Derived from the `Connection` header and the HTTP + /// version default. + pub keep_alive: bool, + /// How we're going to read the body and detect its end. + pub body: BodyStrategy, + /// Body bytes that already arrived in the same read as the head. It must + /// be consumed before reading more from the stream. + /// It can't contain response to another request (we don't do pipelining) + pub leftover: Bytes, +} + +/// Maximum response head size we are willing to buffer. +const MAX_HEAD_SIZE: usize = 64 * 1024; +/// Upper bound on the number of headers we parse per response. +const MAX_HEADERS: usize = 128; + +/// Reads bytes from `stream` until a complete (non-informational) response +/// head is available, then parses it. Informational 1xx responses (100 Continue, +/// 103 Early Hints, ...) are consumed and skipped; the final response is +/// returned. `101 Switching Protocols` is rejected as unsupported. +/// +/// `request_method` is needed to detect HEAD queries, which don't have a body despite +/// their possible content-length. +pub async fn read_head( + stream: &mut R, + request_method: &http::Method, + read_timeout: Duration, +) -> Result +where + R: AsyncRead + Unpin, +{ + let mut buf = BytesMut::with_capacity(8192); + loop { + if buf.len() > MAX_HEAD_SIZE { + return Err(HttpError::HeadTooLarge(MAX_HEAD_SIZE)); + } + let mut headers = [httparse::EMPTY_HEADER; MAX_HEADERS]; + let mut resp = httparse::Response::new(&mut headers); + match resp.parse(&buf) { + Ok(httparse::Status::Complete(head_len)) => { + let status = resp.code.ok_or(HttpError::Parse(httparse::Error::Token))?; + // discard 1xx (they are informational) and followed by an actual + // status. 101 Switching Protocols is not supported. + if (100..200).contains(&status) { + if status == 101 { + return Err(HttpError::Parse(httparse::Error::Version)); + } + buf.advance(head_len); + continue; + } + let mut head = build_head(&resp, request_method)?; + let _ = buf.split_to(head_len); + head.leftover = buf.freeze(); + return Ok(head); + } + Ok(httparse::Status::Partial) => { + if buf.capacity() - buf.len() < 1024 { + buf.reserve(8192); + } + let n = match tokio::time::timeout(read_timeout, stream.read_buf(&mut buf)).await { + Ok(res) => res?, + Err(_) => { + return Err(HttpError::Timeout( + read_timeout, + "response head read".to_string(), + )); + } + }; + if n == 0 { + return Err(HttpError::UnexpectedEof { + read: buf.len(), + expected: 0, + }); + } + } + Err(err) => return Err(HttpError::Parse(err)), + } + } +} + +fn build_head( + resp: &httparse::Response<'_, '_>, + request_method: &http::Method, +) -> Result { + let status = resp.code.ok_or(HttpError::Parse(httparse::Error::Token))?; + let version = resp.version.unwrap_or(0); // 1 for HTTP/1.1, 0 for HTTP/1.0 + + let mut header_map = http::HeaderMap::with_capacity(resp.headers.len()); + let mut connection_close = false; + let mut connection_keepalive = false; + let mut transfer_encoding_chunked: Option = None; + let mut content_length: Option = None; + + for header in resp.headers.iter() { + let name = header.name; + let value = header.value; + if name.eq_ignore_ascii_case("connection") { + for tok in std::str::from_utf8(value) + .unwrap_or("") + .split(',') + .map(str::trim) + { + if tok.eq_ignore_ascii_case("close") { + connection_close = true; + } else if tok.eq_ignore_ascii_case("keep-alive") { + connection_keepalive = true; + } + } + } else if name.eq_ignore_ascii_case("transfer-encoding") { + let value_str = std::str::from_utf8(value).unwrap_or(""); + for tok in value_str.split(',').map(str::trim) { + if !tok.is_empty() { + transfer_encoding_chunked = Some(tok.eq_ignore_ascii_case("chunked")); + } + } + } else if name.eq_ignore_ascii_case("content-length") { + let Some(parsed) = parse_usize(value)? else { + continue; + }; + match content_length { + Some(existing) if existing != parsed => { + return Err(HttpError::InvalidLength(format!( + "conflicting Content-Length headers: {existing} vs {parsed}" + ))); + } + None => content_length = Some(parsed), + _ => {} + } + } + let header_name = http::HeaderName::from_bytes(name.as_bytes()) + .map_err(|_| HttpError::Parse(httparse::Error::HeaderName))?; + let header_value = http::HeaderValue::from_bytes(value) + .map_err(|_| HttpError::Parse(httparse::Error::HeaderName))?; + header_map.append(header_name, header_value); + } + + // keep-alive is default in 1.1 unless `connection: close`. + // with 1.0 keep-alive has to be explicitly mentioned to be used. + let keep_alive = if connection_close { + false + } else if version == 1 { + true + } else { + connection_keepalive + }; + + // 1xx other than 101 are handled in read_head. + // `HEAD` and `304 Not Modified` may have headers suggesting content, but they + // *never* have an actual body. + // 204, 304, and HEAD responses have no body despite any Content-Length. + let body = if status == 204 || status == 304 || request_method == http::Method::HEAD { + BodyStrategy::Empty + } else if let Some(chunked) = transfer_encoding_chunked { + // Transfer-Encoding takes precedence over Content-Length. + // HTTP/1.0 + TE is invalid + if version == 0 { + return Err(HttpError::InvalidLength( + "HTTP/1.0 response with Transfer-Encoding is invalid".to_string(), + )); + } + // The final encoding determines framing: `chunked` -> Chunked, + // anything else -> UntilClose (Content-Length is ignored). + if chunked { + BodyStrategy::Chunked + } else { + BodyStrategy::UntilClose + } + } else if let Some(len) = content_length { + BodyStrategy::Known(len) + } else { + BodyStrategy::UntilClose + }; + + let http_version = if version == 1 { + http::Version::HTTP_11 + } else { + http::Version::HTTP_10 + }; + let response = http::Response::builder() + .status( + http::StatusCode::from_u16(status) + .map_err(|_| HttpError::Parse(httparse::Error::Token))?, + ) + .version(http_version) + .body(()) + .map_err(|_| HttpError::Parse(httparse::Error::Token))?; + let (mut parts, ()) = response.into_parts(); + parts.headers = header_map; + + Ok(ResponseHead { + parts, + keep_alive, + body, + leftover: Bytes::new(), + }) +} + +fn parse_usize(value: &[u8]) -> Result, HttpError> { + let s = std::str::from_utf8(value) + .map_err(|err| HttpError::InvalidLength(format!("non-UTF-8 Content-Length: {err}")))?; + let trimmed = s.trim(); + if trimmed.is_empty() { + return Err(HttpError::InvalidLength("empty Content-Length".to_string())); + } + trimmed + .parse::() + .map(Some) + .map_err(|err| HttpError::InvalidLength(format!("invalid length `{s}`: {err}"))) +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn read_head_from(raw: &[u8]) -> ResponseHead { + read_head(&mut &raw[..], &http::Method::GET, Duration::from_secs(5)) + .await + .unwrap() + } + + async fn read_head_from_err(raw: &[u8]) -> HttpError { + read_head(&mut &raw[..], &http::Method::GET, Duration::from_secs(5)) + .await + .unwrap_err() + } + + async fn read_head_from_method(raw: &[u8], method: http::Method) -> ResponseHead { + read_head(&mut &raw[..], &method, Duration::from_secs(5)) + .await + .unwrap() + } + + #[tokio::test] + async fn known_length_from_content_length() { + let head = read_head_from(b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nhello").await; + assert_eq!(head.body, BodyStrategy::Known(5)); + assert!(head.keep_alive); + assert_eq!(head.parts.status, 200); + assert_eq!(head.parts.version, http::Version::HTTP_11); + assert_eq!(head.leftover, Bytes::from_static(b"hello")); + assert_eq!( + head.parts.headers.get("content-length").unwrap(), + "5".parse::().unwrap() + ); + } + + #[tokio::test] + async fn chunked() { + let head = read_head_from(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n").await; + assert_eq!(head.body, BodyStrategy::Chunked); + assert!(head.keep_alive); + } + + #[tokio::test] + async fn transfer_encoding_chunked_overrides_content_length() { + // transfer-encoding takes precedence over content-length. + let head = read_head_from( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nContent-Length: 5\r\n\r\n", + ) + .await; + assert_eq!(head.body, BodyStrategy::Chunked); + } + + #[tokio::test] + async fn until_close_when_no_length() { + let head = read_head_from(b"HTTP/1.0 200 OK\r\n\r\nbody").await; + assert_eq!(head.body, BodyStrategy::UntilClose); + assert!(!head.keep_alive, "HTTP/1.0 defaults to close"); + } + + #[tokio::test] + async fn http10_keep_alive_header() { + let head = read_head_from(b"HTTP/1.0 200 OK\r\nConnection: keep-alive\r\n\r\n").await; + assert!(head.keep_alive); + } + + #[tokio::test] + async fn http11_connection_close() { + let head = read_head_from(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n").await; + assert!(!head.keep_alive); + } + + #[tokio::test] + async fn empty_body_for_204() { + let head = read_head_from(b"HTTP/1.1 204 No Content\r\n\r\n").await; + assert_eq!(head.body, BodyStrategy::Empty); + assert!(head.keep_alive); + } + + #[tokio::test] + async fn empty_body_for_304() { + let head = read_head_from(b"HTTP/1.1 304 Not Modified\r\n\r\n").await; + assert_eq!(head.body, BodyStrategy::Empty); + } + + #[tokio::test] + async fn conflicting_content_length_headers_error() { + let err = read_head_from_err( + b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\nContent-Length: 6\r\n\r\n", + ) + .await; + assert!(matches!(err, HttpError::InvalidLength(_)), "{err:?}"); + } + + #[tokio::test] + async fn skips_informational_100() { + let head = read_head_from( + b"HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nhi", + ) + .await; + assert_eq!(head.parts.status, 200); + assert_eq!(head.body, BodyStrategy::Known(2)); + assert_eq!(head.leftover, Bytes::from_static(b"hi")); + } + + #[tokio::test] + async fn leftover_only_body_bytes() { + let head = read_head_from(b"HTTP/1.1 200 OK\r\nContent-Length: 3\r\n\r\nabcEXTRA").await; + // this works because we read on a large buffer from an in memory buffer, on an actual + // socket we might get less bytes + assert_eq!(head.leftover, Bytes::from_static(b"abcEXTRA")); + } + + #[tokio::test] + async fn read_head_times_out_when_no_bytes_arrive() { + // The head never arrives: the write side stays open so reads stay + // pending; the per-read idle timeout (20 ms) fires. + let (_client, mut server) = tokio::io::duplex(1024); + let err = read_head(&mut server, &http::Method::GET, Duration::from_millis(20)) + .await + .unwrap_err(); + assert!(err.is_timeout(), "expected a timeout, got {err:?}"); + } + + #[tokio::test] + async fn head_response_with_content_length_is_empty_body() { + let head = read_head_from_method( + b"HTTP/1.1 200 OK\r\nContent-Length: 12345\r\n\r\n", + http::Method::HEAD, + ) + .await; + assert_eq!(head.body, BodyStrategy::Empty); + assert!(head.keep_alive); + assert_eq!(head.parts.headers.get("content-length").unwrap(), "12345"); + } + + #[tokio::test] + async fn head_response_with_chunked_is_empty_body() { + let head = read_head_from_method( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n", + http::Method::HEAD, + ) + .await; + assert_eq!(head.body, BodyStrategy::Empty); + } + + #[tokio::test] + async fn transfer_encoding_non_chunked_is_until_close() { + // Any Transfer-Encoding takes precedence over Content-Length. A + // non-final `chunked` (or no `chunked` at all) means the body is + // delimited by close, not by Content-Length. + let head = read_head_from( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: gzip\r\nContent-Length: 5\r\n\r\n", + ) + .await; + assert_eq!(head.body, BodyStrategy::UntilClose); + } + + #[tokio::test] + async fn transfer_encoding_chunked_final_is_chunked() { + let head = + read_head_from(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: gzip, chunked\r\n\r\n").await; + assert_eq!(head.body, BodyStrategy::Chunked); + } + + #[tokio::test] + async fn transfer_encoding_chunked_not_final_is_until_close() { + let head = + read_head_from(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked, gzip\r\n\r\n").await; + assert_eq!(head.body, BodyStrategy::UntilClose); + } + + #[tokio::test] + async fn transfer_encoding_multiple_headers_combined() { + // determines framing. + let head = read_head_from( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: gzip\r\nTransfer-Encoding: chunked\r\n\r\n", + ) + .await; + assert_eq!(head.body, BodyStrategy::Chunked); + } + + #[tokio::test] + async fn http10_with_transfer_encoding_errors() { + let err = + read_head_from_err(b"HTTP/1.0 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n").await; + assert!( + matches!(err, HttpError::InvalidLength(_)), + "expected an error for HTTP/1.0 + TE, got {err:?}" + ); + } + + #[tokio::test] + async fn transfer_encoding_ignores_content_length() { + let head = read_head_from( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: gzip\r\nContent-Length: 5\r\n\r\n", + ) + .await; + assert_eq!(head.body, BodyStrategy::UntilClose); + // Content-Length is still surfaced in the headers. + assert_eq!(head.parts.headers.get("content-length").unwrap(), "5"); + } +} diff --git a/quickwit/quickwit-http-client/src/tls.rs b/quickwit/quickwit-http-client/src/tls.rs new file mode 100644 index 00000000000..ad536f59086 --- /dev/null +++ b/quickwit/quickwit-http-client/src/tls.rs @@ -0,0 +1,52 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use crate::error::HttpError; + +/// Builds a `rustls::ClientConfig` backed by the aws-lc-rs crypto provider +/// and the OS native root store. +pub fn default_client_config() -> Result, HttpError> { + let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider()); + let mut roots = rustls::RootCertStore::empty(); + let loaded = rustls_native_certs::load_native_certs(); + for cert in loaded.certs { + // `add` only fails on an unparseable DER blob but native-certs + // should have checked these already. + let _ = roots.add(cert); + } + if loaded.errors.is_empty() { + tracing::debug!( + loaded = roots.len(), + "loaded native root certificates for quickwit-http-client" + ); + } else { + tracing::warn!( + + errors = ?loaded.errors, + loaded = roots.len(), + "some native root certificates failed to load for quickwit-http-client" + ); + } + + let mut config = rustls::ClientConfig::builder_with_provider(provider) + .with_safe_default_protocol_versions() + .map_err(|err| HttpError::Tls(format!("unsupported TLS protocol versions: {err}")))? + .with_root_certificates(roots) + .with_no_client_auth(); + // Force HTTP/1.1 ALPN so the server cannot negotiate HTTP/2 + config.alpn_protocols = vec![b"http/1.1".to_vec()]; + Ok(Arc::new(config)) +} diff --git a/quickwit/quickwit-storage/Cargo.toml b/quickwit/quickwit-storage/Cargo.toml index 0e42395287a..2d50f16d081 100644 --- a/quickwit/quickwit-storage/Cargo.toml +++ b/quickwit/quickwit-storage/Cargo.toml @@ -45,6 +45,7 @@ aws-config = { workspace = true } aws-credential-types = { workspace = true } aws-sdk-s3 = { workspace = true } aws-smithy-types = { workspace = true } +aws-smithy-runtime-api = { workspace = true, features = ["client"] } azure_core = { workspace = true, optional = true } azure_identity = { workspace = true, optional = true } @@ -56,6 +57,7 @@ quickwit-common = { workspace = true } quickwit-metrics = { workspace = true } quickwit-config = { workspace = true } quickwit-proto = { workspace = true } +quickwit-http-client = { path = "../quickwit-http-client", features = ["s3"] } opendal = { workspace = true, optional = true } reqwest = { workspace = true, optional = true } diff --git a/quickwit/quickwit-storage/src/lib.rs b/quickwit/quickwit-storage/src/lib.rs index 73ef173eee2..b9fbb575c8b 100644 --- a/quickwit/quickwit-storage/src/lib.rs +++ b/quickwit/quickwit-storage/src/lib.rs @@ -74,7 +74,8 @@ pub use self::local_file_storage::{LocalFileStorage, LocalFileStorageFactory}; #[cfg(feature = "azure")] pub use self::object_storage::{AzureBlobStorage, AzureBlobStorageFactory}; pub use self::object_storage::{ - MultiPartPolicy, S3CompatibleObjectStorage, S3CompatibleObjectStorageFactory, + MultiPartPolicy, S3CompatibleObjectStorage, S3CompatibleObjectStorageFactory, create_s3_client, + create_s3_full_body_client, }; #[cfg(feature = "gcs")] pub use self::opendal_storage::GoogleCloudStorageFactory; diff --git a/quickwit/quickwit-storage/src/object_storage/mod.rs b/quickwit/quickwit-storage/src/object_storage/mod.rs index e914c107291..b5917d6031c 100644 --- a/quickwit/quickwit-storage/src/object_storage/mod.rs +++ b/quickwit/quickwit-storage/src/object_storage/mod.rs @@ -15,7 +15,9 @@ mod error; mod s3_compatible_storage; -pub use self::s3_compatible_storage::S3CompatibleObjectStorage; +pub use self::s3_compatible_storage::{ + S3CompatibleObjectStorage, create_s3_client, create_s3_full_body_client, +}; pub use self::s3_compatible_storage_resolver::S3CompatibleObjectStorageFactory; mod policy; diff --git a/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs b/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs index b26acb9f7b9..88ca3e4e966 100644 --- a/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs +++ b/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs @@ -96,6 +96,8 @@ impl AsyncRead for S3AsyncRead { /// S3-compatible object storage implementation. pub struct S3CompatibleObjectStorage { s3_client: S3Client, + // client to use for GetObject, sometime distinct from s3_client + get_s3_client: S3Client, uri: Uri, bucket: String, prefix: PathBuf, @@ -140,6 +142,7 @@ fn get_region(s3_storage_config: &S3StorageConfig) -> Option { }) } +/// Build a Hyper based `S3Client` pub async fn create_s3_client(s3_storage_config: &S3StorageConfig) -> S3Client { let aws_config = get_aws_config().await; let credentials_provider = @@ -179,6 +182,85 @@ pub async fn create_s3_client(s3_storage_config: &S3StorageConfig) -> S3Client { S3Client::from_conf(s3_config.build()) } +// Small adapter so we can use our custom resolver inside our custom http client +struct CachingDnsResolverBridge(quickwit_aws::dns::CachingDnsResolver); + +impl quickwit_http_client::DnsResolver for CachingDnsResolverBridge { + fn resolve<'a>( + &'a self, + host: &'a str, + ) -> std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, quickwit_http_client::HttpError>, + > + Send + + 'a, + >, + > { + use aws_smithy_runtime_api::client::dns::ResolveDns; + Box::pin(async move { + self.0 + .resolve_dns(host) + .await + .map_err(|err| quickwit_http_client::HttpError::Io(std::io::Error::other(err))) + }) + } +} + +/// Build a `S3Client` based on a custom HTTP client, which should be faster on +/// single-body, mostly through less wakeup and returning bodies as a single +/// ready to use buffer, instead of multiple buffers needing to be concatenated. +pub async fn create_s3_full_body_client(s3_storage_config: &S3StorageConfig) -> S3Client { + let aws_config = get_aws_config().await; + let credentials_provider = + get_credentials_provider(s3_storage_config).or(aws_config.credentials_provider()); + let region = get_region(s3_storage_config).or(aws_config.region().cloned()); + let mut s3_config = aws_sdk_s3::Config::builder() + .behavior_version(aws_behavior_version()) + .region(region); + + if let Some(identity_cache) = aws_config.identity_cache() { + s3_config.set_identity_cache(identity_cache); + } + s3_config.set_credentials_provider(credentials_provider); + s3_config.set_force_path_style(s3_storage_config.force_path_style_access()); + let connector = quickwit_http_client::SingleBufferHttp1HttpClient::builder() + .buffer_hint(quickwit_http_client::BufferHint { + target: 512 * 1024 * 1024, + }) + .dns_resolver(std::sync::Arc::new(CachingDnsResolverBridge( + quickwit_aws::dns::CachingDnsResolver::default(), + ))) + .build() + .map_err(|err| { + tracing::warn!(error = ?err, "failed to build single-buffer HTTP client"); + err + }) + .ok(); + s3_config.set_http_client( + connector + .map(quickwit_http_client::shared_http_client) + .or_else(|| aws_config.http_client()), + ); + s3_config.set_retry_config(aws_config.retry_config().cloned()); + s3_config.set_sleep_impl(aws_config.sleep_impl()); + // stalled stream protection doesn't work with our client, but it implement something similar + // by itself + s3_config.set_stalled_stream_protection(Some(StalledStreamProtectionConfig::disabled())); + s3_config.set_timeout_config(aws_config.timeout_config().cloned()); + + s3_config.set_response_checksum_validation(Some(ResponseChecksumValidation::WhenRequired)); + s3_config.set_request_checksum_calculation(Some(request_checksum_calculation( + s3_storage_config.checksum_algorithm, + ))); + + if let Some(endpoint) = s3_storage_config.endpoint() { + info!(endpoint=%endpoint, "using S3 endpoint defined in storage config or environment variable"); + s3_config.set_endpoint_url(Some(endpoint)); + } + S3Client::from_conf(s3_config.build()) +} + impl S3CompatibleObjectStorage { /// Creates an object storage given a region and an uri. pub async fn from_uri( @@ -202,8 +284,15 @@ impl S3CompatibleObjectStorage { let retry_params = RetryParams::aggressive(); let disable_multi_object_delete = s3_storage_config.disable_multi_object_delete; let disable_multipart_upload = s3_storage_config.disable_multipart_upload; + let get_s3_client = + if quickwit_common::get_bool_from_env_cached!("QW_S3_USE_FULL_BODY_CLIENT", false) { + create_s3_full_body_client(s3_storage_config).await + } else { + s3_client.clone() + }; Ok(Self { s3_client, + get_s3_client, uri: uri.clone(), bucket, prefix, @@ -222,6 +311,7 @@ impl S3CompatibleObjectStorage { pub fn with_prefix(self, prefix: PathBuf) -> Self { Self { s3_client: self.s3_client, + get_s3_client: self.get_s3_client, uri: self.uri, bucket: self.bucket, prefix, @@ -704,7 +794,7 @@ impl S3CompatibleObjectStorage { crate::metrics::OBJECT_STORAGE_GET_TOTAL.inc(); let _timer = HistogramTimer::new(&crate::metrics::OBJECT_STORAGE_GET_OBJECT_DURATION); - self.s3_client + self.get_s3_client .get_object() .bucket(self.bucket.clone()) .key(key) @@ -1323,7 +1413,8 @@ mod tests { let prefix = PathBuf::new(); let mut s3_storage = S3CompatibleObjectStorage { - s3_client, + s3_client: s3_client.clone(), + get_s3_client: s3_client, uri, bucket, prefix, @@ -1439,8 +1530,10 @@ mod tests { .http_client(client.clone()) .credentials_provider(credentials) .build(); + let s3_client = S3Client::from_conf(config); let storage = DebouncedStorage::new(S3CompatibleObjectStorage { - s3_client: S3Client::from_conf(config), + s3_client: s3_client.clone(), + get_s3_client: s3_client, uri: Uri::for_test("s3://bucket/indexes"), bucket: "bucket".to_string(), prefix: PathBuf::from("indexes"), @@ -1497,7 +1590,8 @@ mod tests { let prefix = PathBuf::new(); let s3_storage = S3CompatibleObjectStorage { - s3_client, + s3_client: s3_client.clone(), + get_s3_client: s3_client, uri, bucket, prefix, @@ -1535,7 +1629,8 @@ mod tests { let prefix = PathBuf::new(); let s3_storage = S3CompatibleObjectStorage { - s3_client, + s3_client: s3_client.clone(), + get_s3_client: s3_client, uri, bucket, prefix, @@ -1618,7 +1713,8 @@ mod tests { let prefix = PathBuf::new(); let s3_storage = S3CompatibleObjectStorage { - s3_client, + s3_client: s3_client.clone(), + get_s3_client: s3_client, uri, bucket, prefix, @@ -1710,7 +1806,8 @@ mod tests { let prefix = PathBuf::new(); let s3_storage = S3CompatibleObjectStorage { - s3_client, + s3_client: s3_client.clone(), + get_s3_client: s3_client, uri, bucket, prefix, @@ -1738,8 +1835,10 @@ mod tests { .credentials_provider(credentials) .request_checksum_calculation(request_checksum_calculation(checksum_algorithm)) .build(); + let s3_client = S3Client::from_conf(config); S3CompatibleObjectStorage { - s3_client: S3Client::from_conf(config), + s3_client: s3_client.clone(), + get_s3_client: s3_client, uri: Uri::for_test("s3://bucket/"), bucket: "bucket".to_string(), prefix: PathBuf::new(),