diff --git a/Cargo.lock b/Cargo.lock index e5d56f3f..5de31478 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6159,6 +6159,7 @@ dependencies = [ "object_store 0.14.2", "opentelemetry 0.32.0", "opentelemetry-semantic-conventions", + "opentelemetry_sdk", "postcard", "pretty_assertions", "protobuf", diff --git a/nisshi-broker/tests/it/storage_list_offsets.rs b/nisshi-broker/tests/it/storage_list_offsets.rs index 731304a7..1d4d14e3 100644 --- a/nisshi-broker/tests/it/storage_list_offsets.rs +++ b/nisshi-broker/tests/it/storage_list_offsets.rs @@ -24,7 +24,9 @@ use nisshi_sans_io::{ create_topics_request::CreatableTopic, list_offsets_request::{ListOffsetsPartition, ListOffsetsTopic}, }; -use nisshi_storage::{ArcDynStorage, CreateTopicsService, ListOffsetsService, Storage}; +use nisshi_storage::{ + ArcDynStorage, CreateTopicsService, LatencyIntroducingStorage, ListOffsetsService, Storage, +}; use rama::{Service, extensions::Extensions}; use rand::{prelude::*, rng}; use uuid::Uuid; @@ -196,6 +198,97 @@ async fn timestamp_no_match(storage: impl Storage + Clone, broker_id: i32) -> Re Ok(()) } +/// Every partition of a request for many partitions answers within the +/// ListOffsets deadline, in both isolation levels, when each storage call +/// takes 50 to 150ms. +async fn many_partitions_with_latency(storage: impl Storage + Clone, broker_id: i32) -> Result<()> { + const PARTITIONS: i32 = 64; + + let extensions = Extensions::default(); + let topic = &alphanumeric_string(15)[..]; + + let response = CreateTopicsService { + storage: storage.clone(), + } + .serve(RequestInput { + request: CreateTopicsRequest::default() + .validate_only(Some(false)) + .topics(Some( + [CreatableTopic::default() + .name(topic.into()) + .num_partitions(PARTITIONS) + .replication_factor(1) + .assignments(Some([].into())) + .configs(Some([].into()))] + .into(), + )), + extensions: extensions.clone(), + }) + .await?; + + let topics = response.topics.as_deref().unwrap_or_default(); + assert_eq!(1, topics.len()); + assert_eq!(ErrorCode::None, ErrorCode::try_from(topics[0].error_code)?); + + let service = ListOffsetsService { + storage: LatencyIntroducingStorage::new(storage), + }; + + for isolation_level in [ + IsolationLevel::ReadUncommitted, + IsolationLevel::ReadCommitted, + ] { + let partitions = (0..PARTITIONS) + .map(|partition_index| { + ListOffset::Latest.try_into().map(|timestamp| { + ListOffsetsPartition::default() + .current_leader_epoch(Some(-1)) + .partition_index(partition_index) + .timestamp(timestamp) + }) + }) + .collect::, _>>()?; + + let response = service + .serve(RequestInput { + request: ListOffsetsRequest::default() + .isolation_level(Some(isolation_level.into())) + .replica_id(broker_id) + .topics(Some( + [ListOffsetsTopic::default() + .name(topic.into()) + .partitions(Some(partitions))] + .into(), + )), + extensions: extensions.clone(), + }) + .await?; + + let topics = response.topics.as_deref().unwrap_or_default(); + assert_eq!(1, topics.len()); + + let partitions = topics[0].partitions.as_deref().unwrap_or_default(); + assert_eq!( + (0..PARTITIONS) + .map(|partition| (partition, ErrorCode::None, Some(0))) + .collect::>(), + partitions + .iter() + .map( + |partition| ErrorCode::try_from(partition.error_code).map(|error_code| ( + partition.partition_index, + error_code, + partition.offset + )) + ) + .collect::, _>>()?, + "{isolation_level:?}" + ); + } + + Ok(()) +} + #[cfg(feature = "dynostore")] mod in_memory { use super::*; @@ -234,6 +327,20 @@ mod in_memory { Ok(()) } + + #[tokio::test] + async fn many_partitions_with_latency() -> Result<()> { + let _guard = init_tracing()?; + + let cluster_id = Uuid::now_v7(); + let broker_id = rng().random_range(0..i32::MAX); + + let storage = storage_container(cluster_id, broker_id).await?; + + super::many_partitions_with_latency(storage, broker_id).await?; + + Ok(()) + } } #[cfg(feature = "libsql")] @@ -274,6 +381,20 @@ mod lite { Ok(()) } + + #[tokio::test] + async fn many_partitions_with_latency() -> Result<()> { + let _guard = init_tracing()?; + + let cluster_id = Uuid::now_v7(); + let broker_id = rng().random_range(0..i32::MAX); + + let storage = storage_container(cluster_id, broker_id).await?; + + super::many_partitions_with_latency(storage, broker_id).await?; + + Ok(()) + } } #[cfg(feature = "slatedb")] @@ -314,6 +435,20 @@ mod slatedb { Ok(()) } + + #[tokio::test] + async fn many_partitions_with_latency() -> Result<()> { + let _guard = init_tracing()?; + + let cluster_id = Uuid::now_v7(); + let broker_id = rng().random_range(0..i32::MAX); + + let storage = storage_container(cluster_id, broker_id).await?; + + super::many_partitions_with_latency(storage, broker_id).await?; + + Ok(()) + } } #[cfg(feature = "postgres")] @@ -354,4 +489,18 @@ mod pg { Ok(()) } + + #[tokio::test] + async fn many_partitions_with_latency() -> Result<()> { + let _guard = init_tracing()?; + + let cluster_id = Uuid::now_v7(); + let broker_id = rng().random_range(0..i32::MAX); + + let storage = storage_container(cluster_id, broker_id).await?; + + super::many_partitions_with_latency(storage, broker_id).await?; + + Ok(()) + } } diff --git a/nisshi-storage-sql/src/pg.rs b/nisshi-storage-sql/src/pg.rs index 6a28c717..edf7e6ba 100644 --- a/nisshi-storage-sql/src/pg.rs +++ b/nisshi-storage-sql/src/pg.rs @@ -19,14 +19,18 @@ use std::{ collections::BTreeMap, fmt::Debug, hash::Hash, + io, marker::PhantomData, + pin::Pin, str::FromStr, sync::{Arc, LazyLock}, + task::{Context, Poll, ready}, time::{Duration, SystemTime}, }; use async_trait::async_trait; use bytes::Bytes; +use deadpool::managed; use deadpool_postgres::{ Manager, ManagerConfig, Object, Pool, PoolError, RecyclingMethod, Runtime, TimeoutType, Transaction, @@ -65,19 +69,30 @@ use nisshi_storage::{ BrokerRegistrationRequest, Error, GroupDetail, ListOffsetResponse, METER, MetadataResponse, NamedGroupDetail, OffsetCommitRequest, OffsetStage, ProducerIdResponse, Result, ScramCredential, Storage, TopicId, Topition, TxnAddPartitionsRequest, TxnAddPartitionsResponse, - TxnOffsetCommitRequest, TxnState, UpdateError, Version, check_claim, producer_claim, + TxnOffsetCommitRequest, TxnState, UpdateError, Version, check_claim, deadline, producer_claim, }; use opentelemetry::metrics::Histogram; use opentelemetry::{KeyValue, metrics::Counter}; use rand::{prelude::*, rng}; use serde_json::Value; +use tokio::{ + io::{AsyncRead, AsyncWrite, ReadBuf}, + net::{TcpStream, UnixStream}, + runtime::Handle, + sync::{OwnedSemaphorePermit, Semaphore}, + time::{Instant, timeout, timeout_at}, +}; use tokio_postgres::{ - Config, Row, RowStream, + CancelToken, Config, Row, RowStream, binary_copy::BinaryCopyInWriter, + config::Host, error::SqlState, + tls::MakeTlsConnect, types::{BorrowToSql, ToSql, Type}, }; -use tracing::{debug, error, instrument}; +use tokio_postgres_rustls::MakeRustlsConnect; +use tokio_util::either::Either; +use tracing::{Instrument as _, debug, error, instrument, warn}; use url::Url; use uuid::Uuid; @@ -92,6 +107,13 @@ pub struct Postgres { node: i32, advertised_listener: Url, pool: Pool, + cancel: Option, + /// Permits for the ListOffsets partition reads in flight, see + /// [`LIST_OFFSETS_READS`]. + list_offsets_reads: Arc, + /// How long an abandoned statement's connection stays out of the pool + /// before it is closed, see [`cancel_and_return`]. + abandoned_statement_bound: Duration, schemas: Option, lake: Option, } @@ -103,6 +125,8 @@ pub struct Builder { node: N, advertised_listener: L, pool: P, + cancel: Option, + abandoned_statement_bound: Duration, schemas: Option, lake: Option, } @@ -114,6 +138,8 @@ impl Builder { node: self.node, advertised_listener: self.advertised_listener, pool: self.pool, + cancel: self.cancel, + abandoned_statement_bound: self.abandoned_statement_bound, schemas: self.schemas, lake: self.lake, } @@ -125,6 +151,8 @@ impl Builder { node, advertised_listener: self.advertised_listener, pool: self.pool, + cancel: self.cancel, + abandoned_statement_bound: self.abandoned_statement_bound, schemas: self.schemas, lake: self.lake, } @@ -136,6 +164,8 @@ impl Builder { node: self.node, advertised_listener, pool: self.pool, + cancel: self.cancel, + abandoned_statement_bound: self.abandoned_statement_bound, schemas: self.schemas, lake: self.lake, } @@ -157,12 +187,38 @@ impl Builder { node: self.node, advertised_listener: self.advertised_listener, pool: self.pool, + cancel: self.cancel, + list_offsets_reads: Arc::new(Semaphore::new(LIST_OFFSETS_READS)), + abandoned_statement_bound: self.abandoned_statement_bound, schemas: self.schemas, lake: self.lake, } } } +/// How many connections the pool opens at most. +const POOL_MAX_SIZE: usize = 16; + +/// How many ListOffsets partition reads this engine has in flight at once, +/// across every request, counting an abandoned read until its statement +/// has ended, its connection has failed, or the bound on it has passed, +/// see [`cancel_and_return`]. +/// +/// Each read holds a pooled connection, so this limit leaves half the pool +/// for other requests. A read past this limit waits for a permit, and +/// [`ListOffsetsService`]'s deadline drops it if the wait lasts too long. +/// +/// [`ListOffsetsService`]: nisshi_storage::ListOffsetsService +const LIST_OFFSETS_READS: usize = POOL_MAX_SIZE / 2; + +/// How much longer than its connection's `statement_timeout` an abandoned +/// statement is waited for before its connection is closed, see +/// [`cancel_and_return`]. +/// +/// The server ends the statement itself at `statement_timeout`, and this +/// margin gives its answer time to arrive. +const ABANDONED_STATEMENT_MARGIN: Duration = Duration::from_secs(5); + /// Postgres statement timeout applied to every connection by default, so a stalled query /// (e.g. a lock wait) is aborted server-side instead of holding a pooled connection open /// indefinitely. An operator's own `?options=-c%20statement_timeout%3D...` in the connection @@ -196,6 +252,55 @@ fn merge_statement_timeout_option(existing: Option<&str>, default_ms: u64) -> St } } +/// Reads the `statement_timeout` that `options`, a libpq options startup +/// string, sets on each connection: the last one wins, as on the server. +/// The server reads `-c name=value`, `-cname=value` and `--name=value`. +/// +/// Returns `None` when `options` sets no timeout, disables it with `0`, or +/// sets a value the server rejects, in which case no connection opens. +fn statement_timeout_option(options: &str) -> Option { + let mut tokens = options.split_whitespace(); + let mut statement_timeout = None; + + while let Some(token) = tokens.next() { + let setting = match token { + "-c" => tokens.next(), + short if short.starts_with("-c") => short.strip_prefix("-c"), + long => long.strip_prefix("--"), + }; + + if let Some(value) = setting.and_then(|setting| setting.strip_prefix("statement_timeout=")) + { + statement_timeout = Some(statement_timeout_value(value)); + } + } + + statement_timeout + .flatten() + .filter(|timeout| !timeout.is_zero()) +} + +/// Parses a `statement_timeout` value as the server does: a number of +/// milliseconds, or a number with one of the server's time units, with a +/// fraction allowed either way. +fn statement_timeout_value(value: &str) -> Option { + let number = value.trim_end_matches(|c: char| c.is_ascii_alphabetic()); + let unit = &value[number.len()..]; + let count = number.parse::().ok()?; + + let seconds = match unit { + "us" => count / 1_000_000.0, + "" | "ms" => count / 1_000.0, + "s" => count, + "min" => count * 60.0, + "h" => count * 3_600.0, + "d" => count * 86_400.0, + _ => return None, + }; + + Duration::try_from_secs_f64(seconds).ok() +} + impl FromStr for Builder where C: Default, @@ -208,6 +313,9 @@ where let options = merge_statement_timeout_option(pg_config.get_options(), DEFAULT_STATEMENT_TIMEOUT_MS); + let abandoned_statement_bound = statement_timeout_option(&options) + .unwrap_or(Duration::from_millis(DEFAULT_STATEMENT_TIMEOUT_MS)) + + ABANDONED_STATEMENT_MARGIN; _ = pg_config.options(options); _ = pg_config.connect_timeout(DEFAULT_POOL_CREATE_TIMEOUT); @@ -230,19 +338,25 @@ where .map(|config| config.with_root_certificates(root_store)) .map(|config| config.with_no_client_auth())?; - let tls = tokio_postgres_rustls::MakeRustlsConnect::new(config); + let tls = MakeRustlsConnect::new(config); + let cancel = Some(CancelRequests { + config: pg_config.clone(), + tls: tls.clone(), + }); let mgr = Manager::from_config(pg_config, tls, mgr_config); let advertised_listener = Url::parse("tcp://127.0.0.1/")?; Pool::builder(mgr) - .max_size(16) + .max_size(POOL_MAX_SIZE) .runtime(Runtime::Tokio1) .wait_timeout(Some(DEFAULT_POOL_WAIT_TIMEOUT)) .create_timeout(Some(DEFAULT_POOL_CREATE_TIMEOUT)) .build() .map(|pool| Self { pool, + cancel, + abandoned_statement_bound, advertised_listener, node: N::default(), cluster: C::default(), @@ -253,6 +367,402 @@ where } } +/// What a cancel request needs to reach the server on a connection of its +/// own: the connection config, and the TLS connector. +#[derive(Clone)] +struct CancelRequests { + config: Config, + tls: MakeRustlsConnect, +} + +impl Debug for CancelRequests { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CancelRequests") + .field("config", &self.config) + .finish_non_exhaustive() + } +} + +/// A cancel request that did not reach the server, and the host it was +/// sent to. +#[derive(Debug)] +struct CancelFailed { + host: String, + error: Error, +} + +/// A socket to one configured host, with the host name that TLS verifies +/// and the name a log line gives the host. +struct CancelSocket { + stream: Either, + tls_hostname: Option, + host: String, +} + +impl CancelRequests { + /// Asks the server to cancel the statement that `token` identifies, and + /// returns the host it asked once the server has passed the request on + /// to its backend. + /// + /// [`CancelToken::cancel_query`] returns as soon as it has written the + /// request, so a statement sent on the connection after it could be the + /// one cancelled instead. The server closes the cancel connection once + /// it has signalled the backend, and libpq's `PQcancel` waits for that + /// close; so does this, through [`ClosedByPeer`]. + /// + /// This relies on a direct connection to PostgreSQL: a proxy that + /// closes the cancel connection before forwarding the request breaks + /// the ordering, for libpq as well. + /// + /// The request goes to the first configured host that accepts a + /// connection, as a new connection would. A server other than the one + /// running the statement ignores it, and the statement then runs until + /// it ends on its own or at `statement_timeout`. + async fn cancel(&self, token: &CancelToken) -> std::result::Result { + let CancelSocket { + stream, + tls_hostname, + host, + } = self.connect().await?; + + let tls = match MakeTlsConnect::::make_tls_connect( + &mut self.tls.clone(), + tls_hostname.as_deref().unwrap_or_default(), + ) { + Ok(tls) => tls, + Err(err) => { + return Err(CancelFailed { + host, + error: io::Error::other(err).into(), + }); + } + }; + + match timeout( + DEFAULT_POOL_CREATE_TIMEOUT, + token.cancel_query_raw(ClosedByPeer::new(stream), tls), + ) + .await + { + Ok(Ok(())) => Ok(host), + Ok(Err(err)) => Err(CancelFailed { + host, + error: err.into(), + }), + Err(elapsed) => Err(CancelFailed { + host, + error: io::Error::from(elapsed).into(), + }), + } + } + + /// Opens a socket to the first configured host that accepts one. The + /// error names the last host tried. + async fn connect(&self) -> std::result::Result { + let hosts = self.config.get_hosts(); + let ports = self.config.get_ports(); + let hostaddrs = self.config.get_hostaddrs(); + let connect_timeout = self + .config + .get_connect_timeout() + .copied() + .unwrap_or(DEFAULT_POOL_CREATE_TIMEOUT); + + let mut last = CancelFailed { + host: "none".into(), + error: io::Error::new(io::ErrorKind::NotFound, "no host configured").into(), + }; + + for (index, host) in hosts.iter().enumerate() { + let port = ports.get(index).or(ports.first()).copied().unwrap_or(5432); + + let name = match (hostaddrs.get(index), host) { + (Some(addr), _) => format!("{addr}:{port}"), + (None, Host::Tcp(host)) => format!("{host}:{port}"), + #[cfg(unix)] + (None, Host::Unix(dir)) => { + dir.join(format!(".s.PGSQL.{port}")).display().to_string() + } + }; + + let connect = async { + match (hostaddrs.get(index), host) { + (Some(addr), _) => TcpStream::connect((*addr, port)).await.map(Either::Left), + + (None, Host::Tcp(host)) => TcpStream::connect((host.as_str(), port)) + .await + .map(Either::Left), + + #[cfg(unix)] + (None, Host::Unix(dir)) => { + UnixStream::connect(dir.join(format!(".s.PGSQL.{port}"))) + .await + .map(Either::Right) + } + } + }; + + match timeout(connect_timeout, connect).await { + Ok(Ok(stream)) => { + let tls_hostname = match host { + Host::Tcp(hostname) => Some(hostname.clone()), + #[cfg(unix)] + Host::Unix(_) => None, + }; + + return Ok(CancelSocket { + stream, + tls_hostname, + host: name, + }); + } + + Ok(Err(err)) => { + last = CancelFailed { + host: name, + error: err.into(), + } + } + + Err(_) => { + last = CancelFailed { + host: name, + error: io::Error::new(io::ErrorKind::TimedOut, "connect timed out").into(), + } + } + } + } + + Err(last) + } +} + +/// The stream a cancel request travels on. +type CancelStream = ClosedByPeer>; + +/// A stream whose shutdown also waits for the peer to close its side. +struct ClosedByPeer { + stream: S, + writes_shut: bool, +} + +impl ClosedByPeer { + fn new(stream: S) -> Self { + Self { + stream, + writes_shut: false, + } + } +} + +impl AsyncRead for ClosedByPeer +where + S: AsyncRead + Unpin, +{ + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.stream).poll_read(cx, buf) + } +} + +impl AsyncWrite for ClosedByPeer +where + S: AsyncRead + AsyncWrite + Unpin, +{ + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.stream).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if !self.writes_shut { + match ready!(Pin::new(&mut self.stream).poll_shutdown(cx)) { + // A peer that has already closed makes some systems refuse + // the shutdown, and the close is what this waits for. + Ok(()) | Err(_) => self.writes_shut = true, + } + } + + let mut scratch = [0; 64]; + + loop { + let mut buf = ReadBuf::new(&mut scratch); + + match ready!(Pin::new(&mut self.stream).poll_read(cx, &mut buf)) { + Ok(()) if buf.filled().is_empty() => return Poll::Ready(Ok(())), + Ok(()) => {} + // An abrupt close is still a close. + Err(err) + if matches!( + err.kind(), + io::ErrorKind::ConnectionReset + | io::ErrorKind::ConnectionAborted + | io::ErrorKind::BrokenPipe + | io::ErrorKind::UnexpectedEof + ) => + { + return Poll::Ready(Ok(())); + } + Err(err) => return Poll::Ready(Err(err)), + } + } + } +} + +/// A pooled object that goes back to its pool only through +/// [`Abandonable::release`]. +/// +/// The pool recycles a connection after checking only that it is open +/// ([`RecyclingMethod::Fast`]). A connection dropped while a statement runs +/// on the server would go back to the pool still busy, and the next caller +/// would wait behind that statement. Dropping this guard instead passes the +/// object, still checked out of its pool, to `abandon`, which decides when +/// the object may go back. +struct Abandonable +where + M: managed::Manager, + F: FnOnce(managed::Object), +{ + object: Option>, + abandon: Option, +} + +impl Abandonable +where + M: managed::Manager, + F: FnOnce(managed::Object), +{ + fn new(object: managed::Object, abandon: F) -> Self { + Self { + object: Some(object), + abandon: Some(abandon), + } + } + + fn object(&self) -> &managed::Object { + self.object + .as_ref() + .expect("an abandonable object is held until release or drop") + } + + /// Returns the object to its pool, once its work has finished. + fn release(mut self) { + drop(self.object.take()); + } +} + +impl Drop for Abandonable +where + M: managed::Manager, + F: FnOnce(managed::Object), +{ + fn drop(&mut self) { + if let (Some(object), Some(abandon)) = (self.object.take(), self.abandon.take()) { + abandon(object); + } + } +} + +impl Debug for Abandonable +where + M: managed::Manager, + F: FnOnce(managed::Object), +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Abandonable") + .field("held", &self.object.is_some()) + .finish() + } +} + +/// Cancels the statement running on an abandoned `connection`, and returns +/// the connection to its pool once that statement has ended, so the pool +/// stays within its [`POOL_MAX_SIZE`] server connections. +/// +/// The connection and `permit` are held until the statement has ended, +/// whether or not the cancel request reached the server, so the statement +/// counts against [`LIST_OFFSETS_READS`] meanwhile. A cancel request that +/// misses its statement (a host other than the one running the statement, +/// or a proxy that does not forward it) then shows as ListOffsets reads +/// that wait for a permit, under the `queued` stage, rather than as server +/// connections beyond the pool. +/// +/// The statement has ended once an empty query on `connection` answers, +/// because the server runs the two in order. The connection is closed and +/// taken out of the pool instead, and the pool opens a replacement, when +/// that query fails or has not answered within `bound`. The bound sits +/// past the connection's `statement_timeout` +/// ([`ABANDONED_STATEMENT_MARGIN`]), at which the server ends the statement +/// itself unless an operator has disabled the timeout. deadpool-postgres +/// closes a dropped connection's socket at once, and a backend notices a +/// closed socket only when it next sends a result, so a statement still +/// running then keeps its backend until it ends. +async fn cancel_and_return( + connection: Object, + cancel: Option, + bound: Duration, + permit: OwnedSemaphorePermit, +) { + let _permit = permit; + let settle_by = Instant::now() + bound; + + let host = match cancel { + Some(cancel) => match cancel.cancel(&connection.cancel_token()).await { + Ok(host) => host, + + Err(CancelFailed { host, error }) => { + warn!( + host, + ?error, + ?bound, + "the cancel request for an abandoned statement did not reach the server; waiting for the statement to end" + ); + host + } + }, + + None => { + warn!( + ?bound, + "no cancel request is configured for an abandoned statement; waiting for it to end" + ); + "none".into() + } + }; + + match timeout_at(settle_by, connection.batch_execute("")).await { + Ok(Ok(())) => drop(connection), + + Ok(Err(error)) => { + warn!( + host, + ?error, + "an abandoned statement's connection failed before the statement ended; closing it" + ); + drop(Object::take(connection)); + } + + Err(_elapsed) => { + warn!( + host, + ?bound, + "an abandoned statement did not end within the bound; closing its connection" + ); + drop(Object::take(connection)); + } + } +} + #[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] struct Txn { name: String, @@ -320,6 +830,142 @@ impl Postgres { .map_err(Into::into) } + /// Checks out a connection for a ListOffsets read that a caller may drop + /// mid-statement, holding `permit` from [`LIST_OFFSETS_READS`]. + /// + /// Dropping the guard before [`Abandonable::release`] asks the server to + /// cancel the statement, so it stops instead of running until + /// `statement_timeout`, and keeps the connection and the permit until + /// the statement has ended, see [`cancel_and_return`]. + async fn abandonable_connection( + &self, + permit: OwnedSemaphorePermit, + ) -> Result>> { + let cancel = self.cancel.clone(); + let bound = self.abandoned_statement_bound; + + self.connection().await.map(|connection| { + Abandonable::new(connection, move |connection: Object| { + match Handle::try_current() { + // The task keeps the request's span, so its warnings + // name the peer, the correlation id and the partition. + Ok(runtime) => { + _ = runtime.spawn( + cancel_and_return(connection, cancel, bound, permit).in_current_span(), + ); + } + + // Without a runtime the cancel request cannot be sent. + // Closing the connection still keeps the pool from + // handing out one that is busy. + Err(_) => drop(Object::take(connection)), + } + }) + }) + } + + /// Reads `offsets` one statement at a time on `c`. + async fn list_offsets_on( + &self, + c: &Object, + isolation_level: IsolationLevel, + offsets: &[(Topition, ListOffset)], + ) -> Result> { + let mut responses = vec![]; + + for (topition, offset_type) in offsets { + let query = match (offset_type, isolation_level) { + (ListOffset::Earliest, _) => "list_earliest_offset.sql", + (ListOffset::Latest, IsolationLevel::ReadCommitted) => { + "list_latest_offset_committed.sql" + } + (ListOffset::Latest, IsolationLevel::ReadUncommitted) => { + "list_latest_offset_uncommitted.sql" + } + (ListOffset::Timestamp(_), _) => "list_latest_offset_timestamp.sql", + }; + + debug!(?query); + + let list_offset = match offset_type { + ListOffset::Earliest | ListOffset::Latest => self + .prepare_query_opt( + c, + query, + &[&self.cluster, &topition.topic(), &topition.partition()], + ) + .await + .inspect_err(|err| error!(?err, cluster = self.cluster, ?topition)), + + ListOffset::Timestamp(timestamp) => self + .prepare_query_opt( + c, + query, + &[ + &self.cluster.as_str(), + &topition.topic(), + &topition.partition(), + timestamp, + ], + ) + .await + .inspect_err(|err| error!(?err)), + } + .inspect_err(|err| { + error!(?err, cluster = self.cluster, ?topition); + }) + .inspect(|result| debug!(?result))? + .map_or_else( + || { + // No record row: Earliest/Latest answer 0, a Timestamp + // lookup answers no offset (see `ListOffsetResponse::offset`). + let (offset, timestamp) = match offset_type { + ListOffset::Earliest | ListOffset::Latest => (Some(0), None), + ListOffset::Timestamp(_) => (None, None), + }; + debug!( + cluster = self.cluster, + ?topition, + ?offset_type, + offset, + ?timestamp + ); + + Ok(ListOffsetResponse { + timestamp, + offset, + ..Default::default() + }) + }, + |row| { + debug!(?row); + + row.try_get::<_, i64>(0).map(Some).and_then(|offset| { + row.try_get::<_, SystemTime>(1).map(Some).map(|timestamp| { + debug!( + cluster = self.cluster, + ?topition, + ?offset_type, + offset, + ?timestamp + ); + + ListOffsetResponse { + timestamp, + offset, + ..Default::default() + } + }) + }) + }, + )?; + + responses.push((topition.clone(), list_offset)); + } + + Ok(responses) + } + fn sql_lookup(&self, key: &str) -> Result<&str> { crate::sql::SQL.get(key) } @@ -2729,110 +3375,23 @@ impl Storage for Postgres { Ok(offsets) } - #[instrument(skip_all)] + #[instrument(skip_all, fields(?offsets))] async fn list_offsets( &self, isolation_level: IsolationLevel, offsets: &[(Topition, ListOffset)], ) -> Result> { - debug!(cluster = self.cluster, ?isolation_level, ?offsets); + debug!(cluster = self.cluster, ?isolation_level); - let c = self.connection().await?; + let permit = deadline::queued(self.list_offsets_reads.clone().acquire_owned()).await?; + let connection = deadline::queued(self.abandonable_connection(permit)).await?; + let responses = self + .list_offsets_on(connection.object(), isolation_level, offsets) + .await; + connection.release(); - let mut responses = vec![]; - - for (topition, offset_type) in offsets { - let query = match (offset_type, isolation_level) { - (ListOffset::Earliest, _) => "list_earliest_offset.sql", - (ListOffset::Latest, IsolationLevel::ReadCommitted) => { - "list_latest_offset_committed.sql" - } - (ListOffset::Latest, IsolationLevel::ReadUncommitted) => { - "list_latest_offset_uncommitted.sql" - } - (ListOffset::Timestamp(_), _) => "list_latest_offset_timestamp.sql", - }; - - debug!(?query); - - let list_offset = match offset_type { - ListOffset::Earliest | ListOffset::Latest => self - .prepare_query_opt( - &c, - query, - &[&self.cluster, &topition.topic(), &topition.partition()], - ) - .await - .inspect_err(|err| error!(?err, cluster = self.cluster, ?topition)), - - ListOffset::Timestamp(timestamp) => self - .prepare_query_opt( - &c, - query, - &[ - &self.cluster.as_str(), - &topition.topic(), - &topition.partition(), - timestamp, - ], - ) - .await - .inspect_err(|err| error!(?err)), - } - .inspect_err(|err| { - error!(?err, cluster = self.cluster, ?topition); - }) - .inspect(|result| debug!(?result))? - .map_or_else( - || { - // No record row: Earliest/Latest answer 0, a Timestamp - // lookup answers no offset (see `ListOffsetResponse::offset`). - let (offset, timestamp) = match offset_type { - ListOffset::Earliest | ListOffset::Latest => (Some(0), None), - ListOffset::Timestamp(_) => (None, None), - }; - debug!( - cluster = self.cluster, - ?topition, - ?offset_type, - offset, - ?timestamp - ); - - Ok(ListOffsetResponse { - timestamp, - offset, - ..Default::default() - }) - }, - |row| { - debug!(?row); - - row.try_get::<_, i64>(0).map(Some).and_then(|offset| { - row.try_get::<_, SystemTime>(1).map(Some).map(|timestamp| { - debug!( - cluster = self.cluster, - ?topition, - ?offset_type, - offset, - ?timestamp - ); - - ListOffsetResponse { - timestamp, - offset, - ..Default::default() - } - }) - }) - }, - )?; - - responses.push((topition.clone(), list_offset)); - } - - Ok(responses).inspect(|r| debug!(?r)) - } + responses.inspect(|r| debug!(?r)) + } #[instrument(skip_all)] async fn metadata(&self, topics: Option<&[TopicId]>) -> Result { @@ -4177,6 +4736,87 @@ mod tests { ); } + /// The last `statement_timeout` in the options wins, as on the server, + /// in either spelling and with the server's time units. + #[test] + fn statement_timeout_option_reads_the_last_setting() { + assert_eq!( + Some(Duration::from_secs(30)), + statement_timeout_option("-c statement_timeout=30000") + ); + assert_eq!( + Some(Duration::from_secs(120)), + statement_timeout_option("-c statement_timeout=30000 -c statement_timeout=120s") + ); + assert_eq!( + Some(Duration::from_secs(120)), + statement_timeout_option("-c statement_timeout=30000 --statement_timeout=2min") + ); + assert_eq!( + Some(Duration::from_millis(100)), + statement_timeout_option( + "-c statement_timeout=30000 -c search_path=foo -c statement_timeout=100" + ) + ); + assert_eq!( + Some(Duration::from_secs(120)), + statement_timeout_option("-c statement_timeout=30000 -cstatement_timeout=120s") + ); + assert_eq!( + Some(Duration::from_secs(150)), + statement_timeout_option("-c statement_timeout=2.5min") + ); + assert_eq!( + Some(Duration::from_micros(500)), + statement_timeout_option("-c statement_timeout=500us") + ); + } + + /// No setting, a disabled timeout and a value the server would reject + /// all leave the default in place. + #[test] + fn statement_timeout_option_without_a_usable_setting() { + assert_eq!(None, statement_timeout_option("")); + assert_eq!(None, statement_timeout_option("-c search_path=foo")); + assert_eq!( + None, + statement_timeout_option("-c statement_timeout=30000 -c statement_timeout=0") + ); + assert_eq!(None, statement_timeout_option("-c statement_timeout=soon")); + assert_eq!( + None, + statement_timeout_option("-c statement_timeout=5fortnights") + ); + assert_eq!(None, statement_timeout_option("-c statement_timeout=-1")); + } + + /// The bound on an abandoned statement follows the connection's + /// `statement_timeout`: the default, or an operator's own in the URL. + #[test] + fn abandoned_statement_bound_follows_the_statement_timeout() -> Result<()> { + let storage = Postgres::builder(CONNECTION)? + .cluster("default") + .node(0) + .build(); + assert_eq!( + Duration::from_millis(DEFAULT_STATEMENT_TIMEOUT_MS) + ABANDONED_STATEMENT_MARGIN, + storage.abandoned_statement_bound + ); + + let storage = Postgres::builder(&format!( + "{CONNECTION}?options=-c%20statement_timeout%3D120000" + ))? + .cluster("longer") + .node(0) + .build(); + assert_eq!( + Duration::from_secs(120) + ABANDONED_STATEMENT_MARGIN, + storage.abandoned_statement_bound + ); + + Ok(()) + } + /// Regression test for a real blocker found in review: deadpool's `PoolBuilder::build()` /// returns `Err(BuildError::NoRuntimeSpecified)` if `wait_timeout`/`create_timeout` are /// set without also configuring `.runtime(..)`. This doesn't need a live Postgres: @@ -4204,6 +4844,508 @@ mod tests { assert!(error.to_string().contains("vacuum_into"), "{error}"); } + /// Hands out numbered objects, so a test can tell a fresh object from a + /// recycled one. + #[derive(Debug, Default)] + struct Numbered(std::sync::atomic::AtomicU32); + + impl managed::Manager for Numbered { + type Type = u32; + type Error = std::convert::Infallible; + + async fn create(&self) -> std::result::Result { + Ok(self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst)) + } + + async fn recycle( + &self, + _object: &mut u32, + _metrics: &managed::Metrics, + ) -> managed::RecycleResult { + Ok(()) + } + } + + /// A released object goes back to the pool. A dropped one reaches + /// `abandon` still checked out, so the pool cannot hand it out until + /// `abandon` lets it go. + #[tokio::test] + async fn abandoned_object_is_handed_to_abandon_still_checked_out() -> Result<()> { + let pool = managed::Pool::::builder(Numbered::default()) + .max_size(1) + .build() + .map_err(|err| Error::Message(err.to_string()))?; + + let abandoned = Arc::new(std::sync::Mutex::new(Vec::new())); + + let get = async || { + let abandoned = abandoned.clone(); + pool.get() + .await + .map(|object| { + Abandonable::new(object, move |object| { + abandoned.lock().expect("abandoned").push(object) + }) + }) + .map_err(|err| Error::Message(err.to_string())) + }; + + let first = get().await?; + assert_eq!(0, **first.object()); + first.release(); + assert_eq!(1, pool.status().available); + + let second = get().await?; + assert_eq!(0, **second.object()); + drop(second); + + assert_eq!(1, pool.status().size); + assert_eq!(0, pool.status().available); + assert!( + timeout(Duration::from_millis(50), pool.get()) + .await + .is_err(), + "the abandoned object is still checked out" + ); + + let held = abandoned + .lock() + .expect("abandoned") + .pop() + .expect("abandon received the object"); + assert_eq!(0, *held); + drop(held); + + let third = get().await?; + assert_eq!(0, **third.object()); + third.release(); + + Ok(()) + } + + /// Shutting down a `ClosedByPeer` stream finishes only once the peer + /// has closed its side, and everything the peer sent meanwhile is + /// discarded. + #[tokio::test] + async fn closed_by_peer_shutdown_waits_for_the_peer() -> Result<()> { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + let (ours, mut peer) = tokio::io::duplex(64); + let mut ours = ClosedByPeer::new(ours); + + ours.write_all(b"cancel").await?; + ours.flush().await?; + + let mut shutdown = tokio::spawn(async move { ours.shutdown().await.map(|()| ours) }); + + let mut request = [0; 6]; + _ = peer.read_exact(&mut request).await?; + assert_eq!(b"cancel", &request); + peer.write_all(b"busy").await?; + + assert!( + timeout(Duration::from_millis(100), &mut shutdown) + .await + .is_err(), + "shutdown finished before the peer closed" + ); + + drop(peer); + + let ours = timeout(Duration::from_secs(1), shutdown) + .await + .map_err(|_| Error::Message("shutdown did not finish after the peer closed".into()))? + .map_err(|err| Error::Message(err.to_string()))??; + assert!(ours.writes_shut); + + Ok(()) + } + + /// At most `LIST_OFFSETS_READS` ListOffsets reads reach the pool at + /// once; the rest wait for a permit, and the permits come back when the + /// reads are dropped. + /// + /// The pool has no permits at all, so each read that has a ListOffsets + /// permit waits in the pool, where the pool counts it, without a server. + #[tokio::test] + async fn list_offsets_reads_are_limited_before_the_pool() -> Result<()> { + let pool = Pool::builder(Manager::new( + Config::from_str(CONNECTION)?, + tokio_postgres::NoTls, + )) + .max_size(0) + .runtime(Runtime::Tokio1) + .wait_timeout(Some(DEFAULT_POOL_WAIT_TIMEOUT)) + .build()?; + + let storage = Postgres { + pool, + ..Postgres::builder(CONNECTION)? + .cluster(alphanumeric_string(15).as_str()) + .node(rng().random_range(0..i32::MAX)) + .build() + }; + + let topition = Topition::new("limited", 0); + let reads = (0..LIST_OFFSETS_READS + 3) + .map(|_| { + let storage = storage.clone(); + let topition = topition.clone(); + + tokio::spawn(async move { + storage + .list_offsets( + IsolationLevel::ReadUncommitted, + &[(topition, ListOffset::Latest)], + ) + .await + }) + }) + .collect::>(); + + let settled = Instant::now() + Duration::from_secs(5); + while storage.pool.status().waiting < LIST_OFFSETS_READS { + assert!(Instant::now() < settled, "{:?}", storage.pool.status()); + tokio::time::sleep(Duration::from_millis(10)).await; + } + + assert_eq!(LIST_OFFSETS_READS, storage.pool.status().waiting); + assert_eq!(0, storage.list_offsets_reads.available_permits()); + + for read in &reads { + read.abort(); + } + for read in reads { + let outcome = read.await; + assert!( + outcome.as_ref().is_err_and(|err| err.is_cancelled()), + "{outcome:?}" + ); + } + + assert_eq!(0, storage.pool.status().waiting); + assert_eq!( + LIST_OFFSETS_READS, + storage.list_offsets_reads.available_permits() + ); + + Ok(()) + } + + /// Dropping an abandonable connection mid-statement stops the statement + /// on the server and returns the connection to the pool, with its + /// permit. + #[tokio::test] + async fn abandoned_statement_is_cancelled_and_its_connection_returned() -> Result<()> { + let storage = Postgres::builder(CONNECTION)? + .cluster(alphanumeric_string(15).as_str()) + .node(rng().random_range(0..i32::MAX)) + .build(); + + let monitor = match storage.connection().await { + Ok(monitor) => monitor, + Err(err) => { + eprintln!( + "skipping abandoned_statement_is_cancelled_and_its_connection_returned: {err:?}" + ); + return Ok(()); + } + }; + + let marker = alphanumeric_string(15); + let sleep = format!("SELECT pg_sleep(30) /* {marker} */"); + + let permit = storage.list_offsets_reads.clone().acquire_owned().await?; + let connection = storage.abandonable_connection(permit).await?; + assert_eq!( + LIST_OFFSETS_READS - 1, + storage.list_offsets_reads.available_permits() + ); + + assert!( + timeout( + Duration::from_millis(200), + connection.object().batch_execute(&sleep) + ) + .await + .is_err(), + "pg_sleep(30) is still running" + ); + drop(connection); + + // The cancel task holds the permit, so the statement still counts + // against the limit while it runs. + assert_eq!( + LIST_OFFSETS_READS - 1, + storage.list_offsets_reads.available_permits() + ); + + let settled = Instant::now() + Duration::from_secs(10); + loop { + let running = monitor + .query_one( + "SELECT count(*) FROM pg_stat_activity WHERE state <> 'idle' AND query LIKE $1", + &[&format!("%{marker}%")], + ) + .await? + .get::<_, i64>(0); + let status = storage.pool.status(); + + if running == 0 && status.available == 1 { + break; + } + + assert!( + Instant::now() < settled, + "statement still running on {running} backend(s), pool {status:?}" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + + assert_eq!(2, storage.pool.status().size); + assert_eq!( + LIST_OFFSETS_READS, + storage.list_offsets_reads.available_permits() + ); + + Ok(()) + } + + /// How many backends run a statement marked `marker`, other than + /// `monitor`'s own. + async fn backends_running(monitor: &Object, marker: &str) -> Result { + monitor + .query_one( + "SELECT count(*) FROM pg_stat_activity WHERE state <> 'idle' AND query LIKE $1 AND pid <> pg_backend_pid()", + &[&format!("%{marker}%")], + ) + .await + .map(|row| row.get::<_, i64>(0)) + .map_err(Into::into) + } + + /// Ends every statement marked `marker` through the server, as the + /// cancel request would have, and returns how many it ended. + async fn cancel_backends(monitor: &Object, marker: &str) -> Result { + monitor + .query( + "SELECT pg_cancel_backend(pid) FROM pg_stat_activity WHERE state <> 'idle' AND query LIKE $1 AND pid <> pg_backend_pid()", + &[&format!("%{marker}%")], + ) + .await + .map(|rows| rows.len()) + .map_err(Into::into) + } + + /// Points `storage`'s cancel requests at a port nothing listens on, so + /// each one fails without reaching the server, as one sent to a host + /// other than the one running the statement would. + async fn with_cancel_requests_refused(storage: Postgres) -> Result { + let closed_port = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await? + .local_addr()? + .port(); + + let tls = storage + .cancel + .as_ref() + .expect("cancel requests are configured from a URL") + .tls + .clone(); + + let mut config = Config::new(); + _ = config.host("127.0.0.1").port(closed_port); + + Ok(Postgres { + cancel: Some(CancelRequests { config, tls }), + ..storage + }) + } + + /// Runs `SELECT pg_sleep(30)` marked `marker` on an abandonable + /// connection and drops the connection mid-statement, returning once the + /// cancel task has had time to send its cancel request. + async fn abandon_sleeping_statement(storage: &Postgres, marker: &str) -> Result<()> { + let sleep = format!("SELECT pg_sleep(30) /* {marker} */"); + + let permit = storage.list_offsets_reads.clone().acquire_owned().await?; + let connection = storage.abandonable_connection(permit).await?; + + assert!( + timeout( + Duration::from_millis(200), + connection.object().batch_execute(&sleep) + ) + .await + .is_err(), + "pg_sleep(30) is still running" + ); + drop(connection); + + tokio::time::sleep(Duration::from_millis(500)).await; + + Ok(()) + } + + /// Waits until `settled` holds, failing after 10 seconds. + async fn wait_until(storage: &Postgres, settled: impl Fn(&Postgres) -> bool) { + let give_up = Instant::now() + Duration::from_secs(10); + + while !settled(storage) { + assert!( + Instant::now() < give_up, + "pool {:?}, permits {}", + storage.pool.status(), + storage.list_offsets_reads.available_permits() + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + + /// A cancel request that does not reach the server leaves the + /// statement running, and the connection and its permit stay held + /// until the statement ends; then both come back. + #[tokio::test] + async fn abandoned_statement_outliving_its_cancel_keeps_its_connection_and_permit() -> Result<()> + { + let storage = Postgres::builder(CONNECTION)? + .cluster(alphanumeric_string(15).as_str()) + .node(rng().random_range(0..i32::MAX)) + .build(); + + let monitor = match storage.connection().await { + Ok(monitor) => monitor, + Err(err) => { + eprintln!( + "skipping abandoned_statement_outliving_its_cancel_keeps_its_connection_and_permit: {err:?}" + ); + return Ok(()); + } + }; + + let storage = with_cancel_requests_refused(storage).await?; + let marker = alphanumeric_string(15); + + abandon_sleeping_statement(&storage, &marker).await?; + + assert_eq!(1, backends_running(&monitor, &marker).await?); + assert_eq!( + LIST_OFFSETS_READS - 1, + storage.list_offsets_reads.available_permits() + ); + assert_eq!(0, storage.pool.status().available); + assert_eq!(2, storage.pool.status().size); + + assert_eq!(1, cancel_backends(&monitor, &marker).await?); + + wait_until(&storage, |storage| storage.pool.status().available == 1).await; + + assert_eq!(0, backends_running(&monitor, &marker).await?); + assert_eq!(2, storage.pool.status().size); + assert_eq!( + LIST_OFFSETS_READS, + storage.list_offsets_reads.available_permits() + ); + + Ok(()) + } + + /// A statement still running at the bound has its connection closed + /// and taken out of the pool, and its permit comes back; the backend + /// runs on until the statement ends. + #[tokio::test] + async fn abandoned_statement_still_running_at_the_bound_loses_its_connection() -> Result<()> { + let storage = Postgres::builder(CONNECTION)? + .cluster(alphanumeric_string(15).as_str()) + .node(rng().random_range(0..i32::MAX)) + .build(); + + let monitor = match storage.connection().await { + Ok(monitor) => monitor, + Err(err) => { + eprintln!( + "skipping abandoned_statement_still_running_at_the_bound_loses_its_connection: {err:?}" + ); + return Ok(()); + } + }; + + let storage = Postgres { + abandoned_statement_bound: Duration::from_secs(2), + ..with_cancel_requests_refused(storage).await? + }; + let marker = alphanumeric_string(15); + + abandon_sleeping_statement(&storage, &marker).await?; + + assert_eq!( + LIST_OFFSETS_READS - 1, + storage.list_offsets_reads.available_permits() + ); + assert_eq!(2, storage.pool.status().size); + + wait_until(&storage, |storage| { + storage.list_offsets_reads.available_permits() == LIST_OFFSETS_READS + }) + .await; + + assert_eq!(1, storage.pool.status().size); + assert_eq!(0, storage.pool.status().available); + assert_eq!(1, backends_running(&monitor, &marker).await?); + + assert_eq!(1, cancel_backends(&monitor, &marker).await?); + + Ok(()) + } + + /// A connection that fails while its abandoned statement runs is taken + /// out of the pool, and its permit comes back. + #[tokio::test] + async fn abandoned_statement_whose_connection_fails_loses_it() -> Result<()> { + let storage = Postgres::builder(CONNECTION)? + .cluster(alphanumeric_string(15).as_str()) + .node(rng().random_range(0..i32::MAX)) + .build(); + + let monitor = match storage.connection().await { + Ok(monitor) => monitor, + Err(err) => { + eprintln!("skipping abandoned_statement_whose_connection_fails_loses_it: {err:?}"); + return Ok(()); + } + }; + + let storage = with_cancel_requests_refused(storage).await?; + let marker = alphanumeric_string(15); + + abandon_sleeping_statement(&storage, &marker).await?; + + assert_eq!( + LIST_OFFSETS_READS - 1, + storage.list_offsets_reads.available_permits() + ); + + let terminated = monitor + .query( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE state <> 'idle' AND query LIKE $1 AND pid <> pg_backend_pid()", + &[&format!("%{marker}%")], + ) + .await? + .len(); + assert_eq!(1, terminated); + + wait_until(&storage, |storage| { + storage.list_offsets_reads.available_permits() == LIST_OFFSETS_READS + }) + .await; + + assert_eq!(1, storage.pool.status().size); + assert_eq!(0, storage.pool.status().available); + assert_eq!(0, backends_running(&monitor, &marker).await?); + + Ok(()) + } + #[test] fn pool_error_timeout_maps_to_retriable_error_code() { let error = Error::from(PoolError::Timeout(TimeoutType::Wait)); diff --git a/nisshi-storage/Cargo.toml b/nisshi-storage/Cargo.toml index fd025160..43471d16 100644 --- a/nisshi-storage/Cargo.toml +++ b/nisshi-storage/Cargo.toml @@ -71,6 +71,7 @@ criterion.workspace = true dotenv.workspace = true pretty_assertions.workspace = true nisshi-schema.workspace = true +opentelemetry_sdk.workspace = true tempfile.workspace = true time.workspace = true tracing-subscriber.workspace = true diff --git a/nisshi-storage/src/lib.rs b/nisshi-storage/src/lib.rs index 0c7fc3d5..da29ae0a 100644 --- a/nisshi-storage/src/lib.rs +++ b/nisshi-storage/src/lib.rs @@ -199,6 +199,7 @@ pub use batch::ProduceRequestBatcher; pub use latency::LatencyIntroducingStorage; pub use producer::{check_claim, producer_claim}; pub use proxy::SemaphoreProxy; +pub use service::deadline; pub use service::{ AlterUserScramCredentialsService, ChannelRequestLayer, ChannelRequestService, diff --git a/nisshi-storage/src/proxy.rs b/nisshi-storage/src/proxy.rs index 532306b4..a1380981 100644 --- a/nisshi-storage/src/proxy.rs +++ b/nisshi-storage/src/proxy.rs @@ -42,7 +42,7 @@ use crate::{ BrokerRegistrationRequest, GroupDetail, ListOffsetResponse, METER, MetadataResponse, NamedGroupDetail, OffsetCommitRequest, OffsetStage, ProducerIdResponse, Result, ScramCredential, Storage, TopicId, Topition, TxnAddPartitionsRequest, TxnAddPartitionsResponse, - TxnOffsetCommitRequest, UpdateError, Version, + TxnOffsetCommitRequest, UpdateError, Version, service::deadline, }; static SEMAPHORE_ACQUIRE_DURATION: LazyLock> = LazyLock::new(|| { @@ -215,12 +215,14 @@ where offsets: &[(Topition, ListOffset)], ) -> Result> { let start = SystemTime::now(); - let _permit = self.semaphore.acquire().await.inspect(|_| { - SEMAPHORE_ACQUIRE_DURATION.record( - elapsed_millis(start), - &[KeyValue::new("operation", "list_offsets")], - ) - })?; + let _permit = deadline::queued(self.semaphore.acquire()) + .await + .inspect(|_| { + SEMAPHORE_ACQUIRE_DURATION.record( + elapsed_millis(start), + &[KeyValue::new("operation", "list_offsets")], + ) + })?; self.storage.list_offsets(isolation_level, offsets).await } diff --git a/nisshi-storage/src/service.rs b/nisshi-storage/src/service.rs index f9f8ed6f..86b29bf7 100644 --- a/nisshi-storage/src/service.rs +++ b/nisshi-storage/src/service.rs @@ -16,6 +16,7 @@ mod alter_user_scram_credentials; mod consumer_group_describe; mod create_acls; mod create_topics; +pub mod deadline; mod delete_groups; mod delete_records; mod delete_topics; @@ -422,8 +423,7 @@ impl Service for RequestChannelService { STORAGE_CHANNEL_CAPACITY.record(capacity as u64, &attributes); debug!(operation, capacity); - self.tx - .reserve() + deadline::queued(self.tx.reserve()) .await .map(|permit| permit.send((input, resp_tx))) .inspect(|_| { diff --git a/nisshi-storage/src/service/deadline.rs b/nisshi-storage/src/service/deadline.rs new file mode 100644 index 00000000..5f01068c --- /dev/null +++ b/nisshi-storage/src/service/deadline.rs @@ -0,0 +1,340 @@ +// Copyright ⓒ 2024-2026 Peter Morgan +// +// 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. + +//! Deadlines for storage reads made while answering a client request. +//! +//! A slow storage engine must not hold a request past the point where the +//! client gives up on it: the client reconnects and sends the request again, +//! while the abandoned one keeps running. Each client abandons a request at +//! its own read deadline: +//! +//! | Client | ListOffsets | Fetch | +//! |------------|-------------|--------------------------------| +//! | Java | 30s | 30s (`max_wait` not added) | +//! | librdkafka | 60s | 60s + `fetch.wait.max.ms` | +//! | franz-go | 10s | 10s + `max_wait` | +//! +//! The deadlines here sit under all three with room for the response. + +use std::{cell::Cell, future::Future, pin::pin, sync::LazyLock}; + +use opentelemetry::{KeyValue, metrics::Counter}; +use tokio::time::{Duration, Instant, timeout_at}; + +use crate::METER; + +/// How long ListOffsets may spend in storage, under franz-go's flat 10s. +/// +/// A partition still unread at this deadline answers `REQUEST_TIMED_OUT` +/// with offset and timestamp -1. Kafka 4.0 sends the same answer for a +/// partition whose remote-storage lookup expires in its list-offsets +/// purgatory ([DelayedRemoteListOffsets.scala#L46], [#L144-L149]), so +/// clients already retry it. +/// +/// [DelayedRemoteListOffsets.scala#L46]: https://github.com/apache/kafka/blob/4.0.0/core/src/main/scala/kafka/server/DelayedRemoteListOffsets.scala#L46 +/// [#L144-L149]: https://github.com/apache/kafka/blob/4.0.0/core/src/main/scala/kafka/server/DelayedRemoteListOffsets.scala#L144-L149 +pub(crate) const LIST_OFFSETS_READ_DEADLINE: Duration = Duration::from_secs(5); + +static READ_DEADLINE_EXCEEDED: LazyLock> = LazyLock::new(|| { + METER + .u64_counter("nisshi_storage_read_deadline_exceeded") + .with_description("Storage reads abandoned at the request deadline") + .build() +}); + +/// What a read bounded by [`within`] was doing when its deadline passed. +/// +/// An operator reads it from the `stage` attribute of the +/// `nisshi_storage_read_deadline_exceeded` counter: `reading` points at a +/// slow storage engine, `queued` at an engine's limit on reads in flight +/// being used up by slow reads, and `not_started` at a request with more +/// partitions than could be read before the deadline. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Missed { + /// The deadline had passed before the read was polled, so no storage + /// call was made. + NotStarted, + /// The read was waiting in [`queued`] for its storage engine's limit on + /// reads in flight, for a pooled connection, or for room on its request + /// channel. + Queued, + /// The read was in storage. + Reading, +} + +impl Missed { + pub(crate) fn stage(self) -> &'static str { + match self { + Self::NotStarted => "not_started", + Self::Queued => "queued", + Self::Reading => "reading", + } + } +} + +tokio::task_local! { + /// Whether the read that the enclosing [`within`] bounds is in [`queued`]. + static QUEUED: Cell; +} + +/// Runs `wait`, a storage engine's wait for its own limit on reads in +/// flight, for a pooled connection, or for room on its request channel, so +/// that a read abandoned during it is counted under the `queued` stage +/// rather than `reading`. +/// +/// The stage is a task-local of the task that called the [`Storage`] +/// method, so `wait` counts only when it is awaited on that task. A wait on +/// a task the engine spawns, such as the task that serves the request +/// channel in `mpsc` mode, is not seen and counts as `reading`. +/// +/// Outside a read that a deadline bounds, this only runs `wait`. +/// +/// [`Storage`]: crate::Storage +pub async fn queued(wait: F) -> F::Output +where + F: Future, +{ + _ = QUEUED.try_with(|queued| queued.set(true)); + let output = wait.await; + _ = QUEUED.try_with(|queued| queued.set(false)); + output +} + +/// Runs `read` until `deadline`, returning what it was doing if it did not +/// finish. +/// +/// A read is not started at all once the deadline has passed: +/// [`tokio::time::timeout_at`] polls its future once before looking at the +/// clock, which would send a storage request only to drop it. +/// +/// Dropping the read cancels nisshi's side of it. Work the storage engine +/// runs on its own tasks, or a statement already sent to a database server, +/// carries on until it finishes. +pub(crate) async fn within( + operation: &'static str, + deadline: Instant, + read: F, +) -> Result +where + F: Future, +{ + within_counting(&READ_DEADLINE_EXCEEDED, operation, deadline, read).await +} + +/// Runs [`within`], counting each abandoned read in `exceeded`. +async fn within_counting( + exceeded: &Counter, + operation: &'static str, + deadline: Instant, + read: F, +) -> Result +where + F: Future, +{ + let outcome = if Instant::now() >= deadline { + Err(Missed::NotStarted) + } else { + let mut read = pin!(QUEUED.scope(Cell::new(false), timeout_at(deadline, read))); + + match read.as_mut().await { + Ok(output) => Ok(output), + Err(_elapsed) => Err(if read.take_value().is_some_and(|queued| queued.get()) { + Missed::Queued + } else { + Missed::Reading + }), + } + }; + + if let Err(missed) = outcome { + exceeded.add( + 1, + &[ + KeyValue::new("operation", operation), + KeyValue::new("stage", missed.stage()), + ], + ); + } + + outcome +} + +#[cfg(test)] +mod tests { + use std::{ + future::pending, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + }; + + use opentelemetry::metrics::MeterProvider as _; + use opentelemetry_sdk::metrics::{ + InMemoryMetricExporter, PeriodicReader, SdkMeterProvider, + data::{AggregatedMetrics, MetricData, ResourceMetrics, ScopeMetrics, SumDataPoint}, + }; + + use super::*; + + #[tokio::test(start_paused = true)] + async fn finishes_before_the_deadline() { + let deadline = Instant::now() + Duration::from_secs(1); + assert_eq!(Ok(7), within("test", deadline, async { 7 }).await); + } + + #[tokio::test(start_paused = true)] + async fn gives_up_at_the_deadline() { + let started_at = Instant::now(); + let deadline = started_at + Duration::from_secs(1); + + assert_eq!( + Err(Missed::Reading), + within("test", deadline, pending::<()>()).await + ); + assert_eq!(Duration::from_secs(1), started_at.elapsed()); + } + + /// Once the deadline has passed, the read is never polled. + #[tokio::test(start_paused = true)] + async fn does_not_start_after_the_deadline() { + let deadline = Instant::now(); + tokio::time::advance(Duration::from_millis(1)).await; + + let polled = Arc::new(AtomicBool::new(false)); + + let read = { + let polled = polled.clone(); + async move { polled.store(true, Ordering::SeqCst) } + }; + + assert_eq!( + Err(Missed::NotStarted), + within("test", deadline, read).await + ); + assert!(!polled.load(Ordering::SeqCst)); + } + + /// A read abandoned while it waits in `queued` is a queued miss; one + /// whose wait finished before the deadline is a reading miss. + #[tokio::test(start_paused = true)] + async fn tells_a_queued_read_from_a_reading_one() { + let deadline = Instant::now() + Duration::from_secs(1); + + assert_eq!( + Err(Missed::Queued), + within("test", deadline, queued(pending::<()>())).await + ); + + let deadline = Instant::now() + Duration::from_secs(1); + + assert_eq!( + Err(Missed::Reading), + within("test", deadline, async { + queued(async {}).await; + pending::<()>().await + }) + .await + ); + } + + /// Two reads bounded at once each keep their own stage. + #[tokio::test(start_paused = true)] + async fn reads_bounded_at_once_keep_their_own_stage() { + let deadline = Instant::now() + Duration::from_secs(1); + + let (queued_read, reading) = tokio::join!( + within("test", deadline, queued(pending::<()>())), + within("test", deadline, pending::<()>()) + ); + + assert_eq!(Err(Missed::Queued), queued_read); + assert_eq!(Err(Missed::Reading), reading); + } + + /// Sums the `nisshi_storage_read_deadline_exceeded` counter for + /// `operation` and `stage`. + fn exceeded(exporter: &InMemoryMetricExporter, operation: &str, stage: &str) -> u64 { + exporter + .get_finished_metrics() + .expect("finished metrics") + .iter() + .flat_map(ResourceMetrics::scope_metrics) + .flat_map(ScopeMetrics::metrics) + .filter(|metric| metric.name() == "nisshi_storage_read_deadline_exceeded") + .filter_map(|metric| match metric.data() { + AggregatedMetrics::U64(MetricData::Sum(sum)) => Some( + sum.data_points() + .filter(|point| { + let has = |key: &str, value: &str| { + point.attributes().any(|attribute| { + attribute.key.as_str() == key + && attribute.value.as_str() == value + }) + }; + + has("operation", operation) && has("stage", stage) + }) + .map(SumDataPoint::value) + .sum::(), + ), + _ => None, + }) + .sum() + } + + /// Each abandoned read adds one to the counter under its stage, and a + /// read that finishes in time adds nothing. + #[tokio::test(start_paused = true)] + async fn counts_each_abandoned_read_by_stage() { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_reader(PeriodicReader::builder(exporter.clone()).build()) + .build(); + let counter = provider + .meter("test") + .u64_counter("nisshi_storage_read_deadline_exceeded") + .build(); + + let deadline = || Instant::now() + Duration::from_secs(1); + + assert_eq!( + Ok(7), + within_counting(&counter, "on_time", deadline(), async { 7 }).await + ); + assert_eq!( + Err(Missed::Reading), + within_counting(&counter, "stalled", deadline(), pending::<()>()).await + ); + assert_eq!( + Err(Missed::Reading), + within_counting(&counter, "stalled", deadline(), pending::<()>()).await + ); + assert_eq!( + Err(Missed::Queued), + within_counting(&counter, "stalled", deadline(), queued(pending::<()>())).await + ); + assert_eq!( + Err(Missed::NotStarted), + within_counting(&counter, "stalled", Instant::now(), async {}).await + ); + + provider.force_flush().expect("flush"); + + assert_eq!(0, exceeded(&exporter, "on_time", "reading")); + assert_eq!(2, exceeded(&exporter, "stalled", "reading")); + assert_eq!(1, exceeded(&exporter, "stalled", "queued")); + assert_eq!(1, exceeded(&exporter, "stalled", "not_started")); + } +} diff --git a/nisshi-storage/src/service/list_offsets.rs b/nisshi-storage/src/service/list_offsets.rs index 5806f69c..d051509b 100644 --- a/nisshi-storage/src/service/list_offsets.rs +++ b/nisshi-storage/src/service/list_offsets.rs @@ -12,16 +12,32 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::{collections::BTreeSet, ops::Deref as _}; +use std::{collections::BTreeSet, slice}; +use futures::{StreamExt as _, stream}; use nisshi_sans_io::{ - ApiKey, IsolationLevel, ListOffset, ListOffsetsRequest, ListOffsetsResponse, RequestInput, + ApiKey, ErrorCode, IsolationLevel, ListOffset, ListOffsetsRequest, ListOffsetsResponse, + RequestInput, list_offsets_response::{ListOffsetsPartitionResponse, ListOffsetsTopicResponse}, }; use rama::Service; -use tracing::{debug, error, instrument}; +use tokio::time::Instant; +use tracing::{debug, error, instrument, warn}; -use crate::{Error, Result, Storage, Topition}; +use crate::{ + Error, ListOffsetResponse, Result, Storage, Topition, + service::deadline::{LIST_OFFSETS_READ_DEADLINE, Missed, within}, +}; + +/// How many partitions of one ListOffsets request are read from storage at +/// once. +/// +/// A storage engine with a connection pool adds its own limit across all +/// requests, where it can size that limit from the pool. +const LIST_OFFSETS_CONCURRENCY: usize = 4; + +/// How many timed-out partitions the deadline warning names. +const TIMED_OUT_SAMPLE: usize = 8; /// A [`Service`] using its [`Storage`] taking [`ListOffsetsRequest`] returning [`ListOffsetsResponse`]. /// ```no_run @@ -99,6 +115,91 @@ impl ApiKey for ListOffsetsService { const KEY: i16 = ListOffsetsRequest::KEY; } +impl ListOffsetsService +where + G: Storage, +{ + /// Reads each partition's offset separately, a few at a time, each + /// bounded by one deadline for the whole request. A partition still + /// unanswered at the deadline gets `REQUEST_TIMED_OUT`, which clients + /// retry, so one slow partition neither delays the others nor holds the + /// request past the client's own timeout. A storage error still fails + /// the whole request. + async fn list_offsets( + &self, + isolation_level: IsolationLevel, + offsets: &[(Topition, ListOffset)], + ) -> Result> { + let started_at = Instant::now(); + let deadline = started_at + LIST_OFFSETS_READ_DEADLINE; + + // The reads are built up front (but not started) rather than in a + // `StreamExt::map` closure, whose lifetimes keep `serve` from being `Send`. + let reads = offsets + .iter() + .enumerate() + .map(|(index, offset)| async move { + let answer = within( + "list_offsets", + deadline, + self.storage + .list_offsets(isolation_level, slice::from_ref(offset)), + ) + .await; + + (index, offset, answer) + }) + .collect::>(); + + let mut answers = stream::iter(reads).buffer_unordered(LIST_OFFSETS_CONCURRENCY); + + let mut responses = vec![Vec::new(); offsets.len()]; + let mut timed_out = Vec::new(); + + while let Some((index, (topition, _), answer)) = answers.next().await { + responses[index] = match answer { + Ok(answer) => answer?, + + Err(missed) => { + timed_out.push((topition, missed)); + + vec![( + topition.clone(), + ListOffsetResponse { + error_code: ErrorCode::RequestTimedOut, + timestamp: None, + offset: Some(-1), + }, + )] + } + }; + } + + if !timed_out.is_empty() { + let count = |stage: Missed| { + timed_out + .iter() + .filter(|(_, missed)| *missed == stage) + .count() + }; + + warn!( + deadline = ?LIST_OFFSETS_READ_DEADLINE, + elapsed = ?started_at.elapsed(), + timed_out = timed_out.len(), + partitions = offsets.len(), + reading = count(Missed::Reading), + queued = count(Missed::Queued), + not_started = count(Missed::NotStarted), + sample = ?&timed_out[..timed_out.len().min(TIMED_OUT_SAMPLE)], + "list offsets answered REQUEST_TIMED_OUT for partitions unread at the deadline" + ); + } + + Ok(responses.into_iter().flatten().collect()) + } +} + impl Service for ListOffsetsService where G: Storage, @@ -133,8 +234,7 @@ where } } - self.storage - .list_offsets(isolation_level, offsets.deref()) + self.list_offsets(isolation_level, &offsets) .await .inspect(|r| debug!(?r, ?offsets)) .inspect_err(|err| error!(?err, ?offsets)) @@ -193,3 +293,584 @@ where .inspect(|r| debug!(?r)) } } + +#[cfg(test)] +mod tests { + use std::{ + collections::BTreeMap, + future::pending, + sync::{Arc, Mutex}, + time::SystemTime, + }; + + use async_trait::async_trait; + use nisshi_sans_io::{ + ConfigResource, ErrorCode, IsolationLevel, ListOffset, ListOffsetsRequest, ScramMechanism, + create_topics_request::CreatableTopic, + delete_groups_response::DeletableGroupResult, + delete_records_request::DeleteRecordsTopic, + delete_records_response::DeleteRecordsTopicResult, + describe_cluster_response::DescribeClusterBroker, + describe_configs_response::DescribeConfigsResult, + describe_topic_partitions_response::DescribeTopicPartitionsResponseTopic, + fetch_response::AbortedTransaction, + incremental_alter_configs_request::AlterConfigsResource, + incremental_alter_configs_response::AlterConfigsResourceResponse, + list_groups_response::ListedGroup, + list_offsets_request::{ListOffsetsPartition, ListOffsetsTopic}, + list_offsets_response::ListOffsetsPartitionResponse, + record::deflated, + txn_offset_commit_response::TxnOffsetCommitResponseTopic, + }; + use rama::Service as _; + use tokio::time::{Duration, Instant, timeout}; + use url::Url; + use uuid::Uuid; + + use super::{LIST_OFFSETS_CONCURRENCY, ListOffsetsService}; + use crate::{ + BrokerRegistrationRequest, Error, GroupDetail, ListOffsetResponse, MetadataResponse, + NamedGroupDetail, OffsetCommitRequest, OffsetStage, ProducerIdResponse, Result, + ScramCredential, SemaphoreProxy, Storage, TopicId, Topition, TxnAddPartitionsRequest, + TxnAddPartitionsResponse, TxnOffsetCommitRequest, UpdateError, Version, + service::{ + RequestChannelService, bounded_channel, + deadline::{LIST_OFFSETS_READ_DEADLINE, Missed, within}, + }, + }; + + const TOPIC: &str = "abc"; + + /// Fails a test that would otherwise hang on a storage read that never + /// finishes, as every stalled read did before the deadline. + const HANG: Duration = Duration::from_secs(60); + + #[derive(Clone, Copy, Debug)] + enum Answer { + Offset(i64), + Stall, + Fail, + } + + /// Records the partition of a stalled storage read when nisshi drops it. + #[derive(Debug)] + struct Dropped { + partition: i32, + dropped: Arc>>, + } + + impl Drop for Dropped { + fn drop(&mut self) { + if let Ok(mut dropped) = self.dropped.lock() { + dropped.push(self.partition); + } + } + } + + /// Storage answering `list_offsets` per partition: an offset, a read that + /// never finishes, or an error. Records each call and each dropped read. + #[derive(Clone, Debug, Default)] + struct Stub { + partitions: BTreeMap, + calls: Arc>>, + dropped: Arc>>, + } + + impl Stub { + fn new(answers: &[Answer]) -> Self { + Self { + partitions: (0..).zip(answers.iter().copied()).collect(), + ..Self::default() + } + } + + fn calls(&self) -> Vec { + let mut calls = self.calls.lock().expect("calls").clone(); + calls.sort_unstable(); + calls + } + + fn dropped(&self) -> Vec { + let mut dropped = self.dropped.lock().expect("dropped").clone(); + dropped.sort_unstable(); + dropped + } + } + + #[async_trait] + impl Storage for Stub { + async fn fetch( + &self, + _topition: &Topition, + _offset: i64, + _min_bytes: u32, + _max_bytes: u32, + _isolation_level: IsolationLevel, + _max_wait: Duration, + ) -> Result> { + unimplemented!() + } + + async fn offset_stage(&self, _topition: &Topition) -> Result { + unimplemented!() + } + + async fn list_offsets( + &self, + _isolation_level: IsolationLevel, + offsets: &[(Topition, ListOffset)], + ) -> Result> { + let [(topition, _)] = offsets else { + return Err(Error::Message(format!( + "expected one partition per storage call, got {offsets:?}" + ))); + }; + + let partition = topition.partition(); + self.calls.lock()?.push(partition); + + match self.partitions.get(&partition) { + Some(Answer::Offset(offset)) => Ok(vec![( + topition.clone(), + ListOffsetResponse { + error_code: ErrorCode::None, + timestamp: None, + offset: Some(*offset), + }, + )]), + + Some(Answer::Stall) => { + let _guard = Dropped { + partition, + dropped: self.dropped.clone(), + }; + + pending().await + } + + Some(Answer::Fail) | None => Err(Error::Message(format!( + "storage failed for partition {partition}" + ))), + } + } + + async fn register_broker( + &self, + _broker_registration: BrokerRegistrationRequest, + ) -> Result<()> { + unimplemented!() + } + + async fn brokers(&self) -> Result> { + unimplemented!() + } + + async fn create_topic(&self, _topic: CreatableTopic, _validate_only: bool) -> Result { + unimplemented!() + } + + async fn delete_records( + &self, + _topics: &[DeleteRecordsTopic], + ) -> Result> { + unimplemented!() + } + + async fn delete_topic(&self, _topic: &TopicId) -> Result { + unimplemented!() + } + + async fn incremental_alter_resource( + &self, + _resource: AlterConfigsResource, + ) -> Result { + unimplemented!() + } + + async fn produce( + &self, + _transaction_id: Option<&str>, + _topition: &Topition, + _deflated: deflated::Batch, + ) -> Result { + unimplemented!() + } + + async fn offset_commit( + &self, + _group: &str, + _retention: Option, + _offsets: &[(Topition, OffsetCommitRequest)], + ) -> Result> { + unimplemented!() + } + + async fn committed_offset_topitions( + &self, + _group_id: &str, + ) -> Result> { + unimplemented!() + } + + async fn offset_fetch( + &self, + _group_id: Option<&str>, + _topics: &[Topition], + _require_stable: Option, + ) -> Result> { + unimplemented!() + } + + async fn metadata(&self, _topics: Option<&[TopicId]>) -> Result { + unimplemented!() + } + + async fn describe_config( + &self, + _name: &str, + _resource: ConfigResource, + _keys: Option<&[String]>, + ) -> Result { + unimplemented!() + } + + async fn describe_topic_partitions( + &self, + _topics: Option<&[TopicId]>, + _partition_limit: i32, + _cursor: Option, + ) -> Result> { + unimplemented!() + } + + async fn list_groups(&self, _states_filter: Option<&[String]>) -> Result> { + unimplemented!() + } + + async fn delete_groups( + &self, + _group_ids: Option<&[String]>, + ) -> Result> { + unimplemented!() + } + + async fn describe_groups( + &self, + _group_ids: Option<&[String]>, + _include_authorized_operations: bool, + ) -> Result> { + unimplemented!() + } + + async fn update_group( + &self, + _group_id: &str, + _detail: GroupDetail, + _version: Option, + ) -> Result> { + unimplemented!() + } + + async fn init_producer( + &self, + _transaction_id: Option<&str>, + _transaction_timeout_ms: i32, + _producer_id: Option, + _producer_epoch: Option, + ) -> Result { + unimplemented!() + } + + async fn txn_add_offsets( + &self, + _transaction_id: &str, + _producer_id: i64, + _producer_epoch: i16, + _group_id: &str, + ) -> Result { + unimplemented!() + } + + async fn txn_add_partitions( + &self, + _partitions: TxnAddPartitionsRequest, + ) -> Result { + unimplemented!() + } + + async fn txn_offset_commit( + &self, + _offsets: TxnOffsetCommitRequest, + ) -> Result> { + unimplemented!() + } + + async fn txn_end( + &self, + _transaction_id: &str, + _producer_id: i64, + _producer_epoch: i16, + _committed: bool, + ) -> Result { + unimplemented!() + } + + async fn maintain(&self, _now: SystemTime) -> Result<()> { + unimplemented!() + } + + async fn maintain_transactions(&self, _now: SystemTime) -> Result<()> { + unimplemented!() + } + + async fn aborted_transactions( + &self, + _topition: &Topition, + _offset: i64, + _last_stable_offset: i64, + ) -> Result> { + unimplemented!() + } + + async fn cluster_id(&self) -> Result { + unimplemented!() + } + + async fn node(&self) -> Result { + unimplemented!() + } + + async fn advertised_listener(&self) -> Result { + unimplemented!() + } + + async fn ping(&self) -> Result<()> { + unimplemented!() + } + + async fn delete_user_scram_credential( + &self, + _user: &str, + _mechanism: ScramMechanism, + ) -> Result<()> { + unimplemented!() + } + + async fn upsert_user_scram_credential( + &self, + _user: &str, + _mechanism: ScramMechanism, + _credential: ScramCredential, + ) -> Result<()> { + unimplemented!() + } + + async fn user_scram_credential( + &self, + _user: &str, + _mechanism: ScramMechanism, + ) -> Result> { + unimplemented!() + } + } + + /// Asks for the latest offset of `partitions` partitions of one topic. + async fn list_offsets( + storage: Stub, + partitions: i32, + ) -> Result> { + let request = ListOffsetsRequest::default() + .isolation_level(Some(IsolationLevel::ReadUncommitted.into())) + .replica_id(-1) + .topics(Some(vec![ + ListOffsetsTopic::default() + .name(TOPIC.into()) + .partitions(Some( + (0..partitions) + .map(|partition| { + ListOffsetsPartition::default() + .partition_index(partition) + .current_leader_epoch(Some(-1)) + .timestamp(ListOffset::Latest.try_into().expect("latest")) + }) + .collect(), + )), + ])); + + let response = timeout(HANG, ListOffsetsService { storage }.serve(request)) + .await + .map_err(|_| Error::Message("list offsets hung on a stalled read".into()))??; + + let topics = response.topics.unwrap_or_default(); + assert_eq!(1, topics.len()); + Ok(topics[0].partitions.clone().unwrap_or_default()) + } + + fn answers( + partitions: &[ListOffsetsPartitionResponse], + ) -> Result)>> { + partitions + .iter() + .map(|partition| { + ErrorCode::try_from(partition.error_code) + .map(|error_code| (partition.partition_index, error_code, partition.offset)) + .map_err(Into::into) + }) + .collect() + } + + /// A partition that is slow in storage answers `REQUEST_TIMED_OUT` at + /// the deadline, and the partitions either side of it get their offsets. + #[tokio::test(start_paused = true)] + async fn slow_partition_times_out_alone() -> Result<()> { + let storage = Stub::new(&[Answer::Offset(10), Answer::Stall, Answer::Offset(30)]); + + let started_at = Instant::now(); + let partitions = list_offsets(storage.clone(), 3).await?; + + assert_eq!(LIST_OFFSETS_READ_DEADLINE, started_at.elapsed()); + assert_eq!( + vec![ + (0, ErrorCode::None, Some(10)), + (1, ErrorCode::RequestTimedOut, Some(-1)), + (2, ErrorCode::None, Some(30)), + ], + answers(&partitions)? + ); + assert_eq!(vec![0, 1, 2], storage.calls()); + assert_eq!(vec![1], storage.dropped()); + + Ok(()) + } + + /// Without a slow partition, every partition is answered straight away, + /// in the order asked. + #[tokio::test(start_paused = true)] + async fn answers_in_request_order() -> Result<()> { + let offsets = [50, 40, 30, 20, 10, 0]; + let storage = Stub::new(&offsets.map(Answer::Offset)); + + let started_at = Instant::now(); + let partitions = list_offsets(storage.clone(), 6).await?; + + assert_eq!(Duration::ZERO, started_at.elapsed()); + assert_eq!( + (0..) + .zip(offsets) + .map(|(partition, offset)| (partition, ErrorCode::None, Some(offset))) + .collect::>(), + answers(&partitions)? + ); + + Ok(()) + } + + /// At most `LIST_OFFSETS_CONCURRENCY` reads run at once, and a read not + /// started by the deadline is never sent to storage. + #[tokio::test(start_paused = true)] + async fn bounded_and_not_started_after_the_deadline() -> Result<()> { + let partitions = i32::try_from(LIST_OFFSETS_CONCURRENCY + 2)?; + let storage = Stub::new(&vec![Answer::Stall; partitions as usize]); + + let started_at = Instant::now(); + let answered = list_offsets(storage.clone(), partitions).await?; + + assert_eq!(LIST_OFFSETS_READ_DEADLINE, started_at.elapsed()); + assert!( + answers(&answered)? + .iter() + .all(|(_, error_code, _)| *error_code == ErrorCode::RequestTimedOut) + ); + assert_eq!( + (0..i32::try_from(LIST_OFFSETS_CONCURRENCY)?).collect::>(), + storage.calls() + ); + assert_eq!(storage.calls(), storage.dropped()); + + Ok(()) + } + + /// A storage error still fails the whole request, dropping the reads + /// still in flight without reading any partition twice. + #[tokio::test(start_paused = true)] + async fn storage_error_fails_the_request() -> Result<()> { + let storage = Stub::new(&[Answer::Stall, Answer::Fail, Answer::Offset(30)]); + + let outcome = list_offsets(storage.clone(), 3).await; + + assert!( + matches!(outcome, Err(Error::Message(ref message)) if message.contains("partition 1")), + "{outcome:?}" + ); + assert_eq!(vec![0], storage.dropped()); + + let calls = storage.calls(); + assert!(calls.starts_with(&[0, 1]), "{calls:?}"); + assert!(calls.len() <= 3, "{calls:?}"); + + Ok(()) + } + + /// A read abandoned while it waits for the storage proxy's permit is a + /// queued miss, so a libSQL engine in `semaphore` mode whose permit + /// Produce and Fetch hold shows as queued, not as a slow engine. + #[tokio::test(start_paused = true)] + async fn read_waiting_for_the_proxy_permit_is_queued() -> Result<()> { + let proxy = SemaphoreProxy::new(Stub::new(&[Answer::Stall, Answer::Stall])); + + let permit_holder = { + let proxy = proxy.clone(); + + tokio::spawn(async move { + proxy + .list_offsets( + IsolationLevel::ReadUncommitted, + &[(Topition::new(TOPIC, 0), ListOffset::Latest)], + ) + .await + }) + }; + + // The stalled read holds the proxy's only permit once it has run. + tokio::time::sleep(Duration::from_millis(1)).await; + + let outcome = within( + "test", + Instant::now() + Duration::from_secs(1), + proxy.list_offsets( + IsolationLevel::ReadUncommitted, + &[(Topition::new(TOPIC, 1), ListOffset::Latest)], + ), + ) + .await; + assert!(matches!(outcome, Err(Missed::Queued)), "{outcome:?}"); + + permit_holder.abort(); + + Ok(()) + } + + /// A read abandoned while it waits for room on the storage request + /// channel is a queued miss, so a libSQL engine in `mpsc` mode behind + /// other requests shows as queued, not as a slow engine. + #[tokio::test(start_paused = true)] + async fn read_waiting_for_the_request_channel_is_queued() -> Result<()> { + let (tx, _rx) = bounded_channel(1); + let _only_slot = tx.reserve().await.map_err(|_| Error::UnableToSend)?; + + let service = RequestChannelService::new(tx.clone()); + + let outcome = within( + "test", + Instant::now() + Duration::from_secs(1), + service.list_offsets( + IsolationLevel::ReadUncommitted, + &[(Topition::new(TOPIC, 0), ListOffset::Latest)], + ), + ) + .await; + assert!(matches!(outcome, Err(Missed::Queued)), "{outcome:?}"); + + Ok(()) + } +}