diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 91991cd..1b7eca6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -98,7 +98,7 @@ jobs: sleep 2 done docker compose exec -T postgres psql -U postgres -tAc \ - "SELECT extversion FROM pg_extension WHERE extname='kafgres'" | grep -q 0.3.2 + "SELECT extversion FROM pg_extension WHERE extname='kafgres'" | grep -q 0.3.3 docker compose exec -T postgres psql -U postgres -tAc \ "SELECT count(*) FROM pg_stat_activity WHERE backend_type='kafgres_broker'" | grep -q '^1$' diff --git a/Cargo.lock b/Cargo.lock index 8215562..05a93bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -56,7 +56,7 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "kafgres-codec" -version = "0.3.2" +version = "0.3.3" dependencies = [ "bytes", "crc32c", @@ -68,7 +68,7 @@ dependencies = [ [[package]] name = "kafgres-codec-gen" -version = "0.3.2" +version = "0.3.3" dependencies = [ "serde", "serde_json", diff --git a/README.md b/README.md index 1c9732d..c4622b7 100644 --- a/README.md +++ b/README.md @@ -189,6 +189,7 @@ SELECT kafgres_add_mapping( FROM customers c WHERE c.id = new.customer_id)) $$, key_expr => $$ new.id::text $$, filter_expr => $$ new.status = 'placed' $$); +-- SELECT kafgres_slot_ready(); -- optional: blocks until the CDC slot is ready after the first mapping is added ``` A background worker drains a logical replication slot fed by kafgres's own output plugin, diff --git a/codec-gen/Cargo.toml b/codec-gen/Cargo.toml index 05d1de2..9d51e64 100644 --- a/codec-gen/Cargo.toml +++ b/codec-gen/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "kafgres-codec-gen" -version = "0.3.2" +version = "0.3.3" edition = "2021" license = "Elastic-2.0" description = "Emits kafgres-codec from the vendored Kafka message schemas." diff --git a/codec/Cargo.toml b/codec/Cargo.toml index 31b60d9..8d77fcc 100644 --- a/codec/Cargo.toml +++ b/codec/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "kafgres-codec" -version = "0.3.2" +version = "0.3.3" edition = "2021" license = "Elastic-2.0" description = "Kafka wire protocol codec, generated from vendored schemas. No pgrx, no Postgres." diff --git a/docs/architecture.md b/docs/architecture.md index 30b4a78..c815c83 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -31,6 +31,18 @@ out of band (see below) and is always asynchronous, since `synchronous_standby_n does not cover files outside the WAL; a failover can lose the newest records, and the leader epochs described next make that visible to consumers rather than silent. +On the primary, `kafgres_produce()` still commits the row and the event together. The +loss is in replication: its row and commit marker travel in the WAL, which can be +synchronous, but the record bytes travel with the segment follower. After a failover the new +primary drops the markers whose records never arrived, with a logged warning, so a row can +survive without its event. Synchronous replication does not close this. +`kafgres_follower_lag()` on a standby reports how many records the follower was behind at +its last Fetch, which `pg_stat_replication` does not cover. It returns no row while the +follower cannot reach the leader, or while it is too far behind for one Fetch to cover +every partition, so alert on a missing row as well as on a high one. A follower that stops +or hangs keeps its last reading, so alert too when `observed_at_ms` grows old. The table +engine keeps the bytes in the WAL and has no such gap. + ## Primary-only operation The broker, CDC and archiver workers run on the primary only. They are registered with diff --git a/docs/producing.md b/docs/producing.md index ed64c33..e11d3c7 100644 --- a/docs/producing.md +++ b/docs/producing.md @@ -50,11 +50,15 @@ computed from committed markers. Three properties of that mechanism are worth kn client that re-requests the same offset forever looks like a hung consumer. - Each transaction produces under its own producer id, taken from the transaction's xid. Kafka clients drop everything from a producer id at or after an aborted first offset, - so a shared producer id would let one rollback discard later committed records. For - the same reason, when a savepoint that produced is rolled back, the rest of the - transaction produces under a new producer id derived from the same xid. A transaction - that rolls back more than about four million such savepoints gets an error from - further calls. + so a shared producer id would let one rollback discard later committed records. The id + also carries a flag that keeps it apart from the ids Kafka clients are assigned, and the + number of the server start that produced it, so a transaction id that Postgres hands out + again after a crash does not share an id with the batch an earlier, lost transaction + left in the log. For the same reason, when a savepoint that produced is rolled back, the + rest of the transaction produces under a new producer id derived from the same xid. A + transaction that rolls back more than 1023 such savepoints (for example a loop of + `EXCEPTION` blocks that produce and fail) gets an error from further calls. Commit in + smaller batches instead of rolling back that many times in one transaction. - The last stable offset is the first offset of the oldest transaction still in flight on the partition, so `read_committed` consumers never see records that may yet commit or abort. A commit becomes visible once every transaction that produced to the partition @@ -84,8 +88,25 @@ creation and deletion, ACLs, CDC mappings and slots, retention, archiving, resto checks, replication), so each needs an explicit `GRANT EXECUTE`. The exceptions are `kafgres_produce()` itself and the read-only reports: `kafgres_version()`, `kafgres_kafka_version()`, `kafgres_partition_offsets()`, `kafgres_archive_status()`, -`kafgres_cdc_status()`, `kafgres_cdc_snapshots()` and `kafgres_share_state()`. On an -install upgraded from 0.2.0 this takes effect at `ALTER EXTENSION kafgres UPDATE`. +`kafgres_cdc_status()`, `kafgres_cdc_snapshots()`, `kafgres_share_state()` and +`kafgres_follower_lag()`. On an install upgraded from 0.2.0 this takes effect at +`ALTER EXTENSION kafgres UPDATE`. + +`kafgres_add_mapping` runs with the caller's privileges: it reads `kafgres_topics` and the +source table, and writes `kafgres_cdc_mappings`, so a role needs those grants as well as +`EXECUTE`. It also creates the CDC slot, which Postgres allows only to a superuser or a role +with the `REPLICATION` attribute. A role without that still saves its mapping, with a +warning, and the CDC worker creates the slot (see [CDC mappings](#cdc-mappings)). Use +`ALTER ROLE ... REPLICATION` for a role that should create it. `kafgres_slot_ready()` and +`kafgres_cdc_create_slot()` are administrative like the rest. + +```sql +GRANT EXECUTE ON FUNCTION kafgres_add_mapping(text, text, text, text, text, text) TO cdc_admin; +GRANT SELECT ON kafgres_topics TO cdc_admin; +GRANT SELECT, INSERT, UPDATE ON kafgres_cdc_mappings TO cdc_admin; +GRANT SELECT ON public.orders TO cdc_admin; -- each source table it maps +ALTER ROLE cdc_admin REPLICATION; -- optional: lets kafgres_add_mapping create the slot itself +``` ### Engines @@ -98,7 +119,29 @@ caller holds no lock that other producers need. ## CDC mappings -A mapping copies changes from a table to a topic through the WAL. You register it once: +A mapping copies changes from a table to a topic through the WAL. You register it once. +`kafgres_add_mapping` creates the logical replication slot if it does not exist yet. +Capture starts when the mapping commits: changes committed after that are captured, and +changes other sessions commit while the adding transaction is still open are not, so keep +that transaction short. A slot sees only WAL written after it exists, so call it before the +transaction has written anything, or in a transaction of its own. Creating a slot waits for +every transaction that holds an xid to finish, so the call gives the ones running a second +to finish and usually does not wait on any that take longer (one that takes an xid while the +slot is being built can still hold the call). If the slot cannot be created (such a +transaction is running, the transaction has already written, the role lacks `REPLICATION` +and is not a superuser ([Who may produce](#who-may-produce)), or `wal_level` is not +`logical`), the mapping is saved with a warning and the worker creates the slot later; +changes before then are not captured. If every slot `max_replication_slots` allows is in +use, the call fails and the mapping is not saved. `kafgres_slot_ready()` waits until the slot is ready and returns `slot ready`, or +`slot not ready` if its `timeout_ms` (120 s by default) passes first; a transaction left open +that long keeps the slot from being ready, and the answer says so rather than guessing. In +the transaction that added the mapping it answers at once: `slot ready` if the slot is +already ready, and otherwise `slot not ready`, since the slot cannot be ready before that +transaction ends; then call it again after the commit. Either way, capture still starts at +the commit. `kafgres_cdc_create_slot()` creates the slot and waits, if an enabled, committed +mapping wants one; with none it creates nothing and returns false. The slot stays after the +last `kafgres_remove_mapping`, or a rolled-back `kafgres_add_mapping`; the worker keeps +advancing it, so it holds no WAL while the worker drains. `kafgres_cdc_drop_slot()` drops it. The examples assume a `public.orders (id int, customer_id int, total numeric, status text)` table and a `customers (id int, name text, tier text)` lookup table. @@ -119,6 +162,7 @@ SELECT kafgres_add_mapping( key_expr => $$ new.id::text $$, filter_expr => $$ new.status = 'placed' AND old.status IS DISTINCT FROM 'placed' $$ ); +-- SELECT kafgres_slot_ready(); -- optional: blocks until the CDC slot is ready after the first mapping is added ``` A background worker drains a logical replication slot fed by kafgres's own output @@ -268,6 +312,7 @@ SELECT kafgres_add_mapping('txns', 'kafgres.transaction', 'my-transactions', 'event_count', event_count, 'data_collections', data_collections)$$, $$xid::text$$, NULL); +-- SELECT kafgres_slot_ready(); -- optional: blocks until the CDC slot is ready after the first mapping is added ``` `xid` on each change lets a consumer group the changes of one commit; the summary tells diff --git a/extension/Cargo.lock b/extension/Cargo.lock index a82fdbd..757d6f4 100644 --- a/extension/Cargo.lock +++ b/extension/Cargo.lock @@ -854,7 +854,7 @@ dependencies = [ [[package]] name = "kafgres" -version = "0.3.2" +version = "0.3.3" dependencies = [ "base64", "kafgres-codec", @@ -873,7 +873,7 @@ dependencies = [ [[package]] name = "kafgres-codec" -version = "0.3.2" +version = "0.3.3" dependencies = [ "bytes", "crc32c", diff --git a/extension/Cargo.toml b/extension/Cargo.toml index 90ffb81..e66ef25 100644 --- a/extension/Cargo.toml +++ b/extension/Cargo.toml @@ -5,7 +5,7 @@ [package] name = "kafgres" -version = "0.3.2" +version = "0.3.3" edition = "2021" license = "Elastic-2.0" diff --git a/extension/kafgres.control b/extension/kafgres.control index 4800f2b..a22c161 100644 --- a/extension/kafgres.control +++ b/extension/kafgres.control @@ -1,5 +1,5 @@ comment = 'kafgres: a Kafka broker embedded in PostgreSQL' -default_version = '0.3.2' +default_version = '0.3.3' module_pathname = '$libdir/kafgres' relocatable = false superuser = true diff --git a/extension/sql/kafgres--0.3.2--0.3.3.sql b/extension/sql/kafgres--0.3.2--0.3.3.sql new file mode 100644 index 0000000..4a9401a --- /dev/null +++ b/extension/sql/kafgres--0.3.2--0.3.3.sql @@ -0,0 +1,21 @@ +-- kafgres 0.3.2 -> 0.3.3. +-- +-- kafgres_follower_lag() reports how far the follower is behind the leader. It stays +-- executable by PUBLIC, like the other read-only reports. +CREATE FUNCTION @extschema@."kafgres_follower_lag"() RETURNS TABLE ( + "lag_records" bigint, + "observed_at_ms" bigint +) +STRICT +LANGUAGE c +AS 'MODULE_PATHNAME', 'kafgres_follower_lag_wrapper'; + +-- kafgres_slot_ready() waits for the CDC slot; like the other administrative functions it +-- is not executable by PUBLIC. +CREATE FUNCTION @extschema@."kafgres_slot_ready"( + "timeout_ms" INT DEFAULT 120000 +) RETURNS text +STRICT +LANGUAGE c +AS 'MODULE_PATHNAME', 'kafgres_slot_ready_wrapper'; +REVOKE EXECUTE ON FUNCTION @extschema@."kafgres_slot_ready"(int) FROM PUBLIC; diff --git a/extension/src/cdc.rs b/extension/src/cdc.rs index 40dc689..18e0a58 100644 --- a/extension/src/cdc.rs +++ b/extension/src/cdc.rs @@ -42,12 +42,38 @@ fn kafgres_add_mapping( if !valid_source(source_table) { error!("kafgres: {source_table:?} is not a plain schema-qualified table name"); } - if crate::meta::topic_id_by_name(topic).ok().flatten().is_none() { + // Checked with reads that assign no xid, so the slot can still be created below. + let topic_exists = read_one::( + "SELECT (SELECT topic_id::int FROM kafgres_topics WHERE name = $1)", + &[topic.into()], + ) + .unwrap_or_else(|e| error!("kafgres: {e}")) + .is_some(); + if !topic_exists { error!("kafgres: no such topic {topic:?}; create it before mapping to it"); } render_check(source_table, value_expr, key_expr, filter_expr); + // Created here, before this transaction writes, unless that would wait on a long + // transaction. Otherwise the worker creates it once the mapping commits; not an error, + // but changes before the slot is ready are not captured. + if !slot_exists() { + let late = slot_preconditions().err().or_else(|| { + if !running_transactions_finish(std::time::Duration::from_secs(1)) { + Some("a running transaction would hold the slot's creation".to_string()) + } else { + create_slot().err() + } + }); + if let Some(why) = late { + warning!( + "kafgres: {why}; the mapping is saved, but changes before its slot is ready are \ + not captured (kafgres_slot_ready() waits for it)" + ); + } + } + Spi::run_with_args( "INSERT INTO kafgres_cdc_mappings (mapping_name, source_table, topic, key_expr, value_expr, filter_expr) @@ -281,6 +307,18 @@ fn slot_exists() -> bool { .unwrap_or(false) } +/// Exists and has its snapshot, so it captures from here on. +fn slot_built() -> bool { + read_one::( + "SELECT EXISTS(SELECT 1 FROM pg_replication_slots + WHERE slot_name = $1 AND confirmed_flush_lsn IS NOT NULL)", + &[SLOT.into()], + ) + .ok() + .flatten() + .unwrap_or(false) +} + /// Checked, not attempted: a failed creation leaves an undroppable half-built slot behind. fn slot_preconditions() -> Result<(), String> { let wal_level = read_one::("SELECT current_setting('wal_level')", &[])? @@ -291,8 +329,9 @@ fn slot_preconditions() -> Result<(), String> { )); } + // `pg_settings`, not `current_setting`, which raises for a role that may not read it. if let Some(libs) = read_one::( - "SELECT current_setting('output_plugin_libraries', true)", + "SELECT (SELECT setting FROM pg_settings WHERE name = 'output_plugin_libraries')", &[], )? { if !libs.split(',').any(|l| l.trim() == "kafgres") { @@ -302,7 +341,19 @@ fn slot_preconditions() -> Result<(), String> { } } - if read_one::("SELECT pg_current_xact_id_if_assigned()::text::bigint", &[])?.is_some() { + if !read_one::( + "SELECT (SELECT rolsuper OR rolreplication FROM pg_roles WHERE rolname = current_user)", + &[], + )? + .unwrap_or(false) + { + return Err( + "cannot create the CDC slot: this role needs the REPLICATION attribute or superuser" + .to_string(), + ); + } + + if has_xid()? { return Err( "cannot create the CDC slot: this transaction has already been assigned an xid" .to_string(), @@ -311,6 +362,46 @@ fn slot_preconditions() -> Result<(), String> { Ok(()) } +/// The current transaction has written, so a slot cannot be created in it or wait on it. +fn has_xid() -> Result { + Ok(read_one::("SELECT pg_current_xact_id_if_assigned()::text::bigint", &[])?.is_some()) +} + +/// Whether `done` turns true within `within`, asking every `every`. +fn poll(within: std::time::Duration, every: std::time::Duration, mut done: impl FnMut() -> bool) -> bool { + let deadline = std::time::Instant::now() + within; + loop { + if done() { + return true; + } + if std::time::Instant::now() >= deadline { + return false; + } + pgrx::check_for_interrupts!(); + std::thread::sleep(every); + } +} + +/// Slot creation waits for the transactions holding an xid. If the ones running now finish +/// within `within`, those it waits on are likely as short. Assigns no xid. +fn running_transactions_finish(within: std::time::Duration) -> bool { + let Ok(Some(running)) = + with_latest_snapshot(|| read_one::("SELECT pg_current_snapshot()::text", &[])) + else { + return false; + }; + poll(within, std::time::Duration::from_millis(20), || { + read_one::( + "SELECT NOT EXISTS(SELECT 1 FROM pg_snapshot_xip($1::pg_snapshot) x + WHERE pg_xact_status(x) = 'in progress')", + &[running.as_str().into()], + ) + .ok() + .flatten() + .unwrap_or(false) + }) +} + fn enabled_mappings_exist() -> bool { read_one::( "SELECT EXISTS(SELECT 1 FROM kafgres_cdc_mappings WHERE enabled)", @@ -328,6 +419,12 @@ pub fn ensure_slot() -> Result { if !enabled_mappings_exist() { return Ok(false); } + create_slot()?; + Ok(true) +} + +/// Refused in a transaction that has written. +fn create_slot() -> Result<(), String> { slot_preconditions()?; // Serialise, then re-check: the worker and the SQL entry point race. An advisory lock assigns no xid. @@ -342,7 +439,7 @@ pub fn ensure_slot() -> Result { }) .map_err(|e| format!("could not lock for CDC slot creation: {e}"))?; if slot_exists() { - return Ok(true); + return Ok(()); } // `client.select`, not `Spi::run`: the latter assigns the xid that disqualifies creation. @@ -362,7 +459,7 @@ pub fn ensure_slot() -> Result { .map_err(|e| format!("could not reset CDC delivery positions: {e}"))?; log!("kafgres: created logical replication slot {SLOT:?} for CDC"); - Ok(true) + Ok(()) } /// Peek, not `get_changes`: advancing only after producing costs duplicates on a crash, not losses. @@ -939,6 +1036,13 @@ fn drain_locked(max_changes: i32, mode: Mode, lock: &DrainLock) -> Result, } +fn same_definitions(a: &[Mapping], b: &[Mapping]) -> bool { + a.len() == b.len() + && a.iter().zip(b).all(|(x, y)| { + (&x.name, &x.source, &x.topic, &x.value_expr, &x.key_expr, &x.filter_expr, &x.on_error) + == (&y.name, &y.source, &y.topic, &y.value_expr, &y.key_expr, &y.filter_expr, &y.on_error) + }) +} + +/// A read-only SPI query uses the active snapshot, which in a worker step predates commits +/// since the step began; this makes it see them. +fn with_latest_snapshot(f: impl FnOnce() -> T) -> T { + unsafe { pg_sys::PushActiveSnapshot(pg_sys::GetLatestSnapshot()) }; + let out = f(); + unsafe { pg_sys::PopActiveSnapshot() }; + out +} + fn load_mappings() -> Result, String> { Spi::connect(|client| { let rows = client.select( @@ -1473,6 +1594,27 @@ fn kafgres_cdc_create_slot() -> bool { } } +/// Wait until the slot is ready to capture: "slot ready", or "slot not ready" once +/// `timeout_ms` passes. The slot is created by `kafgres_add_mapping`, or by the worker once +/// the transactions it waits on finish. Creates nothing itself. +#[pg_extern] +fn kafgres_slot_ready(timeout_ms: default!(i32, 120000)) -> &'static str { + // The slot waits for this transaction, and its mapping is not visible to the worker yet. + if has_xid().unwrap_or_else(|e| error!("kafgres: {e}")) && !slot_built() { + warning!( + "kafgres: the slot cannot be ready before this transaction ends, since it has \ + written; call kafgres_slot_ready() after it commits" + ); + return "slot not ready"; + } + let within = std::time::Duration::from_millis(timeout_ms.max(0) as u64); + if poll(within, std::time::Duration::from_millis(100), slot_built) { + "slot ready" + } else { + "slot not ready" + } +} + /// If an enabled mapping remains, the worker recreates the slot; changes in between are not replayed. #[pg_extern] fn kafgres_cdc_drop_slot() -> bool { diff --git a/extension/src/group.rs b/extension/src/group.rs index 9694ba2..1131a8b 100644 --- a/extension/src/group.rs +++ b/extension/src/group.rs @@ -153,16 +153,49 @@ pub fn member(group_id: &str, member_id: &str) -> Result, spi::Er #[allow(clippy::too_many_arguments)] pub fn protocol_type_conflicts(group_id: &str, protocol_type: &str) -> Result { let conflicts: bool = Spi::get_one_with_args( + // A consumer or share group by this id conflicts too, as its heartbeats refuse ours. "SELECT COALESCE( (SELECT protocol_type IS NOT NULL AND protocol_type <> '' AND protocol_type <> $2 FROM kafgres_groups WHERE group_id = $1), - false)", + false) + OR EXISTS(SELECT 1 FROM kafgres_consumer_groups WHERE group_id = $1) + OR EXISTS(SELECT 1 FROM kafgres_share_groups WHERE group_id = $1)", &[group_id.into(), protocol_type.into()], )? .unwrap_or(false); Ok(conflicts) } +/// Whether `protocols` shares a name with every other member's. A static member's earlier +/// registration is the one this join replaces, so it is left out. +pub fn protocols_compatible( + group_id: &str, + member_id: &str, + group_instance_id: Option<&str>, + protocols: &[String], +) -> Result { + // Only the protocol names: a member's metadata and assignment can be large. + let others: Vec> = Spi::connect(|client| { + let rows = client.select( + "SELECT protocols FROM kafgres_group_members + WHERE group_id = $1 AND member_id <> $2 + AND ($3::text IS NULL OR group_instance_id IS DISTINCT FROM $3)", + None, + &[group_id.into(), member_id.into(), group_instance_id.into()], + )?; + let mut out = Vec::new(); + for row in rows { + out.push(row.get::>(1)?.unwrap_or_default()); + } + Ok::<_, spi::Error>(out) + })?; + let mut common: Vec = protocols.to_vec(); + for theirs in &others { + common.retain(|p| theirs.contains(p)); + } + Ok(!common.is_empty()) +} + pub fn join( group_id: &str, member_id: &str, diff --git a/extension/src/handlers/join_sync.rs b/extension/src/handlers/join_sync.rs index c59fe1a..a24b9a7 100644 --- a/extension/src/handlers/join_sync.rs +++ b/extension/src/handlers/join_sync.rs @@ -68,6 +68,18 @@ fn register_and_park( )))); } + if !group::protocols_compatible( + &req.group_id, + member_id, + req.group_instance_id.as_deref(), + &protocols, + )? { + return Ok(JoinOutcome::Reply(Box::new(error_join( + ErrorCode::InconsistentGroupProtocol, + member_id.to_string(), + )))); + } + let metadata = req.protocols[0].metadata.to_vec(); group::join( diff --git a/extension/src/handlers/produce.rs b/extension/src/handlers/produce.rs index a86355e..617135d 100644 --- a/extension/src/handlers/produce.rs +++ b/extension/src/handlers/produce.rs @@ -610,10 +610,12 @@ fn plan_partition( }); let decision = match prior { Some((_, prev)) => decide_after(&batches, prev, idem)?, - None => match producer::check( - producer_id, idem.epoch, idem.first_seq, idem.last_seq, tid, index, - ) - .map_err(|e| AppendError::ProducerState(e.to_string()))? + None => match producer::ensure_window(tid, index) + .map_err(AppendError::ProducerState) + .and_then(|()| { + producer::check(producer_id, idem.epoch, idem.first_seq, idem.last_seq, tid, index) + .map_err(|e| AppendError::ProducerState(e.to_string())) + })? { SequenceCheck::Append => Decision::Append(idem), SequenceCheck::Duplicate { base_offset } if base_offset < 0 => { diff --git a/extension/src/init210.rs b/extension/src/init210.rs new file mode 100644 index 0000000..ed42445 --- /dev/null +++ b/extension/src/init210.rs @@ -0,0 +1,10 @@ +//! One value per server start, folded into `kafgres_produce()` producer ids. + +use crate::ddl::run_ddl; + +pub fn init_210() { + run_ddl( + "CREATE SEQUENCE IF NOT EXISTS kafgres_boot_seq AS bigint START 1", + "create the boot sequence", + ); +} diff --git a/extension/src/lib.rs b/extension/src/lib.rs index 8e579c6..87c7b9c 100644 --- a/extension/src/lib.rs +++ b/extension/src/lib.rs @@ -37,6 +37,7 @@ mod init170; mod init180; mod init190; mod init200; +mod init210; pub mod quota; pub mod meta; pub mod plan; @@ -704,8 +705,12 @@ pub extern "C-unwind" fn _PG_init() { // Registered unconditionally rather than only when `storage_engine=segment`: shared crate::storage::segment::init_shmem(); { - use crate::server::EPOCHS_RAISED; + use crate::replication::{LAG_AT_MS, LAG_RECORDS}; + use crate::server::{BOOT_ID, EPOCHS_RAISED}; pgrx::pg_shmem_init!(EPOCHS_RAISED); + pgrx::pg_shmem_init!(BOOT_ID); + pgrx::pg_shmem_init!(LAG_RECORDS); + pgrx::pg_shmem_init!(LAG_AT_MS); } // The broker starts at `RecoveryFinished` so a standby never serves partitions; the @@ -935,6 +940,7 @@ pub(crate) fn ensure_tables_exist() { init180::init_180(); init190::init_190(); init200::init_200(); + init210::init_210(); } pgrx::extension_sql!( @@ -955,6 +961,7 @@ BEGIN AND d.deptype = 'e' AND p.proname NOT IN ('kafgres_produce', 'kafgres_version', 'kafgres_kafka_version', 'kafgres_partition_offsets', 'kafgres_archive_status', + 'kafgres_follower_lag', 'kafgres_cdc_status', 'kafgres_cdc_snapshots', 'kafgres_share_state') LOOP @@ -1055,12 +1062,14 @@ pub unsafe extern "C-unwind" fn kafgres_follower_worker_main(_arg: pg_sys::Datum reload_config(); if replicate_from().is_none() { log!("kafgres: follower stopping — replicate_from was cleared"); + replication::clear_lag(); return; } } // Stop before the broker starts serving: never two writers on one partition. if !pg_sys::RecoveryInProgress() { log!("kafgres: follower stopping — this node was promoted, the broker takes over"); + replication::clear_lag(); return; } @@ -1081,6 +1090,8 @@ pub unsafe extern "C-unwind" fn kafgres_follower_worker_main(_arg: pg_sys::Datum Ok(_) => delay = Duration::from_millis(50), Err(e) => { log!("kafgres: follower: {e}"); + // A reading from before the failure would show a lag that may have grown since. + replication::clear_lag(); follower.disconnect(); delay = Duration::from_millis(500); } diff --git a/extension/src/produce_sql.rs b/extension/src/produce_sql.rs index 3db37a2..983e90c 100644 --- a/extension/src/produce_sql.rs +++ b/extension/src/produce_sql.rs @@ -92,6 +92,9 @@ fn kafgres_produce( let xid: i64 = unsafe { pg_sys::GetTopFullTransactionId().value as i64 }; // SAFETY: as above. let subid = unsafe { pg_sys::GetCurrentSubTransactionId() }; + // The boot id in the producer id is known once the epochs are raised. + crate::server::wait_for_epochs(); + // Decided before the append: a batch must never carry a pre-rollback producer id. let (producer_id, reserve) = enter_produce(xid, subid, topic_id, partition); @@ -126,9 +129,6 @@ fn kafgres_produce( is_control: false, }; - // After a restart the log may still lose a torn tail and take a new epoch. - crate::server::wait_for_epochs(); - // Through the factory: whether an engine supports this is the engine's call. let mut store = crate::storage::open(); let (base_offset, last_offset) = @@ -187,35 +187,51 @@ struct Producing { reserved: HashSet<(u32, i32)>, /// Bumped when a subtransaction that produced rolls back; part of the producer id. generation: i64, + /// The server start this transaction first produced under; a restart mid-transaction + /// must not change its ids. + boot: Option, /// Highest subtransaction that produced under `generation`, 0 if none. max_subid: pg_sys::SubTransactionId, hooked: bool, } -/// Bits of a generation-tagged producer id below the generation. Bit 62 flags the tag, so an -/// untagged id is the xid itself. +/// Layout of a `kafgres_produce()` producer id, high to low: bit 62 (SQL, outside the wire +/// producer sequence's range), boot, generation, xid. +/// +/// The boot keeps an xid that Postgres reissues after a crash from sharing an id with the +/// orphaned batch of its earlier holder. const XID_BITS: u32 = 40; -const GENERATION_FLAG: i64 = 1 << 62; -const MAX_GENERATION: i64 = (1 << (62 - XID_BITS)) - 1; - -/// The producer id of a batch appended under `generation` of transaction `xid`. +const GENERATION_BITS: u32 = 10; +const BOOT_BITS: u32 = 12; +const SQL_FLAG: i64 = 1 << 62; +const MAX_GENERATION: i64 = (1 << GENERATION_BITS) - 1; +const XID_MASK: i64 = (1 << XID_BITS) - 1; +const BOOT_MASK: i64 = (1 << BOOT_BITS) - 1; + +/// The producer id of a batch appended under `generation` of transaction `xid` in `boot`. /// /// A rolled-back savepoint's orphan never gets an abort marker, so a client would drop every /// later batch of the same producer id; a new id after the rollback keeps them out of that range. -fn sql_producer_id(xid: i64, generation: i64) -> Option { - if generation == 0 { - return Some(xid); - } - if xid >> XID_BITS != 0 || generation > MAX_GENERATION { +/// +/// Only the xid's low bits are kept: a collision needs two transactions 2^40 xids apart in the +/// same boot, both still inside a producer's retention. +fn sql_producer_id(xid: i64, generation: i64, boot: i64) -> Option { + if generation > MAX_GENERATION { return None; } - Some(GENERATION_FLAG | (generation << XID_BITS) | xid) + let xid = xid & XID_MASK; + Some( + SQL_FLAG + | ((boot & BOOT_MASK) << (XID_BITS + GENERATION_BITS)) + | (generation << XID_BITS) + | xid, + ) } /// The top-level xid behind a `kafgres_produce()` batch's producer id. pub fn xid_of_sql_producer_id(producer_id: i64) -> i64 { - if producer_id & GENERATION_FLAG != 0 { - producer_id & ((1 << XID_BITS) - 1) + if producer_id & SQL_FLAG != 0 { + producer_id & XID_MASK } else { producer_id } @@ -228,7 +244,7 @@ fn enter_produce( topic: u32, partition: i32, ) -> (i64, bool) { - let (generation, reserve, hook) = PRODUCING.with(|s| { + let (generation, boot, reserve, hook) = PRODUCING.with(|s| { let mut s = s.borrow_mut(); if s.xid != xid { *s = Producing { xid, ..Producing::default() }; @@ -236,7 +252,8 @@ fn enter_produce( s.max_subid = s.max_subid.max(subid); let hook = !s.hooked; s.hooked = true; - (s.generation, !s.reserved.contains(&(topic, partition)), hook) + let boot = *s.boot.get_or_insert_with(crate::server::sql_boot); + (s.generation, boot, !s.reserved.contains(&(topic, partition)), hook) }); if hook { // pgrx drops subtransaction hooks at top-level commit or abort. @@ -245,7 +262,7 @@ fn enter_produce( move |subid, _parent| rolled_back(xid, subid), ); } - match sql_producer_id(xid, generation) { + match sql_producer_id(xid, generation, boot) { Some(id) => (id, reserve), None => error!( "kafgres_produce(): too many rolled-back savepoints that produced in transaction {xid}" @@ -430,31 +447,40 @@ mod producer_id_tests { use super::*; #[test] - fn generation_zero_is_the_xid_itself() { - assert_eq!(sql_producer_id(1234, 0), Some(1234)); - assert_eq!(xid_of_sql_producer_id(1234), 1234); - let big = (1 << 50) + 7; - assert_eq!(xid_of_sql_producer_id(sql_producer_id(big, 0).unwrap()), big); + fn every_id_is_tagged_and_decodes_to_the_xid() { + for xid in [1, 1234, (1 << XID_BITS) - 1] { + let id = sql_producer_id(xid, 0, 7).unwrap(); + assert!(id >= SQL_FLAG, "outside the wire sequence's range"); + assert_eq!(xid_of_sql_producer_id(id), xid); + } + assert_eq!(xid_of_sql_producer_id(1234), 1234, "an untagged id is the xid itself"); } #[test] - fn each_generation_is_a_distinct_producer_id_that_decodes_to_the_xid() { - let xid = (3i64 << 32) | 0xdead_beef; - let ids: Vec = (0..4).map(|g| sql_producer_id(xid, g).unwrap()).collect(); - for (i, a) in ids.iter().enumerate() { - assert!(*a >= 0, "a Kafka producer id is non-negative"); - assert_eq!(xid_of_sql_producer_id(*a), xid); - for b in &ids[i + 1..] { - assert_ne!(a, b); + fn generation_and_boot_each_make_a_distinct_producer_id() { + let xid = (3i64 << 32) | 0xbeef; + let mut ids = Vec::new(); + for boot in 0..3 { + for g in 0..4 { + let id = sql_producer_id(xid, g, boot).unwrap(); + assert!(id >= 0, "a Kafka producer id is non-negative"); + assert_eq!(xid_of_sql_producer_id(id), xid); + ids.push(id); } } - assert_ne!(sql_producer_id(xid, 1), sql_producer_id(xid + 1, 1)); + let n = ids.len(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), n); } #[test] - fn an_unencodable_generation_is_refused() { - assert_eq!(sql_producer_id(1 << XID_BITS, 1), None); - assert_eq!(sql_producer_id(5, MAX_GENERATION + 1), None); - assert!(sql_producer_id(5, MAX_GENERATION).is_some()); + fn an_xid_beyond_the_id_keeps_its_low_bits_and_a_spent_generation_is_refused() { + assert_eq!( + xid_of_sql_producer_id(sql_producer_id((5 << XID_BITS) | 9, 0, 0).unwrap()), + 9, + ); + assert_eq!(sql_producer_id(5, MAX_GENERATION + 1, 0), None); + assert!(sql_producer_id(5, MAX_GENERATION, BOOT_MASK).is_some()); } } diff --git a/extension/src/producer.rs b/extension/src/producer.rs index 4d2f8bb..4554b16 100644 --- a/extension/src/producer.rs +++ b/extension/src/producer.rs @@ -325,6 +325,244 @@ pub fn forget_past(topic_id: u32, partition: i32, end: i64) -> Result<(), spi::E ) } +/// How far below the log end a rebuild rescans, as (slack in ms before the newest batch's +/// timestamp, fewest offsets, most offsets). A crash loses only what was appended but not yet +/// committed, which is the log's tail; a row commits after its batch, so a lost row can sit +/// below another producer's committed one. Strict commits lose only requests in flight; with +/// `relaxed_produce_commit` the rows since the last WAL flush go too. +fn rebuild_lookback() -> (i64, i64, i64) { + if !crate::relaxed_produce_commit() { + return (1_000, 1_000, 50_000); + } + let wal_delay_ms = crate::plan::get_one::( + "SELECT (SELECT setting::bigint FROM pg_settings WHERE name = 'wal_writer_delay')", + &[], + ) + .ok() + .flatten() + .unwrap_or(200); + ((3 * wal_delay_ms).max(5_000), 5_000, 2_000_000) +} + +/// Where a rebuild starts: the offset `slack` before the log tail's timestamp, within the +/// bounds. Client timestamps can be stale or ahead, so the bounds decide when they mislead. +fn rebuild_start( + store: &dyn crate::storage::LogStore, + topic_id: u32, + partition: i32, + end: i64, +) -> Result { + use kafgres_codec::records::RecordBatch; + + let (slack_ms, fewest, most) = rebuild_lookback(); + let nearest = end.saturating_sub(fewest); + let farthest = end.saturating_sub(most); + let tail = store + .read(topic_id, partition, end - 1, 1 << 20, crate::storage::IsolationLevel::ReadUncommitted) + .map_err(|e| e.to_string())?; + // The first batch of the slice, which holds `end - 1`; later appends may follow it. + let first_len = tail + .bytes + .get(kafgres_codec::records::LENGTH_OFFSET..kafgres_codec::records::LENGTH_OFFSET + 4) + .map(|b| i32::from_be_bytes(b.try_into().expect("4 bytes")).max(0) as usize) + .map(|len| kafgres_codec::records::LENGTH_OFFSET + 4 + len); + let Some(tail_ts) = first_len + .and_then(|n| tail.bytes.get(..n)) + .and_then(|b| RecordBatch::new(kafgres_codec::bytes::Bytes::copy_from_slice(b)).ok()) + .map(|b| b.max_timestamp()) + // One ahead by more than the slack would find little or nothing after it, the + // narrowest start. + .filter(|&t| { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(i64::MAX, |d| d.as_millis() as i64); + t > slack_ms && t <= now.saturating_add(slack_ms) + }) + else { + return Ok(farthest); + }; + let from = store + .offset_for_timestamp(topic_id, partition, tail_ts - slack_ms) + .map_err(|e| e.to_string())? + .map_or(nearest, |(offset, _)| offset); + Ok(from.clamp(farthest, nearest)) +} + +type Partitions = std::collections::HashSet<(u32, i32)>; + +static STARTUP_END: std::sync::Mutex>> = + std::sync::Mutex::new(None); + +/// The log end when this process started, where a window rebuild looks back from. +pub fn note_startup_end(topic_id: u32, partition: i32, end: i64) { + STARTUP_END + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get_or_insert_with(Default::default) + .insert((topic_id, partition), end); +} + +static WINDOW_REBUILT: std::sync::Mutex> = std::sync::Mutex::new(None); + +thread_local! { + /// Rebuilt in the current transaction, not yet committed. + static REBUILT_PENDING: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; +} + +/// Before the first idempotent produce to a partition after this process starts, rebuild its +/// window from the log (segment engine only; the table engine commits a batch and its row +/// together). The partition counts as rebuilt only once the transaction holding the rows +/// commits; a rollback, or a failure, leaves it for the next produce to try again. +pub fn ensure_window(topic_id: u32, partition: i32) -> Result<(), String> { + use pgrx::{PgSubXactCallbackEvent, PgXactCallbackEvent}; + + if crate::storage_engine_guc() != "segment" { + return Ok(()); + } + let key = (topic_id, partition); + if WINDOW_REBUILT + .lock() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .is_some_and(|s| s.contains(&key)) + || REBUILT_PENDING.with(|p| p.borrow().as_ref().is_some_and(|s| s.contains(&key))) + { + return Ok(()); + } + let store = crate::storage::open(); + rebuild_window(&*store, topic_id, partition).map_err(|e| { + format!("could not rebuild the idempotence window for topic {topic_id} partition {partition}: {e}") + })?; + + let first = REBUILT_PENDING.with(|p| { + let mut p = p.borrow_mut(); + let first = p.is_none(); + p.get_or_insert_with(Default::default).insert(key); + first + }); + if first { + pgrx::register_xact_callback(PgXactCallbackEvent::Commit, || { + let done = REBUILT_PENDING.with(|p| p.borrow_mut().take()); + if let Some(done) = done { + WINDOW_REBUILT + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get_or_insert_with(Default::default) + .extend(done); + } + }); + pgrx::register_xact_callback(PgXactCallbackEvent::Abort, || { + REBUILT_PENDING.with(|p| p.borrow_mut().take()); + }); + // A rolled-back subtransaction may hold the rows: forget, and rebuild again. + pgrx::register_subxact_callback(PgSubXactCallbackEvent::AbortSub, |_, _| { + REBUILT_PENDING.with(|p| p.borrow_mut().take()); + }); + } + Ok(()) +} + +/// Re-create window rows for batches the log holds and the window lacks, so a retry after a +/// crash between the append and the row's commit is not appended twice. +pub fn rebuild_window( + store: &dyn crate::storage::LogStore, + topic_id: u32, + partition: i32, +) -> Result { + use kafgres_codec::records::{RecordBatch, LENGTH_OFFSET}; + use std::collections::HashMap; + + // Only what the log held at startup can have lost a row; later appends were recorded. + let end = store + .high_watermark(topic_id, partition) + .map_err(|e| e.to_string())?; + let log_start = store + .log_start_offset(topic_id, partition) + .map_err(|e| e.to_string())?; + // A partition missing from a recorded startup was created since, so it lost nothing. + let end = match STARTUP_END.lock().unwrap_or_else(|e| e.into_inner()).as_ref() { + Some(m) => m.get(&(topic_id, partition)).map_or(log_start, |&at_start| at_start.min(end)), + None => end, + }; + let mut at = if end > log_start { + rebuild_start(store, topic_id, partition, end)?.max(log_start) + } else { + log_start + }; + + // Per producer: (epoch, last_seq) of its newest row, `None` with no window here. + let mut newest: HashMap> = HashMap::new(); + let mut rebuilt = 0u64; + while at < end { + let slice = store + .read(topic_id, partition, at, 4 << 20, crate::storage::IsolationLevel::ReadUncommitted) + .map_err(|e| e.to_string())?; + if slice.bytes.is_empty() { + break; + } + let mut pos = 0usize; + while pos + LENGTH_OFFSET + 4 <= slice.bytes.len() { + let len = i32::from_be_bytes( + slice.bytes[pos + LENGTH_OFFSET..pos + LENGTH_OFFSET + 4] + .try_into() + .expect("4 bytes"), + ); + let total = LENGTH_OFFSET + 4 + len.max(0) as usize; + if len <= 0 || pos + total > slice.bytes.len() { + break; + } + let bytes = kafgres_codec::bytes::Bytes::copy_from_slice(&slice.bytes[pos..pos + total]); + pos += total; + let Ok(batch) = RecordBatch::new(bytes) else { continue }; + let (pid, epoch, first) = (batch.producer_id(), batch.producer_epoch(), batch.base_sequence()); + // A batch straddling the start is read whole and kept: it may be the newest one. + if pid < 0 || epoch < 0 || first < 0 || batch.is_control() { + continue; + } + let last = increment_sequence(first, batch.last_offset_delta()); + let state = match newest.get(&pid) { + Some(s) => *s, + None => { + let known = crate::plan::get_one::( + "SELECT (SELECT true FROM kafgres_producers WHERE producer_id = $1)", + &[pid.into()], + ) + .map_err(|e| e.to_string())? + .unwrap_or(false); + let s = if !known { + // Expired or never registered. + Some((i16::MAX, i32::MIN)) + } else { + window(pid, topic_id, partition) + .map_err(|e| e.to_string())? + .0 + .last() + .map(|r| (r.epoch, r.last_seq)) + }; + newest.insert(pid, s); + s + } + }; + let missing = match state { + None => true, + Some((e, l)) => epoch > e || (epoch == e && in_sequence(l, first)), + }; + if !missing { + continue; + } + record(pid, epoch, topic_id, partition, first, last, batch.base_offset()) + .map_err(|e| e.to_string())?; + newest.insert(pid, Some((epoch, last))); + rebuilt += 1; + } + if pos == 0 { + break; + } + at = slice.next_offset.max(at + 1); + } + Ok(rebuilt) +} + /// Stamp the offset a `base_offset = -1` reservation was appended at. pub fn stamp( producer_id: i64, diff --git a/extension/src/replication.rs b/extension/src/replication.rs index d1a05f0..625858b 100644 --- a/extension/src/replication.rs +++ b/extension/src/replication.rs @@ -38,6 +38,34 @@ const MAX_FETCH_BYTES: i32 = 64 * 1024 * 1024; /// every quiet moment. const FETCH_MAX_WAIT_MS: i32 = 100; +/// Records the leader held beyond the follower's fetch offsets at its last Fetch, and when +/// it was answered (epoch ms, 0 before the first). +pub static LAG_RECORDS: pgrx::PgAtomic = + unsafe { pgrx::PgAtomic::new(c"kafgres_follower_lag_records") }; +pub static LAG_AT_MS: pgrx::PgAtomic = + unsafe { pgrx::PgAtomic::new(c"kafgres_follower_lag_at_ms") }; + +/// Forget the last reading: the node no longer follows, or the last Fetch was incomplete. +pub fn clear_lag() { + use std::sync::atomic::Ordering; + LAG_AT_MS.get().store(0, Ordering::Release); + LAG_RECORDS.get().store(0, Ordering::Release); +} + +/// Follower lag in records at its last Fetch, and that Fetch's time in epoch ms. Empty +/// before the first. `pg_stat_replication` does not cover this. +#[pg_extern] +fn kafgres_follower_lag() -> TableIterator<'static, (name!(lag_records, i64), name!(observed_at_ms, i64))> { + use std::sync::atomic::Ordering; + let at = LAG_AT_MS.get().load(Ordering::Acquire); + let rows = if at == 0 { + Vec::new() + } else { + vec![(LAG_RECORDS.get().load(Ordering::Acquire), at)] + }; + TableIterator::new(rows) +} + /// A round returns to the worker loop after this long, so leadership and the partition /// list are re-read even while the leader always has more to give. const ROUND_BUDGET: Duration = Duration::from_millis(500); @@ -279,6 +307,40 @@ pub fn fetch_from( FetchResponse::decode, )?; + // Against the last stable offset, not the high watermark: the follower copies only what + // has committed, so an open transaction is not lag. A partition the leader did not + // answer, answered with an error, or left without offsets once its response budget ran + // out, leaves the reading incomplete. + let mut lag = 0i64; + let mut complete = true; + for (_, name, partition, from, _) in want { + let part = response + .responses + .iter() + .filter(|t| t.topic == *name) + .flat_map(|t| t.partitions.iter()) + .find(|p| p.partition_index == *partition); + match part { + Some(p) if p.error_code == 0 && p.high_watermark >= 0 => { + let stable = if p.last_stable_offset >= 0 { p.last_stable_offset } else { p.high_watermark }; + lag += (stable - *from).max(0); + } + _ => complete = false, + } + } + if complete { + use std::sync::atomic::Ordering; + LAG_RECORDS.get().store(lag, Ordering::Release); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0); + LAG_AT_MS.get().store(now, Ordering::Release); + } else { + // A partial sum would understate the lag; no reading is better than a low one. + clear_lag(); + } + let mut pulled = Vec::new(); for topic in response.responses { let name = topic.topic; diff --git a/extension/src/retention.rs b/extension/src/retention.rs index aa6477a..800ee46 100644 --- a/extension/src/retention.rs +++ b/extension/src/retention.rs @@ -73,13 +73,18 @@ pub fn sweep(cursor: u32) -> Result { } } } - last = *topic_id; - continue; + } else { + let policy = crate::config::retention_policy(*topic_id)?; + match store.enforce_retention(*topic_id, &policy) { + Ok(n) => reclaimed += n, + Err(e) => pgrx::log!("kafgres: retention on topic {topic_id} failed: {e}"), + } } - let policy = crate::config::retention_policy(*topic_id)?; - match store.enforce_retention(*topic_id, &policy) { - Ok(n) => reclaimed += n, - Err(e) => pgrx::log!("kafgres: retention on topic {topic_id} failed: {e}"), + // Whatever moved the log start, retention or DeleteRecords, under any policy. + if crate::storage_engine_guc() == "segment" { + if let Err(e) = forget_markers_below_start(*topic_id) { + pgrx::log!("kafgres: forgetting markers on topic {topic_id} failed: {e}"); + } } last = *topic_id; } @@ -87,6 +92,14 @@ pub fn sweep(cursor: u32) -> Result { Ok(Batch { next: last, examined: topics.len() as u64, reclaimed }) } +fn forget_markers_below_start(topic_id: u32) -> crate::storage::StoreResult<()> { + use crate::storage::pmeta; + for partition in pmeta::partitions(topic_id)? { + pmeta::forget_markers_below(topic_id, partition, pmeta::log_start_offset(topic_id, partition)?)?; + } + Ok(()) +} + fn compact_topic(store: &mut dyn crate::storage::LogStore, topic_id: u32) -> Result { let partitions: Vec = Spi::connect(|client| { let rows = client.select( diff --git a/extension/src/server/mod.rs b/extension/src/server/mod.rs index ff3105c..df32282 100644 --- a/extension/src/server/mod.rs +++ b/extension/src/server/mod.rs @@ -246,6 +246,14 @@ impl Server { pub static EPOCHS_RAISED: pgrx::PgAtomic = unsafe { pgrx::PgAtomic::new(c"kafgres_epochs_raised") }; +/// This server start's number, from `kafgres_boot_seq`; valid once `EPOCHS_RAISED` is set. +pub static BOOT_ID: pgrx::PgAtomic = + unsafe { pgrx::PgAtomic::new(c"kafgres_boot_id") }; + +pub fn sql_boot() -> i64 { + BOOT_ID.get().load(Ordering::Acquire) +} + /// Wait up to 30 s for the broker to raise the leader epochs, then raise an ERROR. pub fn wait_for_epochs() { let deadline = Instant::now() + Duration::from_secs(30); @@ -1356,6 +1364,17 @@ fn exclusive_holders() -> String { fn raise_leader_epochs() -> bool { let result = BackgroundWorker::transaction(|| { crate::dbtx::guarded(|| { + // Once per process: a retried raise must not spend another number. + static PROCESS_BOOT: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0); + let mut boot = PROCESS_BOOT.load(Ordering::Relaxed); + // Flushed at commit even with asynchronous commit, or a crash could hand the next + // start this number along with the xids it reissues. Every attempt: a failed one + // took the number but flushed nothing. + pgrx::Spi::run("SET LOCAL synchronous_commit = on")?; + if boot == 0 { + boot = pgrx::Spi::get_one("SELECT nextval('kafgres_boot_seq')")?.unwrap_or(0); + PROCESS_BOOT.store(boot, Ordering::Relaxed); + } // The timeline increments on `pg_promote`. let wal_timeline: i32 = pgrx::Spi::get_one( "SELECT (('x' || substr(pg_walfile_name(pg_current_wal_lsn()), 1, 8))::bit(32)::int)", @@ -1392,18 +1411,22 @@ fn raise_leader_epochs() -> bool { .map_err(|e| HandlerError::Internal(e.to_string()))?; crate::producer::forget_past(topic_id, partition, end) .map_err(|e| HandlerError::Internal(e.to_string()))?; + crate::producer::note_startup_end(topic_id, partition, end); } - Ok((highest, raised)) + Ok((highest, raised, boot)) }) }); match result { - Ok((epoch, 0)) => { + Ok((epoch, 0, boot)) => { + // Published once committed: a raise that rolls back never exposes its number. + BOOT_ID.get().store(boot, Ordering::Release); log!("kafgres: leader epoch {epoch} already current for every partition"); EPOCHS_RAISED.get().store(true, Ordering::Release); true } - Ok((epoch, n)) => { + Ok((epoch, n, boot)) => { + BOOT_ID.get().store(boot, Ordering::Release); log!("kafgres: took a new leader epoch (up to {epoch}) for {n} partition(s)"); EPOCHS_RAISED.get().store(true, Ordering::Release); true diff --git a/extension/src/storage/pmeta.rs b/extension/src/storage/pmeta.rs index c366eb8..53af83d 100644 --- a/extension/src/storage/pmeta.rs +++ b/extension/src/storage/pmeta.rs @@ -204,6 +204,8 @@ pub fn partitions(topic: TopicId) -> StoreResult> { .map_err(spi_err) } +const MARKERS_FORGOTTEN_PER_PASS: i64 = 20_000; + /// Move `log_start_offset` forward. Never backward — a consumer that has been told an pub fn advance_log_start(topic: TopicId, partition: i32, offset: i64) -> StoreResult<()> { Spi::run_with_args( @@ -216,6 +218,23 @@ pub fn advance_log_start(topic: TopicId, partition: i32, offset: i64) -> StoreRe forget_aborted_below(topic, partition, offset) } +/// Batches wholly below the log start are unreadable. Every retention pass calls this, a +/// bounded number at a time, so a large backlog (one an upgrade leaves, say) goes over +/// several passes instead of stalling the broker in one. The `base_offset` bound is implied +/// by `last_offset` but lets the planner use the primary key's range instead of scanning. +pub fn forget_markers_below(topic: TopicId, partition: i32, offset: i64) -> StoreResult<()> { + Spi::run_with_args( + "DELETE FROM kafgres_markers + WHERE topic_id = $1::oid AND partition = $2 AND last_offset < $3 + AND base_offset < LEAST($3, (SELECT base_offset FROM kafgres_markers + WHERE topic_id = $1::oid AND partition = $2 + AND base_offset < $3 + ORDER BY base_offset OFFSET $4 LIMIT 1))", + &[(topic as i32).into(), partition.into(), offset.into(), MARKERS_FORGOTTEN_PER_PASS.into()], + ) + .map_err(spi_err) +} + /// Cap on markers loaded per Fetch. Without it a `read_committed` consumer far behind the /// LSO could pull an unbounded set into memory; a Rust allocation failure in a Postgres /// backend is an `abort()`, crashing the whole cluster. @@ -267,15 +286,26 @@ pub fn committed_markers( } /// Whether a `kafgres_produce()` batch's marker exists, and whether its writer has finished, -/// in one statement so both see one snapshot. +/// in one statement so both see one snapshot. `xid` is the low bits of the writer's xid8; the +/// writer is the xid8 with those bits nearest the snapshot's horizon. pub fn marker_state(topic: TopicId, partition: i32, base: i64, xid: i64) -> StoreResult<(bool, bool)> { + // Below 2^40 the id holds the whole xid8, so the plain lookup is exact and cheaper. + let writer = if xid_epoch_is_zero()? { + "$4::bigint" + } else { + "(SELECT CASE WHEN d < (1::bigint << 39) THEN x - d ELSE x + (1::bigint << 40) - d END + FROM (SELECT x, (x - $4::bigint) & ((1::bigint << 40) - 1) AS d + FROM (SELECT pg_snapshot_xmax(pg_current_snapshot())::text::bigint AS x) a) b)" + }; Spi::connect(|client| { let row = client .select( - "SELECT EXISTS (SELECT 1 FROM kafgres_markers - WHERE topic_id = $1::oid AND partition = $2 - AND base_offset = $3), - pg_visible_in_snapshot($4::text::xid8, pg_current_snapshot())", + &format!( + "SELECT EXISTS (SELECT 1 FROM kafgres_markers + WHERE topic_id = $1::oid AND partition = $2 + AND base_offset = $3), + pg_visible_in_snapshot(({writer})::text::xid8, pg_current_snapshot())" + ), Some(1), &[(topic as i32).into(), partition.into(), base.into(), xid.into()], )? @@ -288,6 +318,45 @@ pub fn marker_state(topic: TopicId, partition: i32, base: i64, xid: i64) -> Stor .map_err(spi_err) } +/// Whether every xid8 so far fits the producer id's xid bits. "No" is permanent; "yes" is +/// rechecked every few minutes, since the counter can cross 2^40 while a process runs. +fn xid_epoch_is_zero() -> StoreResult { + use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; + const RECHECK_SECS: i64 = 300; + static CROSSED: AtomicBool = AtomicBool::new(false); + static ZERO_UNTIL: AtomicI64 = AtomicI64::new(0); + + if CROSSED.load(Ordering::Relaxed) { + return Ok(false); + } + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + if now < ZERO_UNTIL.load(Ordering::Relaxed) { + return Ok(true); + } + // `client.select`: `get_one` would assign this read an xid. + let zero = Spi::connect(|client| { + client + .select( + "SELECT (pg_snapshot_xmax(pg_current_snapshot())::text::bigint >> 40) = 0", + Some(1), + &[], + )? + .first() + .get_one::() + }) + .map_err(spi_err)? + .unwrap_or(true); + if zero { + ZERO_UNTIL.store(now + RECHECK_SECS, Ordering::Relaxed); + } else { + CROSSED.store(true, Ordering::Relaxed); + } + Ok(zero) +} + /// Whether an aborted Kafka transaction of `producer_id` covers `offset`. pub fn aborted_at(topic: TopicId, partition: i32, producer_id: i64, offset: i64) -> StoreResult { Spi::get_one_with_args::( @@ -435,3 +504,13 @@ pub fn forget_aborted_below(topic: TopicId, partition: i32, offset: i64) -> Stor ) .map_err(spi_err) } + +/// Markers of batches compaction removed. +pub fn forget_markers(topic: TopicId, partition: i32, bases: &[i64]) -> StoreResult<()> { + Spi::run_with_args( + "DELETE FROM kafgres_markers + WHERE topic_id = $1::oid AND partition = $2 AND base_offset = ANY($3)", + &[(topic as i32).into(), partition.into(), bases.to_vec().into()], + ) + .map_err(spi_err) +} diff --git a/extension/src/storage/segment.rs b/extension/src/storage/segment.rs index 9e7c4c5..3e079ae 100644 --- a/extension/src/storage/segment.rs +++ b/extension/src/storage/segment.rs @@ -867,6 +867,8 @@ struct Rewrite { pos: u64, written: u64, removed: u64, + /// Base offsets of batches the rewrite drops whole. + dropped: Vec, writing: bool, /// This process's replacement file, removed when the rewrite is dropped unfinished. tmp: PathBuf, @@ -1029,6 +1031,8 @@ impl SegmentStore { if let Some(bytes) = rebuild_batch(&view, &kept) { out.write_all_at(&bytes, rw.written)?; rw.written += bytes.len() as u64; + } else { + rw.dropped.push(view.base_offset()); } } } @@ -1037,7 +1041,12 @@ impl SegmentStore { drop(out); Ok(Some( match Self::swap_segment(topic, partition, &rw.info, rw.written == 0, generation)? { - Some(generation) => Cleaned::Swapped(rw.removed, generation), + Some(generation) => { + if let Err(e) = pmeta::forget_markers(topic, partition, &rw.dropped) { + pgrx::log!("kafgres: could not drop the commit markers compaction removed: {e}"); + } + Cleaned::Swapped(rw.removed, generation) + } None => Cleaned::Abandoned, }, )) @@ -1254,6 +1263,7 @@ impl SegmentStore { pos: 0, written: 0, removed: 0, + dropped: Vec::new(), writing: false, tmp: compacting_path(topic, partition, s.base), }) @@ -1433,6 +1443,58 @@ impl SegmentStore { } impl SegmentStore { + /// The base of the first batch ending at or after `from` and starting below `ceiling`. + /// A compacted log is sparse, so it can be far past `from`; markers in between are moot. + fn next_batch_base( + topic: TopicId, + partition: i32, + from: i64, + ceiling: i64, + bases_on_disk: &[i64], + ) -> StoreResult> { + let mut bases = bases_on_disk.to_vec(); + bases.sort_unstable(); + Self::with_slot(topic, partition, |st, hints| { + for (bi, base) in bases.iter().enumerate() { + if bases.get(bi + 1).is_some_and(|next_base| *next_base <= from) { + continue; + } + let path = data_path(&segment_path(topic, partition, *base, "log")); + let Ok(vfd) = Vfd::open(&path, false) else { + continue; // reclaimed while we looked + }; + let data_end = if *base == st.active_base { + st.active_bytes + } else { + vfd.size()? + }; + let mut pos = hints + .get(base) + .and_then(|entries| entries.iter().rev().find(|(b, _)| *b <= from).map(|(_, p)| *p)) + .unwrap_or(0); + let mut header = [0u8; records::RECORD_BATCH_OVERHEAD]; + while pos + header.len() as u64 <= data_end { + if vfd.read_at(&mut header, pos)? != header.len() { + break; + } + let Some(BatchSpan { total, base: batch_base, last: batch_last }) = + batch_span(&header, pos, data_end) + else { + break; + }; + if batch_base >= ceiling { + return Ok(None); + } + if batch_last >= from { + return Ok(Some(batch_base)); + } + pos += total; + } + } + Ok(None) + }) + } + fn segment_bases(topic: TopicId, partition: i32) -> StoreResult> { let dir = partition_dir(topic, partition); let entries = match std::fs::read_dir(data_path(&dir)) { @@ -1893,6 +1955,37 @@ pub fn migrate_time_indexes() -> std::io::Result { Ok(removed) } +/// A batch's byte length and the offsets it spans. +struct BatchSpan { + total: u64, + base: i64, + last: i64, +} + +/// From a batch's fixed header at `pos`; `None` when the header is corrupt or the batch runs +/// past `data_end`. +fn batch_span(header: &[u8], pos: u64, data_end: u64) -> Option { + let length = i32::from_be_bytes( + header[records::LENGTH_OFFSET..records::LENGTH_OFFSET + 4].try_into().ok()?, + ); + if length <= 0 { + return None; + } + let total = records::LENGTH_OFFSET as u64 + 4 + length as u64; + if pos + total > data_end { + return None; + } + let base = i64::from_be_bytes( + header[records::BASE_OFFSET_OFFSET..records::BASE_OFFSET_OFFSET + 8].try_into().ok()?, + ); + let last_delta = i32::from_be_bytes( + header[records::LAST_OFFSET_DELTA_OFFSET..records::LAST_OFFSET_DELTA_OFFSET + 4] + .try_into() + .ok()?, + ); + Some(BatchSpan { total, base, last: base + last_delta as i64 }) +} + /// Whether a transactional batch came from `kafgres_produce()`, which stamps producer epoch /// -1. From the header alone: the producer table forgets idle producers, and a SQL xid can /// equal a Kafka producer id. @@ -1995,154 +2088,154 @@ impl LogStore for SegmentStore { let mut next = offset; let mut aborted: Vec = Vec::new(); + let mut offset = offset; if offset < ceiling { // Enumerate from disk, not the hint map: hints are per-process, so a partition appended by // TODO: cache the per-Fetch `read_dir`; the list changes only on roll/reclaim. let bases_on_disk = Self::segment_bases(topic, partition)?; - // The marker load can lower the ceiling: past its cap, reading beyond it would - // judge a batch committed on missing evidence. - let (committed, ceiling) = match isolation { - IsolationLevel::ReadCommitted => { - let (set, capped) = - pmeta::committed_markers(topic, partition, offset, ceiling)?; - (Some(set), capped) - } - IsolationLevel::ReadUncommitted => (None, ceiling), - }; - Self::with_slot(topic, partition, |st, hints| { - let mut bases = bases_on_disk; - bases.sort_unstable(); - let mut reclaimed: Vec = Vec::new(); - - for (bi, base) in bases.iter().enumerate() { - if let Some(next_base) = bases.get(bi + 1) { - if *next_base <= offset { - continue; - } + loop { + // The marker load can lower the ceiling: past its cap, reading beyond it would + // judge a batch committed on missing evidence. + let (committed, limit) = match isolation { + IsolationLevel::ReadCommitted => { + let (set, capped) = + pmeta::committed_markers(topic, partition, offset, ceiling)?; + (Some(set), capped) } - - let path = data_path(&segment_path(topic, partition, *base, "log")); - let vfd = match Vfd::open(&path, false) { - Ok(v) => v, - // Reclaimed underneath us: retention runs from user backends too, and a missing file is - Err(_) if bases.get(bi + 1).is_some_and(|nb| *nb <= log_start) => { - reclaimed.push(*base); - continue; + IsolationLevel::ReadUncommitted => (None, ceiling), + }; + Self::with_slot(topic, partition, |st, hints| { + let mut bases = bases_on_disk.clone(); + bases.sort_unstable(); + let mut reclaimed: Vec = Vec::new(); + + for (bi, base) in bases.iter().enumerate() { + if let Some(next_base) = bases.get(bi + 1) { + if *next_base <= offset { + continue; + } } - Err(e) => return Err(e), - }; - let data_end = if *base == st.active_base { - st.active_bytes - } else { - vfd.size()? - }; - - let mut pos = hints - .get(base) - .and_then(|entries| { - entries.iter().rev().find(|(b, _)| *b <= next).map(|(_, p)| *p) - }) - .unwrap_or(0); + let path = data_path(&segment_path(topic, partition, *base, "log")); + let vfd = match Vfd::open(&path, false) { + Ok(v) => v, + // Reclaimed underneath us: retention runs from user backends too, and a missing file is + Err(_) if bases.get(bi + 1).is_some_and(|nb| *nb <= log_start) => { + reclaimed.push(*base); + continue; + } + Err(e) => return Err(e), + }; + + let data_end = if *base == st.active_base { + st.active_bytes + } else { + vfd.size()? + }; + + let mut pos = hints + .get(base) + .and_then(|entries| { + entries.iter().rev().find(|(b, _)| *b <= next).map(|(_, p)| *p) + }) + .unwrap_or(0); - let mut header = [0u8; records::RECORD_BATCH_OVERHEAD]; - while pos + header.len() as u64 <= data_end { - if vfd.read_at(&mut header, pos)? != header.len() { - break; - } - let length = i32::from_be_bytes( - header[records::LENGTH_OFFSET..records::LENGTH_OFFSET + 4] - .try_into() - .expect("4 bytes"), - ); - if length <= 0 { - break; - } - let total = records::LENGTH_OFFSET as u64 + 4 + length as u64; - if pos + total > data_end { - break; - } - let batch_base = i64::from_be_bytes( - header[records::BASE_OFFSET_OFFSET..records::BASE_OFFSET_OFFSET + 8] - .try_into() - .expect("8 bytes"), - ); - let last_delta = i32::from_be_bytes( - header[records::LAST_OFFSET_DELTA_OFFSET - ..records::LAST_OFFSET_DELTA_OFFSET + 4] - .try_into() - .expect("4 bytes"), - ); - let batch_last = batch_base + last_delta as i64; - - // Wholly below the request; the batch containing `offset` is kept whole (the consumer - if batch_last < offset { - pos += total; - continue; - } - // Stop at the ceiling rather than truncating the batch: batches - if batch_base >= ceiling { - break; - } + let mut header = [0u8; records::RECORD_BATCH_OVERHEAD]; + while pos + header.len() as u64 <= data_end { + if vfd.read_at(&mut header, pos)? != header.len() { + break; + } + let Some(BatchSpan { total, base: batch_base, last: batch_last }) = + batch_span(&header, pos, data_end) + else { + break; + }; + + // Wholly below the request; the batch containing `offset` is kept whole (the consumer + if batch_last < offset { + pos += total; + continue; + } + // Stop at the ceiling rather than truncating the batch: batches + if batch_base >= limit { + break; + } - // A marker-backed batch with no committed marker is an orphan: a rolled-back - if let Some(committed) = &committed { - let attributes = i16::from_be_bytes( - header[records::ATTRIBUTES_OFFSET - ..records::ATTRIBUTES_OFFSET + 2] - .try_into() - .expect("2 bytes"), - ); - let producer_id = i64::from_be_bytes( - header[records::PRODUCER_ID_OFFSET - ..records::PRODUCER_ID_OFFSET + 8] - .try_into() - .expect("8 bytes"), - ); - let is_txn = attributes & records::TRANSACTIONAL_FLAG != 0; - let is_control = - attributes & records::CONTROL_BATCH_FLAG != 0; - - if is_txn - && !is_control - && is_marker_backed(&header) - && !committed.contains(&batch_base) - { - aborted.push(super::AbortedTxn { - producer_id, - first_offset: batch_base, - }); + // A marker-backed batch with no committed marker is an orphan: a rolled-back + if let Some(committed) = &committed { + let attributes = i16::from_be_bytes( + header[records::ATTRIBUTES_OFFSET + ..records::ATTRIBUTES_OFFSET + 2] + .try_into() + .expect("2 bytes"), + ); + let producer_id = i64::from_be_bytes( + header[records::PRODUCER_ID_OFFSET + ..records::PRODUCER_ID_OFFSET + 8] + .try_into() + .expect("8 bytes"), + ); + let is_txn = attributes & records::TRANSACTIONAL_FLAG != 0; + let is_control = + attributes & records::CONTROL_BATCH_FLAG != 0; + + if is_txn + && !is_control + && is_marker_backed(&header) + && !committed.contains(&batch_base) + { + aborted.push(super::AbortedTxn { + producer_id, + first_offset: batch_base, + }); + } } - } - if !out.is_empty() && out.len() + total as usize > max_bytes { - for b in reclaimed { - hints.remove(&b); + if !out.is_empty() && out.len() + total as usize > max_bytes { + for b in reclaimed { + hints.remove(&b); + } + return Ok(()); } - return Ok(()); - } - let mut buf = vec![0u8; total as usize]; - if vfd.read_at(&mut buf, pos)? != buf.len() { - break; - } - out.extend_from_slice(&buf); - next = batch_last + 1; - pos += total; + let mut buf = vec![0u8; total as usize]; + if vfd.read_at(&mut buf, pos)? != buf.len() { + break; + } + out.extend_from_slice(&buf); + next = batch_last + 1; + pos += total; - if out.len() >= max_bytes { - for b in reclaimed { - hints.remove(&b); + if out.len() >= max_bytes { + for b in reclaimed { + hints.remove(&b); + } + return Ok(()); } - return Ok(()); + } + } + for b in reclaimed { + hints.remove(&b); + } + Ok(()) + })?; + // Nothing below the lowered ceiling (a compacted log is sparse): resume at the + // next batch, not the next window of markers, which may all be stale. + if out.is_empty() && limit < ceiling { + match Self::next_batch_base(topic, partition, limit, ceiling, &bases_on_disk)? { + Some(base) => { + offset = base; + next = base; + continue; + } + None => { + next = ceiling; + break; } } } - for b in reclaimed { - hints.remove(&b); - } - Ok(()) - })?; + break; + } // Kafka's aborts, scoped to what this response returned: bounding by the ceiling rather if matches!(isolation, IsolationLevel::ReadCommitted) { diff --git a/scripts/dev/upgrade-test.py b/scripts/dev/upgrade-test.py index c787f3f..d865de5 100644 --- a/scripts/dev/upgrade-test.py +++ b/scripts/dev/upgrade-test.py @@ -120,9 +120,11 @@ def offset_for_time(topic, partition, when_ms): def timestamps(topic, partition): + """(offset, timestamp) for every record.""" out = kcat("-C", "-t", f"{topic}", "-p", str(partition), "-o", "beginning", "-e", "-q", - "-f", "%T\n") - return [int(t) for t in out.stdout.split() if t.strip().isdigit()] + "-f", "%o %T\n") + return [tuple(map(int, line.split())) for line in out.stdout.splitlines() + if len(line.split()) == 2] def seed(engine): @@ -161,8 +163,10 @@ def seed(engine): def snapshot(engine): topics = ["up-plain", "up-compact"] + (["up-sql"] if engine == "segment" else []) ts = timestamps("up-plain", 0) - probe = ts[len(ts) // 2] + probe = ts[len(ts) // 2][1] return { + # The first record at or after the probe, which is Kafka's answer. + "exact": next(o for o, t in ts if t >= probe), "records": {t: read_topic(t) for t in topics}, "uncommitted": read_topic("up-sql", "read_uncommitted") if engine == "segment" else [], "groups": sql("""SELECT string_agg(partition || ':' || committed_offset, ',' @@ -173,16 +177,25 @@ def snapshot(engine): } +def latest(rows): + """Each key's last record: all that compaction keeps, whenever the broker's own pass runs.""" + return {k: (o, v) for _, o, k, v in rows} + + def compare(before, after, when): for topic, rows in before["records"].items(): got = after["records"][topic] + if topic == "up-compact": + assert latest(got) == latest(rows), f"{when}: {topic} changed a key's last record" + continue assert got == rows, f"{when}: {topic} changed: {len(rows)} records before, {len(got)} after" assert after["uncommitted"] == before["uncommitted"], f"{when}: read_uncommitted view changed" assert after["groups"] == before["groups"], ( f"{when}: group offsets {before['groups']} became {after['groups']}") - assert after["lookup"] == before["lookup"], ( - f"{when}: offsetsForTimes({before['probe']}) was {before['lookup']}, " - f"now {after['lookup']}") + # Against the exact answer: the old image may answer with the start of the batch. + assert after["lookup"] == before["exact"], ( + f"{when}: offsetsForTimes({before['probe']}) is {after['lookup']}, " + f"not {before['exact']}") def epoch(): @@ -230,6 +243,12 @@ def main(): assert new_version != args.old, f"still {new_version}" assert sql("SELECT has_function_privilege('public', 'kafgres_drop_topic(text)', " "'EXECUTE')") == "f", "the update left the admin functions public" + assert sql("SELECT count(*) FROM pg_proc WHERE proname = 'kafgres_follower_lag'") == "1", \ + "the update did not add kafgres_follower_lag()" + assert sql("SELECT count(*) FROM kafgres_follower_lag()") == "0" + assert sql("SELECT has_function_privilege('public', 'kafgres_slot_ready(int)', " + "'EXECUTE')") == "f", "the update left kafgres_slot_ready() public" + assert sql("SELECT kafgres_slot_ready(0)") == "slot not ready" restart() after = snapshot(args.engine) for topic in ("up-plain", "up-sql"): @@ -237,16 +256,18 @@ def main(): missing = set(before["records"][topic]) - set(after["records"][topic]) assert not missing, f"{topic} lost {len(missing)} records after the update" assert after["groups"] == before["groups"] - assert after["lookup"] == before["lookup"] + assert after["lookup"] == before["exact"] step("compacting what 0.2.0 wrote") sql("SELECT kafgres_enforce_retention()") compacted = read_topic("up-compact") - latest = {k: o for _, o, k, _ in compacted} - assert len(latest) == 50, f"a key was lost to compaction: {len(latest)} keys" - assert all(o >= 1500 for o in latest.values()), "a superseded value outlived its key's last" - assert len(compacted) < len(before["records"]["up-compact"]), "nothing was compacted" - print(f" up-compact: {len(before['records']['up-compact'])} -> {len(compacted)} records") + kept = latest(compacted) + assert len(kept) == 50, f"a key was lost to compaction: {len(kept)} keys" + assert all(o >= 1500 for o, _ in kept.values()), "a superseded value outlived its key's last" + assert kept == latest(before["records"]["up-compact"]), "compaction changed a key's last record" + # Against what was written, as the broker's own pass may have run before the snapshot. + assert len(compacted) < 2000, "nothing was compacted" + print(f" up-compact: 2000 -> {len(compacted)} records") bad = [l for l in logs().splitlines() if ("PANIC" in l or "FATAL" in l or "ERROR" in l) and "kafgres" in l] diff --git a/tests/integration/cdc/test_cdc_source.py b/tests/integration/cdc/test_cdc_source.py index c1616e8..8a9d29c 100644 --- a/tests/integration/cdc/test_cdc_source.py +++ b/tests/integration/cdc/test_cdc_source.py @@ -111,6 +111,132 @@ def start_slot(): def values(rows): return [json.loads(v) for _, _, v in rows] +def test_a_change_right_after_add_mapping_is_captured(): + add_mapping("jsonb_build_object('id', new.id)") + sql("INSERT INTO src_orders VALUES (10,1,99.5,'placed','x')") + assert drain() == 1 + assert len(consume("src-events")) == 1 + +def slot_count(): + return int(sql("SELECT count(*) FROM pg_replication_slots WHERE slot_name = 'kafgres_cdc'")) + +def test_a_mapping_added_after_the_transaction_wrote_is_saved_with_a_warning(): + out = psql("BEGIN; INSERT INTO src_orders VALUES (1,1,1,'x','x'); " + "SELECT kafgres_add_mapping('m1', 'public.src_orders', 'src-events', " + "$$jsonb_build_object('id', new.id)$$, NULL, NULL); COMMIT") + assert out.returncode == 0, out.stderr[-500:] + assert "WARNING" in out.stderr, out.stderr[-500:] + assert sql("SELECT count(*) FROM kafgres_cdc_mappings WHERE mapping_name = 'm1'") == "1" + +def test_a_drain_is_not_held_up_by_a_mapping_transaction_left_open(): + add_mapping("jsonb_build_object('id', new.id)", name="m0", source="public.src_customers") + adding = subprocess.Popen( + ["docker", "compose", "exec", "-T", "postgres", "psql", "-U", "postgres", "-c", + "BEGIN; SELECT kafgres_add_mapping('m1', 'public.src_orders', 'src-events', " + "$$jsonb_build_object('id', new.id)$$, 'new.id::text', NULL); " + "SELECT pg_sleep(5); COMMIT"], + cwd=REPO, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + try: + time.sleep(1.5) + sql("INSERT INTO src_customers VALUES (3,'bronze')") + t0 = time.monotonic() + assert drain() == 1 + assert time.monotonic() - t0 < 3, "the drain waited for the open transaction" + finally: + adding.wait(timeout=30) + sql("INSERT INTO src_orders VALUES (10,1,99.5,'placed','x')") + assert drain() == 1, "a change after the mapping committed was not captured" + +def test_a_mapping_does_not_wait_for_a_running_transaction_to_create_the_slot(): + holder = subprocess.Popen( + ["docker", "compose", "exec", "-T", "postgres", "psql", "-U", "postgres", "-c", + "BEGIN; SELECT pg_current_xact_id(); SELECT pg_sleep(8); COMMIT"], + cwd=REPO, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + try: + time.sleep(1.5) + t0 = time.monotonic() + out = psql("SELECT kafgres_add_mapping('m1', 'public.src_orders', 'src-events', " + "$$jsonb_build_object('id', new.id)$$, 'new.id::text', NULL)") + assert out.returncode == 0, out.stderr[-500:] + assert time.monotonic() - t0 < 4, "the mapping waited for the open transaction" + assert "WARNING" in out.stderr, out.stderr[-500:] + assert "kafgres_slot_ready" in out.stderr, "the warning does not say how to wait for it" + assert sql("SELECT count(*) FROM kafgres_cdc_mappings WHERE mapping_name = 'm1'") == "1" + assert sql("SELECT kafgres_slot_ready(0)") == "slot not ready" + finally: + holder.wait(timeout=30) + start_slot() + assert sql("SELECT kafgres_slot_ready(1000)") == "slot ready" + sql("INSERT INTO src_orders VALUES (10,1,99.5,'placed','x')") + assert drain() == 1 + +def test_a_mapping_creates_the_slot_among_short_transactions(): + busy = subprocess.Popen( + ["docker", "compose", "exec", "-T", "postgres", "psql", "-U", "postgres", "-c", + "DO $$ BEGIN FOR i IN 1..120 LOOP PERFORM pg_current_xact_id(); " + "PERFORM pg_sleep(0.05); COMMIT; END LOOP; END $$"], + cwd=REPO, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + try: + time.sleep(1) + out = psql("SELECT kafgres_add_mapping('m1', 'public.src_orders', 'src-events', " + "$$jsonb_build_object('id', new.id)$$, 'new.id::text', NULL)") + assert out.returncode == 0, out.stderr[-500:] + assert "WARNING" not in out.stderr, out.stderr[-500:] + assert sql("SELECT kafgres_slot_ready(0)") == "slot ready" + finally: + busy.wait(timeout=30) + sql("INSERT INTO src_orders VALUES (10,1,99.5,'placed','x')") + assert drain() == 1 + +def test_slot_ready_answers_at_once_in_a_transaction_that_wrote(): + t0 = time.monotonic() + # Last, as psql before 15 prints only the last result; the transaction rolls back on exit. + out = psql("BEGIN; INSERT INTO src_orders VALUES (1,1,1,'x','x'); " + "SELECT kafgres_add_mapping('m1', 'public.src_orders', 'src-events', " + "$$jsonb_build_object('id', new.id)$$, NULL, NULL); " + "SELECT kafgres_slot_ready()") + assert out.returncode == 0, out.stderr[-500:] + assert time.monotonic() - t0 < 5, "kafgres_slot_ready waited on its own transaction" + assert "slot not ready" in out.stdout + assert "after it commits" in out.stderr, out.stderr[-500:] + + +def test_a_slot_no_mapping_wants_still_advances(): + """Left after the last mapping goes, it holds no WAL while the drain runs.""" + add_mapping("jsonb_build_object('id', new.id)") + sql("SELECT kafgres_remove_mapping('m1')") + assert slot_count() == 1 + lsn = "SELECT confirmed_flush_lsn FROM pg_replication_slots WHERE slot_name = 'kafgres_cdc'" + before = sql(lsn) + sql("INSERT INTO src_orders SELECT g, 1, 1, 'x', 'x' FROM generate_series(1, 500) g") + assert drain() == 0 + assert sql(lsn) != before, "the slot did not advance with no mapping" + +def test_a_role_without_replication_saves_a_mapping_with_a_warning(): + sql("DROP ROLE IF EXISTS mapper") + sql("CREATE ROLE mapper") + try: + sql("GRANT EXECUTE ON FUNCTION kafgres_add_mapping(text,text,text,text,text,text) TO mapper") + sql("GRANT SELECT ON kafgres_topics, src_orders TO mapper") + sql("GRANT SELECT, INSERT, UPDATE ON kafgres_cdc_mappings TO mapper") + out = psql("SET ROLE mapper; SELECT kafgres_add_mapping('m1', 'public.src_orders', " + "'src-events', $$jsonb_build_object('id', new.id)$$, NULL, NULL)") + assert out.returncode == 0, out.stderr[-500:] + assert "REPLICATION" in out.stderr, out.stderr[-500:] + assert sql("SELECT count(*) FROM kafgres_cdc_mappings WHERE mapping_name = 'm1'") == "1" + assert slot_count() == 0 + finally: + sql("DELETE FROM kafgres_cdc_mappings") + sql("REVOKE ALL ON kafgres_topics, src_orders, kafgres_cdc_mappings FROM mapper") + sql("REVOKE ALL ON FUNCTION kafgres_add_mapping(text,text,text,text,text,text) FROM mapper") + sql("DROP ROLE mapper") + +def test_a_mapping_to_a_missing_topic_leaves_no_slot(): + out = psql("SELECT kafgres_add_mapping('m1', 'public.src_orders', 'no-such-topic', " + "$$jsonb_build_object('id', new.id)$$, NULL, NULL)") + assert out.returncode != 0 + assert slot_count() == 0 + def test_an_insert_reaches_the_topic_through_the_wal(): add_mapping("jsonb_build_object('id', new.id, 'status', new.status, 'op', op)") start_slot() diff --git a/tests/integration/groups/test_consumer_groups.py b/tests/integration/groups/test_consumer_groups.py index 4fd1902..7014e1b 100644 --- a/tests/integration/groups/test_consumer_groups.py +++ b/tests/integration/groups/test_consumer_groups.py @@ -33,13 +33,13 @@ def kafka_tool(script, *args, timeout=240): capture_output=True, text=True, timeout=timeout, ) -def start_consumer(name, group, topic): +def start_consumer(name, group, topic, *extra): subprocess.run(["docker", "rm", "-f", name], capture_output=True) subprocess.run( ["docker", "run", "--rm", "-d", "--network", "host", "--name", name, CLIENTS, "kcat", "-b", BROKER, "-G", group, topic, "-X", f"session.timeout.ms={SESSION_MS}", - "-X", f"heartbeat.interval.ms={HEARTBEAT_MS}"], + "-X", f"heartbeat.interval.ms={HEARTBEAT_MS}", *extra], capture_output=True, timeout=120, ) @@ -149,6 +149,47 @@ def test_two_consumers_split_the_partitions(topic, group): finally: stop_consumer("p3a") +def test_a_joiner_with_no_common_protocol_is_refused_alone(topic, group): + make(topic, partitions=2) + start_consumer("p3c", group, topic) + try: + await_group(group, members=1, state="Stable") + _, generation, _ = group_state(group) + subprocess.run(["docker", "rm", "-f", "p3d"], capture_output=True) + subprocess.run( + ["docker", "run", "--rm", "-d", "--network", "host", "--name", "p3d", CLIENTS, + "kcat", "-b", BROKER, "-G", group, topic, + "-X", "partition.assignment.strategy=cooperative-sticky"], + capture_output=True, timeout=120, + ) + try: + time.sleep(15) + state, now, members = group_state(group) + assert (state, members) == ("Stable", 1), (state, members) + assert now == generation, "the existing member must not be rebalanced" + finally: + stop_consumer("p3d") + finally: + stop_consumer("p3c") + +def test_a_restarted_static_member_may_change_its_assignor(topic, group): + """The old registration of the same instance id is what the new one replaces.""" + make(topic, partitions=2) + start_consumer("p3e", group, topic, "-X", "group.instance.id=static-1", + "-X", "partition.assignment.strategy=range") + try: + await_group(group, members=1, state="Stable") + subprocess.run(["docker", "kill", "p3e"], capture_output=True, timeout=120) + start_consumer("p3f", group, topic, "-X", "group.instance.id=static-1", + "-X", "partition.assignment.strategy=cooperative-sticky") + try: + _, _, members = await_group(group, members=1, state="Stable", timeout=40) + assert members == 1 + finally: + stop_consumer("p3f") + finally: + stop_consumer("p3e") + def test_killing_a_consumer_rebalances_to_the_survivor(topic, group): """Second half. `docker kill`, not `stop`: a member that gets no chance to send LeaveGroup is the case that can only be recovered by the session-timeout sweep, and diff --git a/tests/integration/groups/test_kip848.py b/tests/integration/groups/test_kip848.py index 7ad70fd..609045f 100644 --- a/tests/integration/groups/test_kip848.py +++ b/tests/integration/groups/test_kip848.py @@ -141,3 +141,30 @@ def test_a_classic_group_id_is_refused_rather_than_forked(topic): assert "m0" not in out.stdout, ( f"refused the group but consumed from it anyway: {out.stdout[-400:]}" ) + +def test_a_consumer_group_id_is_refused_to_a_classic_join(topic): + """The other way round: a classic member joining a consumer group's id would be given + the same partitions, and the group's own members would be told it is gone.""" + name, grp = topic + produce(name, 10) + first = subprocess.run( + ["docker", "run", "--rm", "--network", "host", KAFKA, + "/opt/kafka/bin/kafka-console-consumer.sh", "--bootstrap-server", BROKER, + "--topic", name, "--from-beginning", "--group", grp, + "--consumer-property", "group.protocol=consumer", + "--timeout-ms", "12000", "--max-messages", "10"], + capture_output=True, text=True, timeout=120) + assert "m0" in first.stdout, "the consumer-protocol consumer did not run" + + out = subprocess.run( + ["docker", "run", "--rm", "--network", "host", KAFKA, + "/opt/kafka/bin/kafka-console-consumer.sh", "--bootstrap-server", BROKER, + "--topic", name, "--from-beginning", "--group", grp, + "--consumer-property", "group.protocol=classic", + "--timeout-ms", "12000"], + capture_output=True, text=True, timeout=120) + combined = out.stdout + out.stderr + assert "InconsistentGroupProtocol" in combined, ( + f"a classic join was allowed onto a consumer group id: {combined[-400:]}" + ) + assert "m0" not in out.stdout, f"refused the join but consumed anyway: {out.stdout[-400:]}" diff --git a/tests/integration/produce_fetch/test_idempotent_producer.py b/tests/integration/produce_fetch/test_idempotent_producer.py index a96e8c8..05fee92 100644 --- a/tests/integration/produce_fetch/test_idempotent_producer.py +++ b/tests/integration/produce_fetch/test_idempotent_producer.py @@ -802,3 +802,44 @@ def test_a_restart_forgets_window_rows_past_the_recovered_log(topic): pass time.sleep(2) raise AssertionError("a window row past the log end survived the restart") + +@pytest.mark.skipif(sql("SHOW kafgres.storage_engine") != "segment", + reason="the segment engine rebuilds its window from the log") +@pytest.mark.parametrize("records", [2, 1500]) +def test_a_batch_whose_window_row_was_lost_is_not_appended_twice(topic, conn, records): + """A crash can keep a batch in the log and lose its window row. After the restart the + first idempotent produce rebuilds the row from the log, so the client's retry is + answered as a duplicate even when other writes have landed in between. 1500 records + is more than the rebuild's nearest start looks back, so the scan starts inside it.""" + from conftest import BROKER_HOST, BROKER_PORT, Connection + import socket + + producer_id = init_producer_id(conn, 1) + batch = record_batch([b"a"] + [b"b"] * (records - 1), producer_id=producer_id, + producer_epoch=0, base_sequence=0) + _, err, first = send_produce(conn, topic, batch, 2) + assert err == 0 + sql(f"DELETE FROM kafgres_producer_batches WHERE producer_id = {producer_id}") + # Newer than the batch, so the rebuild's start by timestamp lands past it. + for i in range(3): + sql(f"SELECT kafgres_produce('{topic}', 'k', 'sql{i}')") + + subprocess.run(["docker", "compose", "restart", "postgres"], check=True, timeout=300, + capture_output=True) + deadline = time.time() + 120 + while True: + try: + sql("SELECT 1") + fresh = Connection(socket.create_connection((BROKER_HOST, BROKER_PORT), timeout=10)) + break + except (RuntimeError, OSError): + assert time.time() < deadline, "the broker did not come back" + time.sleep(2) + try: + _, err, again = send_produce(fresh, topic, batch, 3) + finally: + fresh.sock.close() + assert err == 0, f"the retry failed with {err}" + assert again == first, f"the retry was appended at {again}, not answered with {first}" + back = kcat("-t", topic, "-C", "-o", "beginning", "-e", "-f", "%s\n") + assert back.stdout.split().count("a") == 1, f"the batch was stored twice: {back.stdout.split()}" diff --git a/tests/integration/security/test_sql_produce_permissions.py b/tests/integration/security/test_sql_produce_permissions.py index 4e603e4..3062bd7 100644 --- a/tests/integration/security/test_sql_produce_permissions.py +++ b/tests/integration/security/test_sql_produce_permissions.py @@ -14,7 +14,7 @@ PUBLIC_FUNCTIONS = { "kafgres_produce", "kafgres_version", "kafgres_kafka_version", "kafgres_partition_offsets", "kafgres_archive_status", "kafgres_cdc_status", - "kafgres_cdc_snapshots", "kafgres_share_state", + "kafgres_cdc_snapshots", "kafgres_share_state", "kafgres_follower_lag", } def as_role(statement): diff --git a/tests/integration/storage/test_compaction.py b/tests/integration/storage/test_compaction.py index fbb3c09..409db7c 100644 --- a/tests/integration/storage/test_compaction.py +++ b/tests/integration/storage/test_compaction.py @@ -619,3 +619,99 @@ def test_batches_larger_than_segment_bytes_still_compact(byte_arm_topic): assert len(got) < 2000, f"{len(got)} of 2000 records readable: nothing was compacted" latest = {k: int(o) for o, k in got} assert len(latest) == 50, f"a key was lost: {len(latest)} keys" + +@pytest.fixture +def no_background_sweep(): + """Only the test's own passes run, so it can count what each one forgets.""" + sql("ALTER SYSTEM SET kafgres.retention_check_interval_ms = 3600000") + sql("SELECT pg_reload_conf()") + time.sleep(0.5) + yield + sql("ALTER SYSTEM RESET kafgres.retention_check_interval_ms") + sql("SELECT pg_reload_conf()") + +@pytest.mark.skipif(engine() != "segment", reason="kafgres_produce() is segment engine only") +def test_retention_forgets_a_large_marker_backlog_over_several_passes(topic, no_background_sweep): + """Markers below the log start are deleted a bounded number per pass, so a large + backlog goes over several passes instead of stalling the broker in one.""" + sql("ALTER SYSTEM SET kafgres.segment_offsets = 4000") + sql("ALTER SYSTEM SET kafgres.segment_bytes = 1048576") + sql("SELECT pg_reload_conf()") + time.sleep(0.5) + sql(f"SELECT count(kafgres_produce('{topic}', 'k', i::text)) FROM generate_series(1, 45000) i") + below = (f"""SELECT count(*) FROM kafgres_markers m + JOIN kafgres_topics t ON t.topic_id = m.topic_id + JOIN kafgres_partitions p ON p.topic_id = m.topic_id AND p.partition = m.partition + WHERE t.name = '{topic}' AND m.last_offset < p.log_start_offset""") + log_start = f"""SELECT p.log_start_offset FROM kafgres_partitions p + JOIN kafgres_topics t USING (topic_id) WHERE t.name = '{topic}'""" + assert kafka_tool("kafka-configs.sh", "--entity-type", "topics", "--entity-name", topic, + "--alter", "--add-config", "retention.ms=1").returncode == 0 + time.sleep(2) + compact() + moved = int(sql(log_start)) + assert moved > 20_000, f"retention moved the log start only to {moved}" + assert int(sql(below)) == moved - 20_000, "the first pass did not forget a bounded number" + for _ in range(moved // 20_000): + compact() + assert int(sql(below)) == 0, "later passes did not forget the rest" + +@pytest.mark.skipif(engine() != "segment", reason="kafgres_produce() is segment engine only") +def test_a_committed_read_passes_a_long_run_of_markers_left_by_compaction(topic, no_background_sweep): + """Markers whose batches compaction removed, more than one read looks at, as a lost + cleanup leaves them: a read_committed consumer still reaches the records after them.""" + sql("ALTER SYSTEM SET kafgres.segment_offsets = 8000") + sql("ALTER SYSTEM SET kafgres.segment_bytes = 1048576") + sql("SELECT pg_reload_conf()") + time.sleep(0.5) + make_compacted(topic) + sql(f"SELECT count(kafgres_produce('{topic}', 'k', i::text)) FROM generate_series(1, 70000) i") + sql(f"SELECT count(kafgres_produce('{topic}', 'z' || i, i::text)) FROM generate_series(1, 9000) i") + ours = f"m.topic_id = (SELECT topic_id FROM kafgres_topics WHERE name = '{topic}')" + sql(f"DROP TABLE IF EXISTS saved_markers; " + f"CREATE TABLE saved_markers AS SELECT * FROM kafgres_markers m WHERE {ours}") + try: + for _ in range(3): + compact() + restored = int(sql("WITH r AS (INSERT INTO kafgres_markers SELECT * FROM saved_markers " + "ON CONFLICT DO NOTHING RETURNING 1) SELECT count(*) FROM r")) + assert restored > 65_536, f"compaction removed only {restored} batches" + got = read_back(topic) + assert "k" in {line.split()[1] for line in got}, "the newest k was not read" + assert len({line.split()[1] for line in got if line.split()[1].startswith("z")}) == 9000 + finally: + sql("DROP TABLE IF EXISTS saved_markers") + +@pytest.mark.skipif(engine() != "segment", reason="kafgres_produce() is segment engine only") +@pytest.mark.parametrize("config", ["retention.ms=-1", "cleanup.policy=compact"]) +def test_markers_below_a_delete_records_go_on_a_topic_kept_forever(topic, no_background_sweep, config): + """With no time or size limit retention reclaims nothing itself, but DeleteRecords still + moves the log start, and later passes forget the markers left below it.""" + import json + import tempfile + + assert kafka_tool("kafka-configs.sh", "--entity-type", "topics", "--entity-name", topic, + "--alter", "--add-config", config).returncode == 0 + # Distinct keys, so compaction itself forgets none. + sql(f"SELECT count(kafgres_produce('{topic}', i::text, i::text)) FROM generate_series(1, 25000) i") + spec = {"partitions": [{"topic": topic, "partition": 0, "offset": 24000}], "version": 1} + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "offsets.json") + with open(path, "w") as f: + json.dump(spec, f) + os.chmod(d, 0o755) + os.chmod(path, 0o644) + out = subprocess.run( + ["docker", "run", "--rm", "--network", "host", "-v", f"{d}:/w", KAFKA, + "/opt/kafka/bin/kafka-delete-records.sh", "--bootstrap-server", BROKER, + "--offset-json-file", "/w/offsets.json"], + capture_output=True, text=True, timeout=300, + ) + assert out.returncode == 0, out.stderr + below = f"""SELECT count(*) FROM kafgres_markers m + JOIN kafgres_topics t ON t.topic_id = m.topic_id + WHERE t.name = '{topic}' AND m.last_offset < 24000""" + compact() + assert int(sql(below)) == 4000, "a retention pass did not forget a bounded number" + compact() + assert int(sql(below)) == 0, "the next pass did not forget the rest"