diff --git a/nisshi-storage-dynostore/src/dynostore.rs b/nisshi-storage-dynostore/src/dynostore.rs index 0bf7049a..35d977fa 100644 --- a/nisshi-storage-dynostore/src/dynostore.rs +++ b/nisshi-storage-dynostore/src/dynostore.rs @@ -17,9 +17,10 @@ use std::{ collections::{BTreeMap, BTreeSet, btree_map::Entry}, fmt::{Debug, Display}, + future::ready, str::FromStr, sync::{Arc, Mutex}, - time::SystemTime, + time::{Instant, SystemTime}, }; use async_trait::async_trait; @@ -50,6 +51,7 @@ use nisshi_sans_io::{ list_groups_response::ListedGroup, metadata_response::{MetadataResponseBroker, MetadataResponsePartition, MetadataResponseTopic}, record::{Record, deflated, inflated}, + to_system_time, to_timestamp, txn_offset_commit_response::{TxnOffsetCommitResponsePartition, TxnOffsetCommitResponseTopic}, }; use nisshi_schema::{ @@ -76,13 +78,15 @@ use opticon::OptiCon; use percent_encoding::percent_decode_str; use rand::{prelude::*, rng}; use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use time_index::TimeIndex; use tokio::time::Duration; -use tracing::{debug, error, instrument, warn}; +use tracing::{debug, error, info, instrument, warn}; use url::Url; use uuid::Uuid; mod metadata; mod opticon; +mod time_index; #[cfg(test)] mod tests; @@ -123,6 +127,42 @@ const APPLICATION_JSON: &str = "application/json"; /// which is a different string to the sentinel. const EMPTY_GROUP_SENTINEL: &str = "%empty"; +/// How long after a batch object's write, on the store's clock, every +/// produce below it is taken to have written its own batch object or +/// failed. +/// +/// A produce assigns its offset in the watermark document before it writes +/// its batch object, so a listing can see a later batch and not an earlier +/// one that is still in flight. A gap below a batch written this long ago +/// is taken as a produce that failed in between, which holds no records, +/// and the backfill logs a warning for it. The broker does not enforce +/// this bound: the batch object's PUT gives up after the object store +/// client's retry timeout, three minutes by default, and a lake commit or +/// a transaction CAS between the two writes has no limit of its own. A +/// produce that overruns this margin is indexed past, and nothing rebuilds +/// a complete index. +/// +/// The age is the difference between two write times that the store +/// reports, the batch object's and the watermark document's, so the +/// broker's clock does not enter into it. +const BATCH_SETTLED_AFTER: Duration = Duration::from_secs(10 * 60); + +/// A backfill of one partition's time index in progress (see +/// [`DynoStore::begin_time_index_backfill`]). +#[derive(Clone, Copy, Debug)] +struct Backfill { + /// The token that [`TimeIndex::begin_backfill`] set, or found set. + token: u64, + /// The partition's log start when the backfill began. + low: i64, + /// The partition's high watermark when the backfill began: every batch + /// below it had assigned its offset before the backfill's listing. + high: i64, + /// When the backfill began, on the store's clock: the write time the + /// store reports for the watermark document after the begin CAS. + began: SystemTime, +} + /// Encode a Kafka group id as a single, opaque `object_store` path segment. /// /// Every character [`PathPart`] treats as reserved (including `/`) is percent @@ -183,6 +223,10 @@ pub(crate) struct DynoStore { schemas: Option, lake: Option, watermarks: Arc>>>, + /// One gate per partition, so that concurrent `ListOffsets(Timestamp)` + /// requests in this process share one time index backfill instead of + /// each reading every batch of the partition. + backfills: Arc>>>>, meta: OptiCon, object_store: Arc, @@ -435,7 +479,14 @@ struct TopicMetadata { struct Watermark { low: Option, high: Option, - timestamps: Option>, + /// The index that answers `ListOffsets(Timestamp)` for this partition. + /// A document written before the index existed lacks this key, and + /// decodes to an incomplete index that the first lookup backfills (see + /// `dynostore::tests::watermark_decodes_pre_time_index_documents`). + /// Those documents carry a `timestamps` key that was always `null`, and + /// this type ignores it. + #[serde(default)] + time_index: TimeIndex, } impl OptiCon { @@ -467,6 +518,7 @@ impl DynoStore { lake: None, watermarks: Arc::new(Mutex::new(BTreeMap::new())), + backfills: Arc::new(Mutex::new(BTreeMap::new())), meta: OptiCon::::new(cluster), object_store: Arc::new(Cache::new( Metron::new(object_store, cluster), @@ -558,6 +610,467 @@ impl DynoStore { .join(format!("{:0>10}.json", topition.partition))) } + fn records_prefix(&self, topition: &Topition) -> Path { + Path::from(format!( + "clusters/{}/topics/{}/partitions/{:0>10}/records/", + self.cluster, topition.topic, topition.partition, + )) + } + + fn batch_path(&self, topition: &Topition, base_offset: i64) -> Path { + Path::from(format!( + "clusters/{}/topics/{}/partitions/{:0>10}/records/{:0>20}.batch", + self.cluster, topition.topic, topition.partition, base_offset, + )) + } + + /// Returns the (cached, shared) `OptiCon` handle for `topition`'s + /// watermark document, inserting one if this is the first time this + /// partition has been touched in this process. + fn watermark_for(&self, topition: &Topition) -> Result> { + self.watermarks + .lock() + .map(|mut locked| { + locked + .entry(topition.to_owned()) + .or_insert_with(|| OptiCon::::new(self.cluster.as_str(), topition)) + .to_owned() + }) + .map_err(Into::into) + } + + /// Returns the gate that serializes time index backfills of `topition` + /// in this process. + fn backfill_gate_for(&self, topition: &Topition) -> Result>> { + self.backfills + .lock() + .map(|mut locked| locked.entry(topition.to_owned()).or_default().to_owned()) + .map_err(Into::into) + } + + /// Returns the base offset of the batch object at `meta`, or `None` for + /// an object under the records prefix that is not a batch. + fn batch_base_offset(meta: &ObjectMeta) -> Option { + meta.location.parts().next_back().and_then(|name| { + name.as_ref() + .get(0..20) + .and_then(|digits| i64::from_str(digits).ok()) + }) + } + + /// Streams the batch objects of `topition` at or after `start_offset`, + /// in offset order, as the base offset and write time of each. The + /// stream lists lazily, so a consumer that stops early lists only up to + /// where it stopped. + fn batch_stream( + &self, + topition: &Topition, + start_offset: i64, + ) -> BoxStream<'static, Result<(i64, SystemTime)>> { + let prefix = self.records_prefix(topition); + + // A batch key zero-pads its base offset, so key order is offset + // order. A store that supports a start-after LIST then begins at + // `start_offset`, instead of at the first batch of the partition. + let list_stream = if start_offset > 0 { + self.object_store + .list_with_offset(Some(&prefix), &self.batch_path(topition, start_offset - 1)) + } else { + self.object_store.list(Some(&prefix)) + }; + + let topition = topition.to_owned(); + + list_stream + .map_err(move |error| { + error!(?error, ?topition, start_offset, "listing batches"); + Error::Api(ErrorCode::KafkaStorageError) + }) + .try_filter_map(move |meta| { + ready(Ok(Self::batch_base_offset(&meta) + .filter(|&base_offset| base_offset >= start_offset) + .map(|base_offset| { + (base_offset, SystemTime::from(meta.last_modified)) + }))) + }) + .boxed() + } + + /// Lists the batch objects of `topition` with a base offset in + /// `start_offset..end_offset`: the base offset and write time of each. + async fn list_batches( + &self, + topition: &Topition, + start_offset: i64, + end_offset: i64, + ) -> Result> { + self.batch_stream(topition, start_offset) + .try_take_while(|&(base_offset, _)| ready(Ok(base_offset < end_offset))) + .try_collect() + .await + } + + /// Reads and decodes the batch object at `base_offset`, or returns + /// `Ok(None)` when the object does not exist. + /// + /// A missing object is a batch that was deleted after a LIST returned + /// it, so a caller can skip it. Every other failure is an error, because + /// a caller that skips a batch after a transient GET failure answers + /// from an incomplete log and gives no sign of it. The error is + /// `KAFKA_STORAGE_ERROR`, the code Kafka answers for a log it cannot + /// read, which a client retries; `UNKNOWN_SERVER_ERROR` would make it + /// give up. + async fn read_batch( + &self, + topition: &Topition, + base_offset: i64, + ) -> Result> { + let path = self.batch_path(topition, base_offset); + + let encoded = match self.object_store.get(&path).await { + Ok(get_result) => get_result.bytes().await, + Err(object_store::Error::NotFound { .. }) => { + warn!( + ?topition, + base_offset, "skipping a batch deleted after listing" + ); + return Ok(None); + } + Err(error) => Err(error), + }; + + encoded + .map_err(Error::from) + .and_then(|encoded| self.decode(encoded)) + .map(Some) + .inspect_err(|error| error!(?error, ?topition, base_offset, "reading batch")) + .map_err(|_| Error::Api(ErrorCode::KafkaStorageError)) + } + + /// Marks a backfill of `topition`'s time index in progress, with a CAS + /// write of its token (see [`TimeIndex::begin_backfill`]), and returns + /// the token with the partition's log start and high watermark at that + /// moment, and the store's write time of that CAS. The three steps of a + /// backfill are this, then [`DynoStore::collect_time_index_candidate`], + /// then [`DynoStore::commit_time_index_backfill`], so that a produce + /// landing between any two of them is never lost (see + /// [`TimeIndex::merge`]). + async fn begin_time_index_backfill(&self, topition: &Topition) -> Result { + let token = rng().random::(); + let watermark = self.watermark_for(topition)?; + + let (token, low, high) = watermark + .with_mut(&self.object_store, |w| { + Ok(( + w.time_index.begin_backfill(token), + w.low.unwrap_or(0), + w.high.unwrap_or(0), + )) + }) + .await?; + + // The store's write time of the document the CAS just wrote is the + // backfill's start on the store's clock, the clock that stamps every + // batch object's write time as well. + let began = self + .object_store + .head(watermark.location()) + .await + .map(|meta| SystemTime::from(meta.last_modified))?; + + Ok(Backfill { + token, + low, + high, + began, + }) + } + + /// Lists every batch object of `topition` below the high watermark of + /// `backfill` and replays the time index's append rule over their + /// headers, in offset order, into a candidate index. Returns that index + /// and the number of batches listed. Pure read: never writes anything. + /// + /// The candidate is complete when the listed batches cover every offset + /// from the log start to that high watermark. A gap in the listing is a + /// produce that assigned its offset and has not yet written its batch + /// object, or that failed in between; a batch written after the gap + /// that is older than [`BATCH_SETTLED_AFTER`] settles it as the latter. + /// The batches past a gap that is not settled get no entry, because an + /// entry past the batch in flight would start a scan after it, and the + /// candidate is left incomplete, so that the next lookup lists again. + /// Their timestamps still raise the candidate's maximum, so that a + /// lookup does not answer no match for a listed batch. An empty listing + /// is complete: a candidate without an entry cannot misplace a scan. + /// + /// Returns an error when a batch cannot be read, so that the caller + /// does not commit an index that misses that batch. + async fn collect_time_index_candidate( + &self, + topition: &Topition, + backfill: &Backfill, + ) -> Result<(TimeIndex, usize)> { + let listed = self.list_batches(topition, 0, backfill.high).await?; + let batches = listed.len(); + + info!(?topition, batches, "time index backfill: starting"); + + let mut candidate = TimeIndex::default(); + + // The offset after the last batch read, or `None` after a batch + // deleted between the listing and its read, whose length is unknown, + // so the gap check is off for the batch after it. + let mut next_offset = Some(backfill.low); + let mut settled = true; + + for (base_offset, written) in listed { + if settled + && let Some(expected) = next_offset.filter(|&expected| base_offset > expected) + { + // The batch's age when the backfill began, both times the + // store's own, so the broker's clock does not enter into it. + let age = backfill + .began + .duration_since(written) + .unwrap_or(Duration::ZERO); + + if age >= BATCH_SETTLED_AFTER { + warn!( + ?topition, + gap = ?(expected..base_offset), + age_secs = age.as_secs(), + "time index backfill: taking a gap below a settled batch as a failed produce" + ); + } else { + settled = false; + } + } + + let Some(batch) = self.read_batch(topition, base_offset).await? else { + next_offset = None; + continue; + }; + + if settled { + candidate.append( + batch.max_timestamp, + base_offset, + u64::try_from(batch.batch_length).unwrap_or_default(), + ); + } else { + candidate.observe(batch.max_timestamp); + } + + next_offset = Some(base_offset + i64::from(batch.last_offset_delta) + 1); + } + + // A gap after the last listed batch has no later batch to settle it, + // so it is in flight until a later batch is written and ages. A + // partition whose newest produce failed in between therefore lists + // again on each lookup until its next produce. An empty listing is + // taken as complete, so that a partition without batch objects (a + // lake sink's) does not list again on every lookup; the one case + // this misses is a partition's first produce from a binary without + // the index, still in flight during this listing. + if batches > 0 && next_offset.is_some_and(|expected| expected < backfill.high) { + settled = false; + } + + if settled { + candidate.mark_complete(); + } + + Ok((candidate, batches)) + } + + /// Commits a backfill `candidate` (from + /// [`DynoStore::collect_time_index_candidate`]) via a CAS write that + /// merges it with whatever is live at write time (see + /// [`TimeIndex::merge`]). Returns the merged index actually persisted, + /// which is complete unless the candidate was incomplete or the + /// backfill's token is gone. + async fn commit_time_index_backfill( + &self, + topition: &Topition, + candidate: &TimeIndex, + token: u64, + ) -> Result { + let watermark = self.watermark_for(topition)?; + + watermark + .with_mut(&self.object_store, |w| { + w.time_index = TimeIndex::merge(candidate.clone(), &w.time_index, token); + Ok(w.time_index.clone()) + }) + .await + } + + /// Returns this partition's time index and log start (`low`) when the + /// index is complete, and `None` when a backfill must run first. + /// + /// A partition without a watermark document has no batches, so its + /// empty index is complete. This read does not create the document, + /// because a read must not create state for a partition. + async fn complete_time_index(&self, topition: &Topition) -> Result> { + self.watermark_for(topition)? + .with(&self.object_store, |w| { + Ok((w.time_index.is_complete() || w.high.is_none()) + .then(|| (w.time_index.clone(), w.low.unwrap_or(0)))) + }) + .await + } + + /// Returns this partition's time index and log start (`low`) for a + /// lookup, running a backfill first if the index isn't complete yet (see + /// [`TimeIndex::is_complete`]). + /// + /// The backfill can leave the index incomplete, when its listing found a + /// batch in flight or when a binary without the index rewrote the + /// watermark document meanwhile (see + /// [`DynoStore::collect_time_index_candidate`] and [`TimeIndex::merge`]). + /// Every entry of that index still keeps the index's guarantee, and its + /// maximum timestamp covers every listed batch, so the lookup answers + /// from it, and the next lookup runs the backfill again. + async fn time_index_for_lookup(&self, topition: &Topition) -> Result<(TimeIndex, i64)> { + if let Some(found) = self.complete_time_index(topition).await? { + return Ok(found); + } + + let gate = self.backfill_gate_for(topition)?; + let _backfill = gate.lock().await; + + // Another request in this process can complete the backfill while + // this one waits for the gate. + if let Some(found) = self.complete_time_index(topition).await? { + return Ok(found); + } + + let started = Instant::now(); + let backfill = self.begin_time_index_backfill(topition).await?; + let (candidate, batches) = self + .collect_time_index_candidate(topition, &backfill) + .await?; + let merged = self + .commit_time_index_backfill(topition, &candidate, backfill.token) + .await?; + + info!( + ?topition, + batches, + entries = merged.entries().len(), + complete = merged.is_complete(), + elapsed = ?started.elapsed(), + "time index backfill: committed" + ); + + // The log start is read after the commit, not taken from the read + // before the backfill, so that a lookup never filters with a log + // start older than its index. Nothing moves the log start today; + // DeleteRecords will. + let low = self + .watermark_for(topition)? + .with(&self.object_store, |w| Ok(w.low.unwrap_or(0))) + .await?; + + Ok((merged, low)) + } + + /// Forward sequential scan of real batch objects in `topition`, starting + /// at `start_offset`, for the first logical record, one whose offset is + /// `>= low` (so a record `DeleteRecords` has logically removed is never + /// answered as a match, even if its batch object is still there) and + /// whose timestamp is `>= target`. The scan decodes the records of each + /// batch, because the answer is the first record at or after the target + /// and a batch header's `max_timestamp` cannot place it inside the + /// batch. + /// + /// Returns `Ok(None)` - not an error - when nothing qualifies: an empty + /// partition, a `start_offset` past every real batch (the lake-sink + /// produce path never writes a `.batch` object, see `DynoStore::produce`), + /// or every surviving record's timestamp below `target`. Returns an + /// error when a batch cannot be read, because the skipped batch can hold + /// the match. + async fn sequential_timestamp_scan( + &self, + topition: &Topition, + start_offset: i64, + target: i64, + low: i64, + ) -> Result> { + let mut batches = self.batch_stream(topition, start_offset); + + while let Some((base_offset, _)) = batches.try_next().await? { + let Some(deflated) = self.read_batch(topition, base_offset).await? else { + continue; + }; + + let inflated = inflated::Batch::try_from(&deflated) + .inspect_err(|error| error!(?error, ?topition, base_offset, "inflating batch")) + .map_err(|_| Error::Api(ErrorCode::KafkaStorageError))?; + + for record in &inflated.records { + let offset = base_offset + i64::from(record.offset_delta); + if offset < low { + continue; + } + + let timestamp = inflated.base_timestamp + record.timestamp_delta; + if timestamp >= target { + return Ok(Some((offset, timestamp))); + } + } + } + + Ok(None) + } + + /// Answers `ListOffsets(Timestamp)` for one partition: the first record + /// at or after `target`, below `last_stable` when the lookup is under + /// `READ_COMMITTED` and the partition has an open transaction, or no + /// match. Runs the time index backfill first when the index is + /// incomplete (see [`DynoStore::time_index_for_lookup`]). + async fn timestamp_lookup( + &self, + topition: &Topition, + target: i64, + last_stable: Option, + ) -> Result { + let (time_index, low) = self.time_index_for_lookup(topition).await?; + + // No record is at or after a target above the greatest timestamp of + // every batch, so there is no scan to run. + if target > time_index.max_timestamp() { + return Ok(ListOffsetResponse { + error_code: ErrorCode::None, + offset: None, + ..Default::default() + }); + } + + let start_offset = time_index.floor_offset(target); + + // Under READ_COMMITTED, Kafka answers only a record before the last + // stable offset, and answers no match otherwise + // (https://github.com/apache/kafka/blob/3.9.1/core/src/main/scala/kafka/cluster/Partition.scala#L1589). + let found = self + .sequential_timestamp_scan(topition, start_offset, target, low) + .await? + .filter(|(offset, _)| last_stable.is_none_or(|last_stable| *offset < last_stable)); + + Ok(match found { + Some((offset, timestamp)) => ListOffsetResponse { + error_code: ErrorCode::None, + offset: Some(offset), + timestamp: Some(to_system_time(timestamp)?), + }, + None => ListOffsetResponse { + error_code: ErrorCode::None, + offset: None, + ..Default::default() + }, + }) + } + async fn get(&self, location: &Path) -> Result<(V, Version)> where V: DeserializeOwned, @@ -736,17 +1249,13 @@ impl Storage for DynoStore { for partition in 0..topic.num_partitions { let topition = Topition::new(topic.name.as_str(), partition); - let watermark = self.watermarks.lock().map(|mut locked| { - locked - .entry(topition.to_owned()) - .or_insert(OptiCon::::new(self.cluster.as_str(), &topition)) - .to_owned() - })?; + let watermark = self.watermark_for(&topition)?; watermark .with_mut(&self.object_store, |watermark| { _ = watermark.high.take(); _ = watermark.low.take(); + watermark.time_index = TimeIndex::complete(); Ok(()) }) @@ -938,12 +1447,7 @@ impl Storage for DynoStore { .unwrap_or(false) { // Get watermark to calculate proper offset for lake sink - let watermark = self.watermarks.lock().map(|mut locked| { - locked - .entry(topition.to_owned()) - .or_insert_with(|| OptiCon::::new(self.cluster.as_str(), topition)) - .to_owned() - })?; + let watermark = self.watermark_for(topition)?; let offset = watermark .with_mut(&self.object_store, |watermark| { @@ -955,7 +1459,13 @@ impl Storage for DynoStore { |high| Some(high + deflated.last_offset_delta as i64 + 1i64), ); - watermark.timestamps = None; + // Deliberately not touched: this branch (tansu.lake.sink) + // never writes a `records/*.batch` object, so there is + // nothing for a `time_index` entry here to point at. A + // sequential scan starting from such an entry would list + // nothing and must handle that cleanly (see + // `DynoStore::sequential_timestamp_scan`), but the + // simpler fix is to just never create the entry. debug!(?watermark); @@ -1071,24 +1581,33 @@ impl Storage for DynoStore { } } - let watermark = self.watermarks.lock().map(|mut locked| { - locked - .entry(topition.to_owned()) - .or_insert_with(|| OptiCon::::new(self.cluster.as_str(), topition)) - .to_owned() - })?; + let watermark = self.watermark_for(topition)?; let offset = watermark .with_mut(&self.object_store, |watermark| { debug!(?watermark); let offset = watermark.high.unwrap_or_default(); + let is_new_partition = watermark.high.is_none(); + watermark.high = watermark.high.map_or_else( || Some(deflated.last_offset_delta as i64 + 1i64), |high| Some(high + deflated.last_offset_delta as i64 + 1i64), ); - watermark.timestamps = None; + if is_new_partition { + watermark.time_index.mark_complete(); + } + + // The index is appended to whether or not it is complete. + // A produce that waited for a concurrent backfill would + // leave a gap: the backfill reads the index before this + // entry exists, and merges nothing in for it. + watermark.time_index.append( + deflated.max_timestamp, + offset, + u64::try_from(deflated.batch_length).unwrap_or_default(), + ); debug!(?watermark); @@ -1353,12 +1872,7 @@ impl Storage for DynoStore { debug!(?stable); - let watermark = self.watermarks.lock().map(|mut locked| { - locked - .entry(topition.to_owned()) - .or_insert(OptiCon::::new(self.cluster.as_str(), topition)) - .to_owned() - })?; + let watermark = self.watermark_for(topition)?; watermark .with(&self.object_store, |watermark| { @@ -1423,6 +1937,32 @@ impl Storage for DynoStore { let mut responses = vec![]; for (topition, offset_request) in offsets { + if let ListOffset::Timestamp(system_time) = offset_request { + let target = to_timestamp(system_time)?; + + // A lookup that fails answers the error in its own + // partition, as Kafka does, so the other partitions of the + // request still get their answers. + let response = self + .timestamp_lookup(topition, target, stable.get(topition).copied()) + .await + .unwrap_or_else(|error| { + error!(?error, ?topition, target, "list offsets by timestamp"); + + ListOffsetResponse { + error_code: match error { + Error::Api(error_code) => error_code, + _ => ErrorCode::UnknownServerError, + }, + offset: None, + ..Default::default() + } + }); + + responses.push((topition.to_owned(), response)); + continue; + } + let location = Path::from(format!( "clusters/{}/topics/{}/partitions/{:0>10}/records", self.cluster, topition.topic, topition.partition, @@ -1481,14 +2021,6 @@ impl Storage for DynoStore { _ = candidate.replace(meta); } - ListOffset::Timestamp(system_time) - if SystemTime::from(meta.last_modified) > *system_time - && candidate.as_ref().is_none_or(|found| { - found.last_modified > meta.last_modified - }) => - { - _ = candidate.replace(meta); - } _ => continue, } } @@ -1516,18 +2048,11 @@ impl Storage for DynoStore { }, )) } else { - // No matching object: Earliest/Latest answer 0, a Timestamp - // lookup answers no offset (see `ListOffsetResponse::offset`). - let offset = match offset_request { - ListOffset::Earliest | ListOffset::Latest => Some(0), - ListOffset::Timestamp(_) => None, - }; - responses.push(( topition.to_owned(), ListOffsetResponse { error_code: ErrorCode::None, - offset, + offset: Some(0), ..Default::default() }, )) @@ -3114,6 +3639,16 @@ where self.object_store.list(prefix) } + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + debug!(?prefix, %offset); + + self.object_store.list_with_offset(prefix, offset) + } + async fn list_with_delimiter( &self, prefix: Option<&Path>, diff --git a/nisshi-storage-dynostore/src/dynostore/metadata.rs b/nisshi-storage-dynostore/src/dynostore/metadata.rs index 55710787..b743b80e 100644 --- a/nisshi-storage-dynostore/src/dynostore/metadata.rs +++ b/nisshi-storage-dynostore/src/dynostore/metadata.rs @@ -373,6 +373,17 @@ where self.object_store.list(prefix) } + #[instrument(skip_all, fields(prefix))] + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + debug!(?prefix, %offset); + REQUESTS.add(1, &[KeyValue::new("method", "list_with_offset")]); + self.object_store.list_with_offset(prefix, offset) + } + #[instrument(skip_all, fields(prefix))] async fn list_with_delimiter( &self, diff --git a/nisshi-storage-dynostore/src/dynostore/opticon.rs b/nisshi-storage-dynostore/src/dynostore/opticon.rs index 867bc90f..9db82b63 100644 --- a/nisshi-storage-dynostore/src/dynostore/opticon.rs +++ b/nisshi-storage-dynostore/src/dynostore/opticon.rs @@ -76,6 +76,11 @@ impl OptiCon { data_version: Default::default(), } } + + /// Returns the location of the document in the object store. + pub(super) fn location(&self) -> &Path { + &self.path + } } impl OptiCon diff --git a/nisshi-storage-dynostore/src/dynostore/tests.rs b/nisshi-storage-dynostore/src/dynostore/tests.rs index 91e70e29..11000f16 100644 --- a/nisshi-storage-dynostore/src/dynostore/tests.rs +++ b/nisshi-storage-dynostore/src/dynostore/tests.rs @@ -26,6 +26,7 @@ use tracing_subscriber::EnvFilter; use super::{DynoStore, EMPTY_GROUP_SENTINEL, decode_group_segment, group_path_part}; mod latency; +mod list_offsets; mod ping; mod stale_watermark; @@ -216,3 +217,54 @@ fn committed_offset_key_keeps_the_topic_segment_rule() -> Result<()> { } Ok(()) } + +/// Checked-in version of the serde-compatibility claim in +/// [`super::Watermark::time_index`]'s doc comment: every shape of +/// pre-time-index watermark document on disk decodes cleanly into the +/// current [`super::Watermark`], with an empty, incomplete time index. +#[test] +fn watermark_decodes_pre_time_index_documents() -> Result<()> { + use super::{Watermark, time_index::TimeIndex}; + + // Every watermark document written before the time index has + // `timestamps` explicitly `null`. + let null_timestamps: Watermark = + serde_json::from_str(r#"{"low":6,"high":66,"timestamps":null}"#)?; + assert_eq!(Some(6), null_timestamps.low); + assert_eq!(Some(66), null_timestamps.high); + assert_eq!(TimeIndex::default(), null_timestamps.time_index); + assert!(!null_timestamps.time_index.is_complete()); + + // Documents written before the `timestamps` key existed at all lack it + // entirely. + let key_omitted: Watermark = serde_json::from_str(r#"{"low":6,"high":66}"#)?; + assert_eq!(Some(6), key_omitted.low); + assert_eq!(Some(66), key_omitted.high); + assert_eq!(TimeIndex::default(), key_omitted.time_index); + + // The current code writes no `timestamps` key, and a binary without the + // time index, whose watermark has the shape of `Legacy`, reads the + // document it writes. + #[derive(Deserialize)] + struct Legacy { + low: Option, + high: Option, + timestamps: Option>, + } + + let written = serde_json::to_string(&Watermark { + low: Some(6), + high: Some(66), + time_index: TimeIndex::complete(), + })?; + assert!(!written.contains("timestamps"), "{written}"); + let round_trip: Watermark = serde_json::from_str(&written)?; + assert!(round_trip.time_index.is_complete()); + let legacy: Legacy = serde_json::from_str(&written)?; + assert_eq!( + (Some(6), Some(66), None), + (legacy.low, legacy.high, legacy.timestamps) + ); + + Ok(()) +} diff --git a/nisshi-storage-dynostore/src/dynostore/tests/list_offsets.rs b/nisshi-storage-dynostore/src/dynostore/tests/list_offsets.rs new file mode 100644 index 00000000..3631fcb4 --- /dev/null +++ b/nisshi-storage-dynostore/src/dynostore/tests/list_offsets.rs @@ -0,0 +1,1414 @@ +// 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. + +//! Tests for the `time_index`-backed `ListOffsets(Timestamp)` implementation. +//! Calls `DynoStore`'s private helpers directly (this module is +//! a descendant of `dynostore`, so it can see them) rather than only going +//! through the `Storage` trait, because some of the scenarios below need to +//! control the exact interleaving between a backfill's read and its write. + +use std::{ + collections::BTreeMap, + fmt::Display, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, + }, + time::{Duration, SystemTime}, +}; + +use crate::dynostore::{ + BATCH_SETTLED_AFTER, DynoStore, Txn, TxnDetail, TxnProduceOffset, Watermark, + tests::init_tracing, + time_index::{INITIAL_INTERVAL_BYTES, MAX_ENTRIES, TimeIndex}, +}; +use async_trait::async_trait; +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt, stream::BoxStream}; +use nisshi_sans_io::{ + ErrorCode, IsolationLevel, ListOffset, create_topics_request::CreatableTopic, to_system_time, + to_timestamp, +}; +use nisshi_storage::{Error, ListOffsetResponse, Result, Storage, Topition, TxnState}; +use object_store::{ + Attributes, CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, ObjectStoreExt, PutMode, PutMultipartOptions, PutOptions, PutPayload, PutResult, + memory::InMemory, path::Path, +}; + +/// An arbitrary fixed point in 2020, in Kafka's millisecond-since-epoch +/// timestamp representation. A real, far-from-`now` value keeps a test from +/// passing by coincidence the way a `SystemTime::now()`-based one could. +const T0: i64 = 1_577_836_800_000; // 2020-01-01T00:00:00Z + +fn storage() -> DynoStore { + DynoStore::new("nisshi", 111, InMemory::new()) +} + +async fn create_topic(storage: &DynoStore, topic: &str, num_partitions: i32) -> Result<()> { + _ = storage + .create_topic( + CreatableTopic::default() + .name(topic.into()) + .num_partitions(num_partitions) + .replication_factor(1) + .assignments(Some([].into())) + .configs(Some([].into())), + false, + ) + .await?; + + Ok(()) +} + +/// Builds a single-record batch with an explicit, fully-controlled +/// `base_timestamp`/`max_timestamp`, carrying `last_offset_delta` extra +/// records beyond the first so a caller can build a multi-record batch by +/// supplying more than one timestamp. +fn batch_with_timestamps(base_offset: i64, timestamps: &[i64]) -> Result { + deflated_batch::build(base_offset, timestamps, Bytes::from_static(b"v")) +} + +/// A batch that fills the index's initial interval on its own, so that it +/// gets an entry whenever it raises the maximum timestamp. +fn interval_batch(base_offset: i64, timestamp: i64) -> Result { + deflated_batch::build( + base_offset, + &[timestamp], + Bytes::from(vec![b'v'; usize::try_from(INITIAL_INTERVAL_BYTES)?]), + ) +} + +// Keeps the `deflated`/`inflated` plumbing out of the test bodies above. +mod deflated_batch { + use super::*; + use nisshi_sans_io::record::{Record, deflated, inflated}; + + pub(super) use deflated::Batch; + + pub(super) fn build(base_offset: i64, timestamps: &[i64], value: Bytes) -> Result { + let base_timestamp = timestamps[0]; + let max_timestamp = timestamps.iter().copied().max().unwrap_or(base_timestamp); + + let mut builder = inflated::Batch::builder() + .base_offset(base_offset) + .base_timestamp(base_timestamp) + .max_timestamp(max_timestamp) + .last_offset_delta(i32::try_from(timestamps.len() - 1).unwrap_or_default()); + + for (i, timestamp) in timestamps.iter().enumerate() { + builder = builder.record( + Record::builder() + .value(Some(value.clone())) + .timestamp_delta(timestamp - base_timestamp) + .offset_delta(i32::try_from(i).unwrap_or_default()), + ); + } + + builder + .build() + .and_then(Batch::try_from) + .map_err(Into::into) + } +} + +/// Produces one batch through the normal path (so `time_index` is maintained +/// the way a real produce maintains it), returning its base offset. +async fn produce(storage: &DynoStore, topition: &Topition, timestamps: &[i64]) -> Result { + let batch = batch_with_timestamps(0, timestamps)?; + storage.produce(None, topition, batch).await +} + +/// Writes a raw batch object directly to the object store, bypassing +/// `produce` entirely - simulating a batch written before this code (and its +/// `time_index` maintenance) existed. +async fn write_legacy_batch( + storage: &DynoStore, + topition: &Topition, + base_offset: i64, + timestamp: i64, +) -> Result<()> { + write_batch_object( + storage, + topition, + batch_with_timestamps(base_offset, &[timestamp])?, + ) + .await +} + +/// Writes `batch`'s object directly to the object store, at its own base +/// offset, as the last step of `produce` does. +async fn write_batch_object( + storage: &DynoStore, + topition: &Topition, + batch: deflated_batch::Batch, +) -> Result<()> { + let base_offset = batch.base_offset; + let path = storage.batch_path(topition, base_offset); + let payload = storage.encode(batch)?; + + _ = storage + .object_store + .put_opts( + &path, + payload, + PutOptions { + mode: PutMode::Create, + attributes: Attributes::new(), + ..Default::default() + }, + ) + .await + .map_err(Error::from)?; + + Ok(()) +} + +/// Directly sets a partition's watermark fields, bypassing both `produce` +/// and `create_topic`'s reset - used to set up pre-migration ("legacy") or +/// post-`DeleteRecords` (`low > 0`) states that nothing in the normal write +/// path produces on its own. +async fn seed_watermark( + storage: &DynoStore, + topition: &Topition, + f: impl Fn(&mut Watermark) -> Result<()>, +) -> Result<()> { + let watermark = storage.watermark_for(topition)?; + watermark.with_mut(&storage.object_store, f).await +} + +async fn list_offsets_timestamp( + storage: &DynoStore, + topition: &Topition, + target: i64, +) -> Result<(Option, Option)> { + list_offsets_timestamp_at(storage, IsolationLevel::ReadUncommitted, topition, target).await +} + +/// The offset and timestamp a lookup answers, after checking that it +/// answered without an error. +async fn list_offsets_timestamp_at( + storage: &DynoStore, + isolation_level: IsolationLevel, + topition: &Topition, + target: i64, +) -> Result<(Option, Option)> { + let responses = + list_offsets_timestamp_responses(storage, isolation_level, &[(topition, target)]).await?; + + assert_eq!(1, responses.len()); + assert_eq!(ErrorCode::None, responses[0].error_code); + + Ok(( + responses[0].offset, + responses[0] + .timestamp + .map(|t| to_timestamp(&t)) + .transpose()?, + )) +} + +/// One `ListOffsets` request with a timestamp `target` for each partition +/// of `lookups`, answered in the same order. +async fn list_offsets_timestamp_responses( + storage: &DynoStore, + isolation_level: IsolationLevel, + lookups: &[(&Topition, i64)], +) -> Result> { + let offsets = lookups + .iter() + .map(|&(topition, target)| { + to_system_time(target) + .map(|target| (topition.to_owned(), ListOffset::Timestamp(target))) + .map_err(Error::from) + }) + .collect::>>()?; + + let responses = storage.list_offsets(isolation_level, &offsets).await?; + + assert_eq!( + offsets + .iter() + .map(|(topition, _)| topition) + .collect::>(), + responses + .iter() + .map(|(topition, _)| topition) + .collect::>() + ); + + Ok(responses + .into_iter() + .map(|(_, response)| response) + .collect()) +} + +/// The error code a lookup answers for its partition. +async fn list_offsets_timestamp_error( + storage: &DynoStore, + topition: &Topition, + target: i64, +) -> Result { + let responses = list_offsets_timestamp_responses( + storage, + IsolationLevel::ReadUncommitted, + &[(topition, target)], + ) + .await?; + + assert_eq!(1, responses.len()); + assert_eq!(None, responses[0].offset); + + Ok(responses[0].error_code) +} + +/// Splits the backfill into its two steps and interleaves a concurrent +/// produce between them, to confirm a produce landing mid-backfill is never +/// lost: a write-side guard that defers inserting into `time_index` until +/// the index is complete would miss this entry (the produce never touches +/// `time_index`, so the backfill's commit has nothing to merge it with), +/// which is why `produce` carries no such guard. +#[tokio::test] +async fn backfill_commit_picks_up_concurrent_produce() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "legacy"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + // Two legacy batches, written directly (not through `produce`): off0 has + // the higher timestamp, off1 a lower one that the monotonic rule must + // not let override it. + write_legacy_batch(&storage, &topition, 0, T0).await?; + write_legacy_batch(&storage, &topition, 1, T0 - 50).await?; + + seed_watermark(&storage, &topition, |w| { + w.low = Some(0); + w.high = Some(2); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + + // Step 1: collect a candidate from the pre-concurrent-produce state. + let backfill = storage.begin_time_index_backfill(&topition).await?; + assert_eq!((0, 2), (backfill.low, backfill.high)); + let (candidate, batches) = storage + .collect_time_index_candidate(&topition, &backfill) + .await?; + assert_eq!(2, batches); + assert_eq!(&BTreeMap::from([(T0, 0)]), candidate.entries()); + assert_eq!(T0, candidate.max_timestamp()); + assert!(candidate.is_complete()); + + // Step 2: a concurrent produce lands before the backfill commits, with a + // new running max above everything seen so far. + let concurrent_timestamp = T0 + 50; + let offset = produce(&storage, &topition, &[concurrent_timestamp]).await?; + assert_eq!(2, offset); + + // Step 3: commit the now-stale candidate. + let merged = storage + .commit_time_index_backfill(&topition, &candidate, backfill.token) + .await?; + assert_eq!( + &BTreeMap::from([(T0, 0), (concurrent_timestamp, 2)]), + merged.entries() + ); + assert_eq!(concurrent_timestamp, merged.max_timestamp()); + assert!(merged.is_complete()); + + // Step 4: the concurrent produce's entry must still be answerable. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 25).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(concurrent_timestamp), timestamp); + + Ok(()) +} + +/// A plain `BTreeMap` union of a backfill candidate and whatever is live +/// would keep offset 2's entry below offset 0's already-established higher +/// one. A floor lookup between the two would then start its scan at offset +/// 2 and miss offset 0's record. The correct merge (replay the append rule +/// over the union, sorted by offset) drops that entry, so the lookup +/// answers offset 0. +#[tokio::test] +async fn merge_does_not_let_a_lower_concurrent_entry_outrank_an_established_one() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "legacy-union"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + write_legacy_batch(&storage, &topition, 0, T0).await?; // the eventual floor + write_legacy_batch(&storage, &topition, 1, T0 - 50).await?; + + seed_watermark(&storage, &topition, |w| { + w.low = Some(0); + w.high = Some(2); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + + // A post-upgrade produce lands before any read ever triggers a backfill: + // `time_index` is still `None` at this point, so this insert goes into a + // fresh map, exactly the scenario that makes a plain union wrong. + let offset = produce(&storage, &topition, &[T0 - 20]).await?; + assert_eq!(2, offset); + + // The first read now triggers the backfill. The target is above off2's + // T0 - 20, so an entry for off2 would be the floor, and a scan from + // there finds nothing. Off0's T0 is the first record at or after it. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 - 10).await?; + assert_eq!(Some(0), offset); + assert_eq!(Some(T0), timestamp); + + Ok(()) +} + +/// A `low > 0` straddling batch: one batch holds records at offsets 0, 1 and +/// 2, and `DeleteRecords` (simulated here by setting `low` directly, since +/// `delete_records` is still `todo!()` on this branch) has logically removed +/// offsets 0 and 1. A Timestamp lookup must skip the deleted prefix and only +/// match a record at or after `low`, even though the matching batch's own +/// base offset (0) is below `low`. +#[tokio::test] +async fn timestamp_lookup_skips_deleted_prefix_of_a_straddling_batch() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "straddling"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + // One batch, three records: offsets 0, 1, 2 with increasing timestamps. + let timestamps = [T0, T0 + 50, T0 + 100]; + let offset = produce(&storage, &topition, ×tamps).await?; + assert_eq!(0, offset); + + // Simulate DeleteRecords(cutoff = 2): offsets 0 and 1 are logically + // gone, and the index still points at the batch's base offset (0). + seed_watermark(&storage, &topition, |w| { + w.low = Some(2); + Ok(()) + }) + .await?; + + // Target below every surviving record's timestamp: must still answer the + // first surviving record (offset 2), not offset 0 or 1. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(T0 + 100), timestamp); + + Ok(()) +} + +/// The sequential scan helper must return `Ok(None)`, not error, when +/// `start_offset` is past every real batch - the case both an out-of-range +/// query and the lake-sink produce path (which never writes a `.batch` +/// object at all) hit. +#[tokio::test] +async fn sequential_scan_past_every_real_batch_returns_none() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "empty"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + _ = produce(&storage, &topition, &[T0]).await?; + + let found = storage + .sequential_timestamp_scan(&topition, 1_000_000, T0, 0) + .await?; + assert_eq!(None, found); + + Ok(()) +} + +/// A Timestamp lookup against a partition with no batches at all (never +/// produced to) must answer "no match", not error and not offset 0. +#[tokio::test] +async fn empty_partition_no_match() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "never-produced"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0).await?; + assert_eq!(None, offset); + assert_eq!(None, timestamp); + + Ok(()) +} + +/// A Timestamp lookup above every record's timestamp must answer "no +/// match", found directly from the index's greatest timestamp - no scan +/// required. +#[tokio::test] +async fn future_timestamp_no_match() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "future"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + _ = produce(&storage, &topition, &[T0]).await?; + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 1_000_000).await?; + assert_eq!(None, offset); + assert_eq!(None, timestamp); + + Ok(()) +} + +/// A brand-new partition's first produce must mark its index complete +/// immediately: nothing predates offset 0, so no backfill is ever needed for +/// it. +#[tokio::test] +async fn new_partition_first_produce_marks_index_complete() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "brand-new"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + _ = produce(&storage, &topition, &[T0]).await?; + + let watermark = storage.watermark_for(&topition)?; + let complete = watermark + .with(&storage.object_store, |w| Ok(w.time_index.is_complete())) + .await?; + assert!(complete); + + Ok(()) +} + +/// Pins the ticket's two remaining named scenarios, both inside one +/// multi-record batch: a target equal to a record's own timestamp (the +/// inclusive boundary) and a target that falls strictly between two records' +/// timestamps (a genuine mid-batch match). One batch, offsets 0/1/2 at +/// T0/T0+50/T0+100, is indexed under a single `time_index` entry keyed by its +/// `max_timestamp` (T0+100), so every assertion here is answered by the +/// sequential scan walking records inside that one batch, not by a floor +/// lookup landing on a different batch. +#[tokio::test] +async fn equal_and_mid_batch_timestamps_match_inside_a_batch() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "boundary"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + let timestamps = [T0, T0 + 50, T0 + 100]; + let offset = produce(&storage, &topition, ×tamps).await?; + assert_eq!(0, offset); + + // Inclusive boundary: target equals offset 1's own timestamp exactly. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 50).await?; + assert_eq!(Some(1), offset); + assert_eq!(Some(T0 + 50), timestamp); + + // Mid-batch match: target falls strictly between offset 0's and offset + // 1's timestamps, so the first record at or after it is offset 1. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 25).await?; + assert_eq!(Some(1), offset); + assert_eq!(Some(T0 + 50), timestamp); + + // Mid-batch match on the last record: target falls strictly between + // offset 1's and offset 2's timestamps. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 51).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(T0 + 100), timestamp); + + Ok(()) +} + +/// Each produce into a freshly created topic maintains the index under the +/// monotonic rule: a batch whose `max_timestamp` is below the running max +/// adds no entry. An entry for offset 1 here would be the floor of a lookup +/// between the two timestamps, and a scan from offset 1 would miss offset +/// 0, which holds the first record at or after the target. +#[tokio::test] +async fn produce_skips_an_entry_below_the_running_max() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "fresh-out-of-order"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + seed_dense_index(&storage, &topition).await?; + + assert_eq!(0, produce(&storage, &topition, &[T0]).await?); + assert_eq!(1, produce(&storage, &topition, &[T0 - 50]).await?); + + assert_eq!( + BTreeMap::from([(T0, 0)]), + time_index_entries(&storage, &topition).await? + ); + + for target in [T0 - 60, T0 - 40] { + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, target).await?; + assert_eq!(Some(0), offset, "target {target}"); + assert_eq!(Some(T0), timestamp, "target {target}"); + } + + Ok(()) +} + +/// Sets a fresh partition's index interval to one byte, so that every batch +/// that raises the maximum timestamp is indexed. +async fn seed_dense_index(storage: &DynoStore, topition: &Topition) -> Result<()> { + seed_index_interval(storage, topition, 1).await +} + +async fn seed_index_interval( + storage: &DynoStore, + topition: &Topition, + interval_bytes: u64, +) -> Result<()> { + seed_watermark(storage, topition, |w| { + w.time_index = TimeIndex::complete_with_interval(interval_bytes); + Ok(()) + }) + .await +} + +async fn time_index_entries( + storage: &DynoStore, + topition: &Topition, +) -> Result> { + storage + .watermark_for(topition)? + .with(&storage.object_store, |w| { + Ok(w.time_index.entries().clone()) + }) + .await +} + +/// The index is sparse: a batch that raises the maximum timestamp inside +/// the interval gets no entry. A lookup for a target between two entries +/// must scan from the greatest entry at or below the target (the floor), so +/// that it finds the unindexed batch. A scan from the entry above the +/// target would answer that entry's batch, a later offset than Kafka's. +#[tokio::test] +async fn lookup_scans_from_the_floor_entry_past_an_unindexed_batch() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "sparse"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + // Every batch here is one record of the same value, so they are all the + // same size. The interval admits an entry every third batch. + let batch_bytes = u64::try_from(batch_with_timestamps(0, &[T0])?.batch_length)?; + seed_index_interval(&storage, &topition, 2 * batch_bytes + 1).await?; + + for (expected, timestamp) in (0..).zip([T0, T0 + 100, T0 + 200, T0 + 300]) { + assert_eq!(expected, produce(&storage, &topition, &[timestamp]).await?); + } + + assert_eq!( + BTreeMap::from([(T0, 0), (T0 + 300, 3)]), + time_index_entries(&storage, &topition).await? + ); + + // Offset 2 holds the first record at or after the target, and has no + // entry. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 150).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(T0 + 200), timestamp); + + // A target equal to an entry's timestamp answers that entry's batch. + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 300).await?; + assert_eq!(Some(3), offset); + assert_eq!(Some(T0 + 300), timestamp); + + // A target between the maximum timestamp and the last entry's timestamp + // is a match; only a target above the maximum is no match. + assert_eq!(4, produce(&storage, &topition, &[T0 + 400]).await?); + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 350).await?; + assert_eq!(Some(4), offset); + assert_eq!(Some(T0 + 400), timestamp); + assert_eq!( + (None, None), + list_offsets_timestamp(&storage, &topition, T0 + 401).await? + ); + + Ok(()) +} + +/// The index holds at most `MAX_ENTRIES` entries however many batches the +/// partition has, and still answers each lookup with the first record at or +/// after the target once it has been thinned. +#[tokio::test] +async fn index_stays_within_the_cap_and_answers_after_thinning() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "capped"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + seed_dense_index(&storage, &topition).await?; + + let batches = i64::try_from(MAX_ENTRIES)? + 1; + for i in 0..batches { + assert_eq!(i, produce(&storage, &topition, &[T0 + i]).await?); + } + + let entries = time_index_entries(&storage, &topition).await?; + assert!(entries.len() <= MAX_ENTRIES, "{} entries", entries.len()); + assert!( + entries.len() >= MAX_ENTRIES / 2, + "{} entries", + entries.len() + ); + + for target in [T0, T0 + 1, T0 + 2, T0 + batches / 2, T0 + batches - 1] { + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, target).await?; + assert_eq!(Some(target - T0), offset, "target {target}"); + assert_eq!(Some(target), timestamp, "target {target}"); + } + + Ok(()) +} + +/// A lookup on a partition without a watermark document answers no match, +/// and does not create the document. +#[tokio::test] +async fn lookup_does_not_create_a_watermark_document() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topition = Topition::new("absent", 0); + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0).await?; + assert_eq!(None, offset); + assert_eq!(None, timestamp); + + let path = Path::from("clusters/nisshi/topics/absent/partitions/0000000000/watermark.json"); + assert!(matches!( + storage.object_store.head(&path).await, + Err(object_store::Error::NotFound { .. }) + )); + + Ok(()) +} + +/// Under READ_COMMITTED, a match at or after the first offset of an open +/// transaction is no match, as in Kafka. READ_UNCOMMITTED still answers it. +#[tokio::test] +async fn read_committed_lookup_stops_at_the_last_stable_offset() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "open-transaction"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + assert_eq!(0, produce(&storage, &topition, &[T0]).await?); + assert_eq!(1, produce(&storage, &topition, &[T0 + 50]).await?); + + storage + .meta + .with_mut(&storage.object_store, |meta| { + _ = meta.transactions.insert( + "txn".into(), + Txn { + producer: 1, + epochs: BTreeMap::from([( + 0, + TxnDetail { + state: Some(TxnState::Begin), + produces: BTreeMap::from([( + topic.into(), + BTreeMap::from([( + 0, + Some(TxnProduceOffset { + offset_start: 1, + offset_end: 2, + }), + )]), + )]), + ..Default::default() + }, + )]), + }, + ); + Ok(()) + }) + .await?; + + let committed = + list_offsets_timestamp_at(&storage, IsolationLevel::ReadCommitted, &topition, T0 + 25) + .await?; + assert_eq!((None, None), committed); + + let committed = + list_offsets_timestamp_at(&storage, IsolationLevel::ReadCommitted, &topition, T0).await?; + assert_eq!((Some(0), Some(T0)), committed); + + let uncommitted = list_offsets_timestamp(&storage, &topition, T0 + 25).await?; + assert_eq!((Some(1), Some(T0 + 50)), uncommitted); + + Ok(()) +} + +/// The error that [`Faulty`] returns for the next GET of its armed path. +#[derive(Clone, Copy, Debug)] +enum Fault { + NotFound, + Unavailable, +} + +/// An object store that fails the next GET of one armed path, counts the +/// GETs of batch objects, once aged lists every object as written longer +/// than [`BATCH_SETTLED_AFTER`] ago, and can report every write time from +/// a clock behind the test's own. Every GET first yields to the runtime, +/// so that concurrent requests in a test interleave. +#[derive(Clone, Debug, Default)] +struct Faulty { + inner: Arc, + armed: Arc>>, + batch_gets: Arc, + aged: Arc, + /// How far the store's clock is behind the test's, in seconds, in every + /// write time the store reports. + clock_behind_secs: Arc, + /// Conditional PUTs that failed because the object had changed. + put_conflicts: Arc, +} + +impl Faulty { + fn arm(&self, path: Path, fault: Fault) -> Result<()> { + self.armed + .lock() + .map(|mut armed| _ = armed.replace((path, fault))) + .map_err(Into::into) + } + + fn batch_gets(&self) -> usize { + self.batch_gets.load(Ordering::SeqCst) + } + + fn put_conflicts(&self) -> usize { + self.put_conflicts.load(Ordering::SeqCst) + } + + fn age_listing(&self) { + self.aged.store(true, Ordering::SeqCst); + } + + fn set_clock_behind(&self, behind: Duration) { + self.clock_behind_secs + .store(behind.as_secs(), Ordering::SeqCst); + } + + /// Shifts a write time the inner store reports onto this store's + /// clock. + fn on_store_clock(&self, meta: &mut ObjectMeta) { + let behind = Duration::from_secs(self.clock_behind_secs.load(Ordering::SeqCst)); + meta.last_modified = (SystemTime::from(meta.last_modified) - behind).into(); + } +} + +impl Display for Faulty { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Faulty").finish() + } +} + +#[async_trait] +impl ObjectStore for Faulty { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + self.inner + .put_opts(location, payload, opts) + .await + .inspect_err(|error| { + if matches!(error, object_store::Error::Precondition { .. }) { + _ = self.put_conflicts.fetch_add(1, Ordering::SeqCst); + } + }) + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result, object_store::Error> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> Result { + tokio::task::yield_now().await; + + if location.as_ref().ends_with(".batch") { + _ = self.batch_gets.fetch_add(1, Ordering::SeqCst); + } + + let fault = self + .armed + .lock() + .ok() + .and_then(|mut armed| armed.take_if(|(path, _)| path == location)) + .map(|(_, fault)| fault); + + let source = || "injected by test".into(); + + match fault { + Some(Fault::NotFound) => Err(object_store::Error::NotFound { + path: location.to_string(), + source: source(), + }), + Some(Fault::Unavailable) => Err(object_store::Error::Generic { + store: "Faulty", + source: source(), + }), + None => self + .inner + .get_opts(location, options) + .await + .map(|mut get_result| { + self.on_store_clock(&mut get_result.meta); + get_result + }), + } + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + self.inner.delete_stream(locations) + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> BoxStream<'static, Result> { + let aged = self.aged.load(Ordering::SeqCst); + let faulty = self.clone(); + + self.inner + .list(prefix) + .map_ok(move |mut meta| { + if aged { + meta.last_modified = (SystemTime::now() - 2 * BATCH_SETTLED_AFTER).into(); + } + faulty.on_store_clock(&mut meta); + meta + }) + .boxed() + } + + async fn list_with_delimiter( + &self, + prefix: Option<&Path>, + ) -> Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + opts: CopyOptions, + ) -> Result<(), object_store::Error> { + self.inner.copy_opts(from, to, opts).await + } +} + +/// Writes three legacy batches, with the running max at offset 1, into a +/// partition whose index is incomplete. +async fn legacy_partition_with_max_at_offset_1( + storage: &DynoStore, + topic: &str, +) -> Result { + let topition = Topition::new(topic, 0); + + create_topic(storage, topic, 1).await?; + + write_legacy_batch(storage, &topition, 0, T0).await?; + write_legacy_batch(storage, &topition, 1, T0 + 100).await?; + write_legacy_batch(storage, &topition, 2, T0 + 50).await?; + + seed_watermark(storage, &topition, |w| { + w.low = Some(0); + w.high = Some(3); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + + Ok(topition) +} + +async fn time_index_complete(storage: &DynoStore, topition: &Topition) -> Result { + storage + .watermark_for(topition)? + .with(&storage.object_store, |w| Ok(w.time_index.is_complete())) + .await +} + +/// A produce into a legacy partition before its first lookup indexes its +/// batch against an empty live index, whatever the legacy batches hold. +/// The backfill's candidate is sparse, so its maximum (offset 1 here) has +/// no entry of its own. The merge must still drop the live entry (offset +/// 3), which is below that maximum: kept, it would be the floor of a lookup +/// between the two, and a scan from offset 3 would miss offset 1. +#[tokio::test] +async fn backfill_merge_drops_a_live_entry_below_an_unindexed_legacy_batch() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topition = legacy_partition_with_max_at_offset_1(&storage, "legacy-sparse").await?; + + assert_eq!(3, produce(&storage, &topition, &[T0 + 70]).await?); + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 80).await?; + assert_eq!(Some(1), offset); + assert_eq!(Some(T0 + 100), timestamp); + + assert_eq!( + BTreeMap::from([(T0, 0)]), + time_index_entries(&storage, &topition).await? + ); + + Ok(()) +} + +/// A backfill that cannot read a batch answers an error for the partition +/// and leaves the index incomplete, so the next request runs the backfill +/// again. An index committed without offset 1 would answer no match for +/// this target. +#[tokio::test] +async fn backfill_get_failure_leaves_the_index_incomplete() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topition = legacy_partition_with_max_at_offset_1(&storage, "backfill-failure").await?; + + faulty.arm(storage.batch_path(&topition, 1), Fault::Unavailable)?; + + assert_eq!( + ErrorCode::KafkaStorageError, + list_offsets_timestamp_error(&storage, &topition, T0 + 60).await? + ); + assert!(!time_index_complete(&storage, &topition).await?); + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 60).await?; + assert_eq!(Some(1), offset); + assert_eq!(Some(T0 + 100), timestamp); + assert!(time_index_complete(&storage, &topition).await?); + + Ok(()) +} + +/// A backfill skips a batch that was deleted after the LIST, and completes +/// the index. +#[tokio::test] +async fn backfill_skips_a_batch_deleted_after_listing() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topition = legacy_partition_with_max_at_offset_1(&storage, "backfill-not-found").await?; + + faulty.arm(storage.batch_path(&topition, 2), Fault::NotFound)?; + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 60).await?; + assert_eq!(Some(1), offset); + assert_eq!(Some(T0 + 100), timestamp); + assert!(time_index_complete(&storage, &topition).await?); + + Ok(()) +} + +/// A scan that cannot read a batch answers an error for the partition, +/// instead of answering a later offset from the next batch. +#[tokio::test] +async fn scan_get_failure_answers_an_error_for_the_partition() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topic = "scan-failure"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + assert_eq!(0, produce(&storage, &topition, &[T0]).await?); + assert_eq!(1, produce(&storage, &topition, &[T0 + 50]).await?); + assert_eq!(2, produce(&storage, &topition, &[T0 + 100]).await?); + + faulty.arm(storage.batch_path(&topition, 1), Fault::Unavailable)?; + + assert_eq!( + ErrorCode::KafkaStorageError, + list_offsets_timestamp_error(&storage, &topition, T0 + 25).await? + ); + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 25).await?; + assert_eq!(Some(1), offset); + assert_eq!(Some(T0 + 50), timestamp); + + Ok(()) +} + +/// A lookup that fails answers the error in its own partition's +/// `error_code`, as Kafka does, and the other partitions of the same +/// request still get their answers. +#[tokio::test] +async fn failed_lookup_answers_only_its_own_partition_with_an_error() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topic = "partial-failure"; + let failing = Topition::new(topic, 0); + let healthy = Topition::new(topic, 1); + + create_topic(&storage, topic, 2).await?; + assert_eq!(0, produce(&storage, &failing, &[T0]).await?); + assert_eq!(0, produce(&storage, &healthy, &[T0]).await?); + + faulty.arm(storage.batch_path(&failing, 0), Fault::Unavailable)?; + + let responses = list_offsets_timestamp_responses( + &storage, + IsolationLevel::ReadUncommitted, + &[(&failing, T0), (&healthy, T0)], + ) + .await?; + + assert_eq!(ErrorCode::KafkaStorageError, responses[0].error_code); + assert_eq!(None, responses[0].offset); + assert_eq!(ErrorCode::None, responses[1].error_code); + assert_eq!(Some(0), responses[1].offset); + + Ok(()) +} + +/// Concurrent lookups on one partition in one process share one backfill: +/// the backfill reads each of the three batches once, and each lookup's +/// scan then reads the one batch that holds its match. +#[tokio::test] +async fn concurrent_lookups_share_one_backfill() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topition = legacy_partition_with_max_at_offset_1(&storage, "single-flight").await?; + + let before = faulty.batch_gets(); + + let (first, second) = tokio::join!( + list_offsets_timestamp(&storage, &topition, T0), + list_offsets_timestamp(&storage, &topition, T0), + ); + assert_eq!((Some(0), Some(T0)), first?); + assert_eq!((Some(0), Some(T0)), second?); + + assert_eq!(3 + 2, faulty.batch_gets() - before); + + Ok(()) +} + +/// `produce` assigns a batch's offset in the watermark document before it +/// writes the batch object, so a backfill's listing can see a later batch +/// and not an earlier one still in flight. Here offset 4 (T0 + 200) is in +/// flight, with no entry because it is inside the live index's interval, +/// and offset 5 (T0 + 150) is written. A candidate entry for offset 5 +/// would be the floor of a lookup for T0 + 180, and a scan from offset 5 +/// would miss offset 4 for good. The batches past the gap get no entry +/// instead, and the index is left incomplete, so the lookup after the +/// batch lands rebuilds it. +#[tokio::test] +async fn backfill_stops_its_listing_at_a_batch_in_flight() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topic = "in-flight"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + // The interval batch at offset 1 fills the candidate's interval, so + // offset 1 gets an entry. + write_legacy_batch(&storage, &topition, 0, T0).await?; + write_batch_object(&storage, &topition, interval_batch(1, T0 + 100)?).await?; + write_legacy_batch(&storage, &topition, 2, T0 + 50).await?; + + seed_watermark(&storage, &topition, |w| { + w.low = Some(0); + w.high = Some(3); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + + // A produce by this binary before any lookup: the live index's first + // entry. + assert_eq!(3, produce(&storage, &topition, &[T0 + 10]).await?); + + // Offset 4 has assigned its offset and raised the live maximum without + // an entry, and has not yet written its batch object. + let in_flight = batch_with_timestamps(4, &[T0 + 200])?; + let in_flight_bytes = u64::try_from(in_flight.batch_length)?; + seed_watermark(&storage, &topition, |w| { + w.high = Some(5); + w.time_index.append(T0 + 200, 4, in_flight_bytes); + Ok(()) + }) + .await?; + + // Offset 5 is written, and fills the candidate's interval since its + // entry at offset 1. + assert_eq!( + 5, + storage + .produce(None, &topition, interval_batch(0, T0 + 150)?) + .await? + ); + + // The batch in flight is not readable yet, so no match is the answer + // for now. The index stays incomplete, with no entry past the gap. + let (offset, _) = list_offsets_timestamp(&storage, &topition, T0 + 180).await?; + assert_eq!(None, offset); + assert!(!time_index_complete(&storage, &topition).await?); + assert_eq!( + BTreeMap::from([(T0, 0), (T0 + 100, 1)]), + time_index_entries(&storage, &topition).await? + ); + + write_batch_object(&storage, &topition, in_flight).await?; + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 180).await?; + assert_eq!(Some(4), offset); + assert_eq!(Some(T0 + 200), timestamp); + assert!(time_index_complete(&storage, &topition).await?); + + Ok(()) +} + +/// A gap below a batch written longer than `BATCH_SETTLED_AFTER` ago is a +/// produce that failed between assigning its offset and writing its batch +/// object, and holds no records. The listing continues past it and the +/// index completes. The same gap below a batch written just now is taken +/// as in flight, and the index stays incomplete. +#[tokio::test] +async fn backfill_completes_past_a_settled_gap() -> Result<()> { + let _guard = init_tracing()?; + + async fn partition_with_a_gap_at_offset_1(faulty: &Faulty) -> Result<(DynoStore, Topition)> { + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topic = "gap"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + write_legacy_batch(&storage, &topition, 0, T0).await?; + write_legacy_batch(&storage, &topition, 2, T0 + 100).await?; + + seed_watermark(&storage, &topition, |w| { + w.low = Some(0); + w.high = Some(3); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + + Ok((storage, topition)) + } + + let settled = Faulty::default(); + let (storage, topition) = partition_with_a_gap_at_offset_1(&settled).await?; + settled.age_listing(); + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 50).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(T0 + 100), timestamp); + assert!(time_index_complete(&storage, &topition).await?); + + let fresh = Faulty::default(); + let (storage, topition) = partition_with_a_gap_at_offset_1(&fresh).await?; + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 50).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(T0 + 100), timestamp); + assert!(!time_index_complete(&storage, &topition).await?); + + Ok(()) +} + +/// The age of a batch below a gap is measured between two write times of +/// the store, the batch's and the watermark document's, not against the +/// broker's clock. Here the store's clock is behind the broker's by twice +/// `BATCH_SETTLED_AFTER`, and both batches were written just now on the +/// store's clock: the gap between them is in flight, and the index stays +/// incomplete. An age taken from the broker's clock would settle it. +#[tokio::test] +async fn settled_gap_age_is_measured_on_the_store_clock() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + faulty.set_clock_behind(2 * BATCH_SETTLED_AFTER); + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topic = "store-clock"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + write_legacy_batch(&storage, &topition, 0, T0).await?; + write_legacy_batch(&storage, &topition, 2, T0 + 100).await?; + + seed_watermark(&storage, &topition, |w| { + w.low = Some(0); + w.high = Some(3); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 50).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(T0 + 100), timestamp); + assert!(!time_index_complete(&storage, &topition).await?); + + Ok(()) +} + +/// A gap after the last listed batch has no later batch to settle it, +/// however old the listed batches are, so the backfill leaves the index +/// incomplete and answers from the listed batches. The lookup after the +/// batch lands lists again and completes the index. +#[tokio::test] +async fn backfill_is_left_incomplete_by_a_gap_after_the_last_listed_batch() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + let storage = DynoStore::new("nisshi", 111, faulty.clone()); + let topic = "trailing-gap"; + let topition = Topition::new(topic, 0); + + create_topic(&storage, topic, 1).await?; + + write_legacy_batch(&storage, &topition, 0, T0).await?; + write_legacy_batch(&storage, &topition, 1, T0 + 100).await?; + + // Offset 2 has assigned its offset and not yet written its batch + // object. + seed_watermark(&storage, &topition, |w| { + w.low = Some(0); + w.high = Some(3); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + faulty.age_listing(); + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 50).await?; + assert_eq!(Some(1), offset); + assert_eq!(Some(T0 + 100), timestamp); + assert!(!time_index_complete(&storage, &topition).await?); + + write_legacy_batch(&storage, &topition, 2, T0 + 200).await?; + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 150).await?; + assert_eq!(Some(2), offset); + assert_eq!(Some(T0 + 200), timestamp); + assert!(time_index_complete(&storage, &topition).await?); + + Ok(()) +} + +/// A produce from another process between a backfill's listing and its +/// commit changes the watermark document under the committing process, so +/// the commit's conditional PUT fails. The commit then reads the document +/// again and merges against the index that the produce appended to, so +/// that produce's batch is not lost. +#[tokio::test] +async fn backfill_commit_merges_again_after_a_cas_conflict() -> Result<()> { + let _guard = init_tracing()?; + let faulty = Faulty::default(); + let committing = DynoStore::new("nisshi", 111, faulty.clone()); + let producing = DynoStore::new("nisshi", 111, faulty.clone()); + let topition = legacy_partition_with_max_at_offset_1(&committing, "cas-conflict").await?; + + let backfill = committing.begin_time_index_backfill(&topition).await?; + let (candidate, _) = committing + .collect_time_index_candidate(&topition, &backfill) + .await?; + assert!(candidate.is_complete()); + + assert_eq!(3, produce(&producing, &topition, &[T0 + 200]).await?); + + let conflicts = faulty.put_conflicts(); + let merged = committing + .commit_time_index_backfill(&topition, &candidate, backfill.token) + .await?; + assert_eq!(1, faulty.put_conflicts() - conflicts); + + assert!(merged.is_complete()); + assert_eq!(&BTreeMap::from([(T0, 0), (T0 + 200, 3)]), merged.entries()); + assert_eq!(T0 + 200, merged.max_timestamp()); + + let (offset, timestamp) = list_offsets_timestamp(&committing, &topition, T0 + 150).await?; + assert_eq!(Some(3), offset); + assert_eq!(Some(T0 + 200), timestamp); + + Ok(()) +} + +/// A binary without the index rewrites the watermark document without it, +/// and with the batch it produced in neither the backfill's listing nor +/// the live index. A commit that finds its token gone leaves the index +/// incomplete, so the next lookup backfills again and finds that batch. +#[tokio::test] +async fn backfill_is_left_incomplete_by_an_older_broker_write() -> Result<()> { + let _guard = init_tracing()?; + let storage = storage(); + let topition = legacy_partition_with_max_at_offset_1(&storage, "older-broker").await?; + + let backfill = storage.begin_time_index_backfill(&topition).await?; + let (candidate, _) = storage + .collect_time_index_candidate(&topition, &backfill) + .await?; + assert!(candidate.is_complete()); + + // The older broker's produce lands between the listing and the commit. + write_legacy_batch(&storage, &topition, 3, T0 + 200).await?; + seed_watermark(&storage, &topition, |w| { + w.high = Some(4); + w.time_index = TimeIndex::default(); + Ok(()) + }) + .await?; + + let merged = storage + .commit_time_index_backfill(&topition, &candidate, backfill.token) + .await?; + assert!(!merged.is_complete()); + + let (offset, timestamp) = list_offsets_timestamp(&storage, &topition, T0 + 150).await?; + assert_eq!(Some(3), offset); + assert_eq!(Some(T0 + 200), timestamp); + assert!(time_index_complete(&storage, &topition).await?); + + Ok(()) +} diff --git a/nisshi-storage-dynostore/src/dynostore/time_index.rs b/nisshi-storage-dynostore/src/dynostore/time_index.rs new file mode 100644 index 00000000..3ec44cc2 --- /dev/null +++ b/nisshi-storage-dynostore/src/dynostore/time_index.rs @@ -0,0 +1,469 @@ +// 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. + +//! The per-partition time index that answers `ListOffsets(Timestamp)`. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +/// Kafka's own `NO_TIMESTAMP` sentinel: a batch whose `max_timestamp` is at +/// or below this is never indexed. +pub(super) const NO_TIMESTAMP: i64 = -1; + +/// The most entries a [`TimeIndex`] holds. An append past this count drops +/// every other entry and doubles the interval, so a partition's watermark +/// document stays about this size however long the partition lives. +pub(super) const MAX_ENTRIES: usize = 1024; + +/// The initial interval, in batch bytes, between two entries: Apache Kafka's +/// default `index.interval.bytes`. +pub(super) const INITIAL_INTERVAL_BYTES: u64 = 4096; + +/// A sparse time index of one partition: `max_timestamp -> base_offset` of +/// some of its batches, in the partition's watermark document. +/// +/// Every entry keeps one guarantee: each batch before the entry's offset has +/// a `max_timestamp` below the entry's timestamp. A lookup takes the greatest +/// entry at or below its target (see [`TimeIndex::floor_offset`]) and scans +/// batches forward from that offset, as Kafka's +/// `LogSegment.findOffsetByTimestamp` does, because a record at or after +/// the target is never before that offset. Removing an entry keeps the +/// guarantee, so the index can be thinned, and pruned below a new log start, +/// without a lookup answering wrong; a longer scan is the only cost. +/// +/// An entry is appended under Kafka's `TimeIndex.maybeAppend` rule, only +/// when a batch's `max_timestamp` exceeds every one appended before it, and +/// only once `interval_bytes` of batches have been appended since the last +/// entry (see [`TimeIndex::append`]). The index therefore grows with the +/// partition's bytes and not with its batch count, and [`MAX_ENTRIES`] bounds +/// it from there. +#[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +pub(super) struct TimeIndex { + /// `max_timestamp -> base_offset`. + #[serde(default)] + entries: BTreeMap, + + /// The greatest `max_timestamp` of every batch appended so far, indexed + /// or not. A lookup above it has no match. + #[serde(default = "no_timestamp")] + max_timestamp: i64, + + /// Batch bytes appended since the last entry. + #[serde(default)] + bytes_since_entry: u64, + + /// Batch bytes between two entries. Doubles each time the index is + /// thinned. + #[serde(default = "initial_interval_bytes")] + interval_bytes: u64, + + /// Whether the index has seen every batch of the partition from offset 0 + /// onward. A watermark document written before the index existed lacks + /// this field, so it defaults to `false`, and the first lookup backfills + /// the index from the partition's batches. + #[serde(default)] + complete: bool, + + /// The token of the backfill in progress, set before its listing and + /// cleared by its commit. A binary without the index rewrites the + /// document without this field, so a commit after such a write finds + /// its token gone and leaves the index incomplete (see + /// [`TimeIndex::merge`]). + #[serde(default)] + backfill: Option, +} + +fn no_timestamp() -> i64 { + NO_TIMESTAMP +} + +fn initial_interval_bytes() -> u64 { + INITIAL_INTERVAL_BYTES +} + +impl Default for TimeIndex { + fn default() -> Self { + Self { + entries: BTreeMap::new(), + max_timestamp: NO_TIMESTAMP, + bytes_since_entry: 0, + interval_bytes: INITIAL_INTERVAL_BYTES, + complete: false, + backfill: None, + } + } +} + +impl TimeIndex { + /// Returns the empty index of a partition without a batch, which is + /// complete from offset 0 onward. + pub(super) fn complete() -> Self { + Self { + complete: true, + ..Self::default() + } + } + + /// Returns the complete, empty index of a new partition with an + /// `interval_bytes` other than [`INITIAL_INTERVAL_BYTES`], for a test + /// that needs a denser or a sparser index than the default. + #[cfg(test)] + pub(super) fn complete_with_interval(interval_bytes: u64) -> Self { + Self { + interval_bytes, + ..Self::complete() + } + } + + pub(super) fn is_complete(&self) -> bool { + self.complete + } + + /// Marks the index complete from offset 0 onward: for a partition whose + /// first batch is being appended, nothing predates it; for a backfill + /// candidate, its listing covered every batch below the high watermark. + pub(super) fn mark_complete(&mut self) { + self.complete = true; + } + + /// Marks a backfill in progress and returns its token: `token` when + /// none was in progress, else the token of the backfill another process + /// began, which then covers this one's listing as well. + pub(super) fn begin_backfill(&mut self, token: u64) -> u64 { + *self.backfill.get_or_insert(token) + } + + pub(super) fn entries(&self) -> &BTreeMap { + &self.entries + } + + pub(super) fn max_timestamp(&self) -> i64 { + self.max_timestamp + } + + #[cfg(test)] + pub(super) fn interval_bytes(&self) -> u64 { + self.interval_bytes + } + + /// Records a batch of `bytes` with `max_timestamp` appended at + /// `base_offset`, and indexes it when it raises the partition's maximum + /// timestamp and the interval since the last entry is full. + /// + /// Batches must be appended in offset order, because the guarantee of an + /// entry is over every batch before it. + pub(super) fn append(&mut self, max_timestamp: i64, base_offset: i64, bytes: u64) { + self.bytes_since_entry = self.bytes_since_entry.saturating_add(bytes); + + if max_timestamp <= self.max_timestamp { + return; + } + + self.max_timestamp = max_timestamp; + + if self.entries.is_empty() || self.bytes_since_entry >= self.interval_bytes { + _ = self.entries.insert(max_timestamp, base_offset); + self.bytes_since_entry = 0; + + if self.entries.len() > MAX_ENTRIES { + self.thin(); + } + } + } + + /// Raises the maximum timestamp for a batch of `max_timestamp` that gets + /// no entry, so that a lookup above every indexed timestamp and at or + /// below this one still scans for it. + pub(super) fn observe(&mut self, max_timestamp: i64) { + self.max_timestamp = self.max_timestamp.max(max_timestamp); + } + + /// Drops every other entry, keeping the newest, and doubles the interval + /// so the index fills again at half the rate. + fn thin(&mut self) { + let mut keep = self.entries.len().is_multiple_of(2); + + self.entries.retain(|_, _| { + keep = !keep; + keep + }); + + self.interval_bytes = self.interval_bytes.saturating_mul(2); + } + + /// Returns the offset to scan from for the first record at or after + /// `target`: the offset of the greatest entry whose timestamp is at or + /// below `target`, or offset 0 when `target` is below every entry. + pub(super) fn floor_offset(&self, target: i64) -> i64 { + self.entries + .range(..=target) + .next_back() + .map_or(0, |(_, &base_offset)| base_offset) + } + + /// Merges a `candidate` index built from the partition's listed batches + /// with the `live` index that produces have maintained meanwhile. The + /// result keeps the index's guarantee, and is complete only under the + /// conditions in the last paragraph. + /// + /// The merge is not a union of the two maps. The live index of a + /// partition that predates the index starts from no maximum timestamp, + /// so an entry it holds can be below a batch that the candidate listed, + /// indexed or not. A live entry is kept only when it is above the + /// candidate's maximum timestamp: the candidate listed every batch up + /// to its maximum, and the live index saw every batch after the + /// listing in offset order, because a produce appends to the index + /// whether or not it is complete. The entries of both are then replayed + /// in offset order under the append rule. Each entry of the result is + /// above every batch before it, so the result keeps the index's + /// guarantee. + /// + /// The result is complete only when the candidate is complete (its + /// listing covered every batch below the high watermark) and the live + /// index is already complete or still carries the backfill's `token`. + /// A binary without the index rewrites the document without the token + /// and without the batches it appended, so a commit that finds the + /// token gone leaves the index incomplete, and the next lookup + /// backfills again. + pub(super) fn merge(candidate: Self, live: &Self, token: u64) -> Self { + let candidate_max = candidate.max_timestamp; + + let mut by_offset: BTreeMap = candidate + .entries + .into_iter() + .map(|(timestamp, base_offset)| (base_offset, timestamp)) + .collect(); + + for (×tamp, &base_offset) in live.entries.range(candidate_max.saturating_add(1)..) { + _ = by_offset.insert(base_offset, timestamp); + } + + let mut merged = Self { + entries: BTreeMap::new(), + max_timestamp: NO_TIMESTAMP, + bytes_since_entry: 0, + interval_bytes: candidate.interval_bytes.max(live.interval_bytes), + complete: candidate.complete && (live.complete || live.backfill == Some(token)), + backfill: None, + }; + + for (base_offset, timestamp) in by_offset { + // Every entry of either side counts, whatever its interval, so + // the merge passes the interval's worth of bytes with each. + merged.append(timestamp, base_offset, merged.interval_bytes); + } + + merged.max_timestamp = merged + .max_timestamp + .max(candidate.max_timestamp) + .max(live.max_timestamp); + merged.bytes_since_entry = live.bytes_since_entry; + + merged + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const T0: i64 = 1_577_836_800_000; + + fn dense() -> TimeIndex { + TimeIndex { + interval_bytes: 1, + ..TimeIndex::complete() + } + } + + #[test] + fn append_skips_a_batch_below_the_maximum_timestamp() { + let mut index = dense(); + index.append(T0, 0, 1); + index.append(T0 - 50, 1, 1); + index.append(T0 + 50, 2, 1); + + assert_eq!(&BTreeMap::from([(T0, 0), (T0 + 50, 2)]), index.entries()); + assert_eq!(T0 + 50, index.max_timestamp()); + } + + #[test] + fn append_never_indexes_no_timestamp() { + let mut index = dense(); + index.append(NO_TIMESTAMP, 0, 1); + index.append(NO_TIMESTAMP - 1, 1, 1); + + assert!(index.entries().is_empty()); + assert_eq!(NO_TIMESTAMP, index.max_timestamp()); + } + + /// A batch that raises the maximum timestamp inside the interval raises + /// `max_timestamp` without an entry, so a later batch with a lower + /// timestamp is still not indexed. + #[test] + fn append_tracks_the_maximum_between_entries() { + let mut index = TimeIndex { + interval_bytes: 100, + ..TimeIndex::complete() + }; + index.append(T0, 0, 10); + index.append(T0 + 50, 1, 10); + index.append(T0 + 25, 2, 10); + index.append(T0 + 75, 3, 100); + + assert_eq!(&BTreeMap::from([(T0, 0), (T0 + 75, 3)]), index.entries()); + assert_eq!(T0 + 75, index.max_timestamp()); + } + + #[test] + fn floor_offset_takes_the_greatest_entry_at_or_below_the_target() { + let mut index = dense(); + index.append(100, 0, 1); + index.append(150, 2, 1); + index.append(200, 5, 1); + + assert_eq!(0, index.floor_offset(99)); + assert_eq!(0, index.floor_offset(100)); + assert_eq!(0, index.floor_offset(149)); + assert_eq!(2, index.floor_offset(150)); + assert_eq!(5, index.floor_offset(200)); + assert_eq!(5, index.floor_offset(10_000)); + } + + #[test] + fn thinning_keeps_the_index_within_the_cap_and_doubles_the_interval() { + let mut index = dense(); + let cap = i64::try_from(MAX_ENTRIES).expect("fits"); + + for i in 0..cap { + index.append(T0 + i, i, 1); + } + assert_eq!(MAX_ENTRIES, index.entries().len()); + assert_eq!(1, index.interval_bytes()); + + // The append past the cap thins the index, keeping the newest entry. + index.append(T0 + cap, cap, 1); + assert_eq!(MAX_ENTRIES / 2 + 1, index.entries().len()); + assert_eq!(2, index.interval_bytes()); + assert_eq!(Some((&(T0 + cap), &cap)), index.entries().last_key_value()); + + let count = cap * 3; + for i in cap + 1..count { + index.append(T0 + i, i, 1); + assert!(index.entries().len() <= MAX_ENTRIES, "after {i}"); + } + assert_eq!(T0 + count - 1, index.max_timestamp()); + + // Every surviving entry keeps its guarantee: the floor for a target + // is at or before the target's own batch. + for target in [T0, T0 + 1, T0 + count / 2, T0 + count - 1] { + assert!(index.floor_offset(target) <= target - T0); + } + } + + /// A live entry below a candidate entry of a lower offset is dropped, so + /// the merged index cannot start a scan past the earlier, matching batch. + #[test] + fn merge_drops_a_live_entry_below_an_earlier_candidate_entry() { + let mut candidate = dense(); + candidate.append(T0, 0, 1); + candidate.append(T0 - 50, 1, 1); + + let mut live = TimeIndex::default(); + assert_eq!(7, live.begin_backfill(7)); + live.append(T0 - 20, 2, 1); + + let merged = TimeIndex::merge(candidate, &live, 7); + + assert!(merged.is_complete()); + assert_eq!(&BTreeMap::from([(T0, 0)]), merged.entries()); + assert_eq!(T0, merged.max_timestamp()); + } + + /// The merged index is complete only when the candidate's listing was + /// complete and the live index still carries the backfill's token. A + /// live index rewritten without the token (by a binary without the + /// index), or under another backfill's token, leaves it incomplete. A + /// live index that is already complete stays complete. + #[test] + fn merge_completes_only_a_complete_candidate_under_its_own_backfill() { + let candidate = dense(); + + let mut live = TimeIndex::default(); + assert_eq!(7, live.begin_backfill(7)); + assert_eq!( + 7, + live.begin_backfill(8), + "a second backfill joins the first" + ); + assert!(TimeIndex::merge(candidate.clone(), &live, 7).is_complete()); + assert!(!TimeIndex::merge(candidate.clone(), &live, 8).is_complete()); + assert!(!TimeIndex::merge(candidate.clone(), &TimeIndex::default(), 7).is_complete()); + assert!(TimeIndex::merge(candidate.clone(), &TimeIndex::complete(), 7).is_complete()); + + let incomplete = TimeIndex::default(); + assert!(!TimeIndex::merge(incomplete, &live, 7).is_complete()); + + let merged = TimeIndex::merge(candidate, &live, 7); + assert_eq!(None, merged.backfill, "the commit clears the token"); + } + + /// The candidate's maximum can belong to a batch inside the interval, + /// with no entry. A live entry below that maximum is still below a + /// listed batch, so it is dropped. + #[test] + fn merge_drops_a_live_entry_below_the_candidate_maximum() { + let mut candidate = TimeIndex { + interval_bytes: 100, + ..TimeIndex::complete() + }; + candidate.append(T0, 0, 10); + candidate.append(T0 + 100, 1, 10); + candidate.append(T0 + 50, 2, 10); + assert_eq!(&BTreeMap::from([(T0, 0)]), candidate.entries()); + + let mut live = TimeIndex::default(); + live.append(T0 + 70, 3, 10); + + let merged = TimeIndex::merge(candidate, &live, 7); + + assert_eq!(&BTreeMap::from([(T0, 0)]), merged.entries()); + assert_eq!(T0 + 100, merged.max_timestamp()); + } + + #[test] + fn merge_keeps_a_live_entry_above_the_candidate() { + let mut candidate = dense(); + candidate.append(T0, 0, 1); + + let mut live = TimeIndex::default(); + live.append(T0 + 50, 2, 1); + + let merged = TimeIndex::merge(candidate, &live, 7); + + assert_eq!(&BTreeMap::from([(T0, 0), (T0 + 50, 2)]), merged.entries()); + assert_eq!(T0 + 50, merged.max_timestamp()); + } + + /// A watermark document written before the index existed decodes into + /// an incomplete, empty index. + #[test] + fn decodes_an_empty_document() -> Result<(), serde_json::Error> { + let index: TimeIndex = serde_json::from_str("{}")?; + assert_eq!(TimeIndex::default(), index); + assert!(!index.is_complete()); + Ok(()) + } +}