From 380dd4c530b433c8a0dbd736459d061773010350 Mon Sep 17 00:00:00 2001 From: Michael Netshipise Date: Fri, 17 Jul 2026 07:04:14 +0200 Subject: [PATCH] fix(lock): expose authoritative renewal lifecycle --- VERSION | 2 +- rust/Cargo.toml | 2 +- rust/README.md | 22 +- rust/src/lib.rs | 1 + rust/src/lock/mod.rs | 763 ++++++++++++++++++++++++++++++------ rust/src/stream/consumer.rs | 9 +- rust/src/stream/mod.rs | 9 +- ts/package-lock.json | 4 +- ts/package.json | 2 +- 9 files changed, 681 insertions(+), 133 deletions(-) diff --git a/VERSION b/VERSION index baec65a..db7a480 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.1.28 +0.1.31 diff --git a/rust/Cargo.toml b/rust/Cargo.toml index ce4ef08..a05a298 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "waymaker-client" -version = "0.1.28" +version = "0.1.31" edition = "2021" description = "Official Rust client for waymaker — locks, streams, KV, collections, sketches, cache, object store" repository = "https://git.awesomike.com/pub/waymaker-client" diff --git a/rust/README.md b/rust/README.md index ac41f1a..e996828 100644 --- a/rust/README.md +++ b/rust/README.md @@ -4,7 +4,7 @@ Official Rust client for [waymaker](https://git.awesomike.com/dev/waymaker). ```toml [dependencies] -waymaker-client = { git = "https://git.awesomike.com/pub/waymaker-client", tag = "v0.1.27" } +waymaker-client = { git = "https://git.awesomike.com/pub/waymaker-client", tag = "v0.1.31" } # (the crate lives in the rust/ subdir; cargo resolves it automatically) ``` @@ -17,7 +17,10 @@ dependency on the waymaker server workspace. keeps a background task that **transparently re-binds** its event stream across a primary bounce (reusing the `request_id`) and exposes live state: `lock.fence_token()`, `lock.watch()` (a `watch::Receiver`), - `lock.is_lost()`. Re-read the fence before every fenced side effect. + `lock.is_lost()`. `LockState::generation` distinguishes a same-generation + re-bind from a fresh transparent acquire, while `LockState::renewal` exposes + typed renewal success/failure and the authoritative renewed deadline. + Pin the id/token/generation before protected work and reject any change. - `stream` — JetStream-lite publish + pull/push consumers. - `kv` — Put/Get/Create/Update(CAS)/Delete/Keys/History/Watch. - `cache` — Redis-shape Hash / Set / Queue (collections). @@ -42,11 +45,20 @@ let lock = client.acquire_lock("leader:reports", lock::Config { }).await?; let _renewal = lock.spawn_renewal(Duration::from_secs(15)); -// Re-read the fence before each fenced write; stop if leadership is lost. +// Pin one ownership generation; a transparent re-win cannot inherit work. let mut state = lock.watch(); +let admitted = state.borrow().clone(); loop { - if lock.is_lost() { break; } - do_fenced_write(lock.fence_token()).await?; + let current = state.borrow().clone(); + if current.lost + || current.generation != admitted.generation + || current.id != admitted.id + || current.fence_token != admitted.fence_token + || matches!(current.renewal, lock::RenewalState::Failed(_)) + { + break; + } + do_fenced_write(admitted.fence_token).await?; state.changed().await.ok(); } ``` diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 81cb610..f6539aa 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -23,6 +23,7 @@ pub mod server { /// Streams proto stubs (`WaymakerStreamsService` + internal /// replication RPCs). Generated from `proto/waymaker_streams.proto`. pub mod streams_server { + #![allow(clippy::doc_lazy_continuation)] tonic::include_proto!("waymaker.streams"); } diff --git a/rust/src/lock/mod.rs b/rust/src/lock/mod.rs index b96fca7..e972866 100644 --- a/rust/src/lock/mod.rs +++ b/rust/src/lock/mod.rs @@ -35,11 +35,12 @@ use crate::client::Client; use crate::error::{Error, Result}; use crate::server::{ - waymaker_service_client::WaymakerServiceClient, ExtendLeaseRequest, LeaseStatusRequest, - LockEvent, LockEventType, LockRequest, MultiLockKey, MultiLockRequest, UnLockRequest, + waymaker_service_client::WaymakerServiceClient, ExtendLeaseRequest, ExtendLeaseResponse, + LeaseStatusRequest, LockEvent, LockEventType, LockRequest, MultiLockKey, MultiLockRequest, + UnLockRequest, }; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::{watch, Notify}; use tonic::transport::Channel; use tonic::Request; @@ -65,12 +66,13 @@ use uuid::Uuid; /// durably committed. Even [`Scope::Quorum`] does not let you skip /// that check — an all-at-once cluster restart can still drop an /// in-memory token. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum Scope { /// Counter in RAM on the owning node; resets on that node's /// restart or a hash-ring rebalance. Fastest (no I/O). Right for /// advisory locks / rate limiting where a fence reset across /// failure is tolerable. + #[default] Ephemeral, /// Counter persisted to disk on the owning node. Survives a /// process restart on the same node; still resets if the ring @@ -91,12 +93,6 @@ pub enum Scope { Quorum, } -impl Default for Scope { - fn default() -> Self { - Self::Ephemeral - } -} - impl Scope { fn to_pb(self) -> i32 { use crate::server::FenceScope as F; @@ -207,12 +203,39 @@ pub struct LockState { pub fence_token: u64, /// Lease expiry (epoch ms), refreshed from heartbeats / re-acquire. pub lease_expires_at_ms: i64, + /// Monotonic local ownership-generation number. A same-generation event + /// stream re-bind preserves this value; a transparent fresh acquire with + /// a different id/token increments it. Callers can pin this number to + /// reject work that spans a transparent re-win. + pub generation: u64, + /// Result of the most recent background renewal attempt. Renewal success + /// advances `lease_expires_at_ms` from the authoritative response; a + /// failure leaves the last proven deadline unchanged. + pub renewal: RenewalState, /// `true` once the client could no longer prove it holds the lock /// (the lease was taken by someone else, or expired and could not /// be re-won). A `lost` holder MUST stop acting as the holder. pub lost: bool, } +/// Observable outcome of the most recent background lease renewal. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RenewalState { + NotStarted, + Succeeded, + Failed(RenewalFailure), +} + +/// Failure classes are deliberately stable so callers and telemetry do not +/// have to parse transport or server error strings. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RenewalFailure { + Timeout, + Transport, + Rejected, + Malformed, +} + /// An acquired lock. Dropping the handle **does not** auto-release /// the lock — call [`Lock::unlock`] explicitly (or let the lease /// expire). This matches the underlying RPC semantics; auto-release @@ -232,6 +255,8 @@ pub struct Lock { pub key: String, /// Live lock state. `borrow()` yields the current snapshot. state: watch::Receiver, + state_tx: watch::Sender, + scope: Scope, /// Signals the background hold task to stop. Set by `unlock` and /// `Drop` so the task never re-acquires a deliberately-released /// lock. @@ -239,6 +264,225 @@ pub struct Lock { _hold: tokio::task::JoinHandle<()>, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct ValidatedGrant { + id: String, + fence_token: u64, + lease_expires_at_ms: i64, + acquired_at_ms: i64, + priority: u32, +} + +fn unix_now_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis().min(i64::MAX as u128) as i64) + .unwrap_or(i64::MAX) +} + +fn validate_grant_event( + event: &LockEvent, + expected_key: &str, + scope: Scope, + now_ms: i64, +) -> Result { + if !event.success { + return Err(Error::server( + "malformed_acquired", + "Acquired event has success=false", + )); + } + validate_generation( + &event.key, + &event.id, + event.fence_token, + event.lease_expires_at, + event.acquired_at, + 0, + expected_key, + None, + None, + scope, + now_ms, + ) +} + +#[allow(clippy::too_many_arguments)] +fn validate_generation( + key: &str, + id: &str, + fence_token: u64, + lease_expires_at_ms: i64, + acquired_at_ms: i64, + priority: u32, + expected_key: &str, + expected_id: Option<&str>, + expected_fence_token: Option, + scope: Scope, + now_ms: i64, +) -> Result { + if key != expected_key { + return Err(Error::server( + "malformed_lease", + format!("lease key {key:?} does not match requested key {expected_key:?}"), + )); + } + if id.is_empty() { + return Err(Error::server("malformed_lease", "lease id is empty")); + } + if expected_id.is_some_and(|expected| expected != id) { + return Err(Error::server( + "generation_changed", + "renewal returned a different lease id", + )); + } + if expected_fence_token.is_some_and(|expected| expected != fence_token) { + return Err(Error::server( + "generation_changed", + "renewal returned a different fence token", + )); + } + if scope == Scope::Quorum && fence_token == 0 { + return Err(Error::server( + "malformed_lease", + "quorum-fenced lease returned a zero fence token", + )); + } + if lease_expires_at_ms <= now_ms { + return Err(Error::server( + "malformed_lease", + "lease expiry is not in the future", + )); + } + Ok(ValidatedGrant { + id: id.to_string(), + fence_token, + lease_expires_at_ms, + acquired_at_ms, + priority, + }) +} + +fn replace_state(state_tx: &watch::Sender, next: LockState) { + state_tx.send_if_modified(|current| { + if *current == next { + false + } else { + *current = next.clone(); + true + } + }); +} + +fn publish_grant(state_tx: &watch::Sender, grant: ValidatedGrant) { + let current = state_tx.borrow().clone(); + let same_generation = current.id == grant.id && current.fence_token == grant.fence_token; + replace_state( + state_tx, + LockState { + id: grant.id, + fence_token: grant.fence_token, + lease_expires_at_ms: grant.lease_expires_at_ms, + generation: if same_generation { + current.generation + } else { + current.generation.saturating_add(1) + }, + renewal: if same_generation { + current.renewal + } else { + RenewalState::NotStarted + }, + lost: false, + }, + ); +} + +fn validate_renewal_response( + response: ExtendLeaseResponse, + expected_key: &str, + expected_id: &str, + expected_fence_token: u64, + scope: Scope, + now_ms: i64, +) -> Result { + if !response.success { + return Err(Error::server(response.result_code, response.message)); + } + let lease = response + .lease + .ok_or_else(|| Error::server("malformed_lease", "renewal response is missing its lease"))?; + validate_generation( + &lease.key, + &lease.id, + lease.fence_token, + lease.lease_expires_at, + lease.created_at, + lease.priority, + expected_key, + Some(expected_id), + Some(expected_fence_token), + scope, + now_ms, + ) +} + +fn renewal_failure(error: &Error) -> (RenewalFailure, bool) { + match error { + Error::Connect(_) | Error::Rpc(_) => (RenewalFailure::Transport, false), + Error::Server { code, .. } if code == "malformed_lease" || code == "generation_changed" => { + (RenewalFailure::Malformed, true) + } + Error::Server { .. } => (RenewalFailure::Rejected, false), + Error::Invalid(_) => (RenewalFailure::Malformed, true), + } +} + +fn publish_renewal_failure( + state_tx: &watch::Sender, + expected_generation: u64, + expected_id: &str, + expected_fence_token: u64, + failure: RenewalFailure, + fatal: bool, +) { + state_tx.send_if_modified(|current| { + if current.generation != expected_generation + || current.id != expected_id + || current.fence_token != expected_fence_token + { + return false; + } + let next = RenewalState::Failed(failure); + let changed = current.renewal != next || (fatal && !current.lost); + current.renewal = next; + current.lost |= fatal; + changed + }); +} + +fn publish_renewal_success( + state_tx: &watch::Sender, + expected_generation: u64, + expected_id: &str, + expected_fence_token: u64, + grant: ValidatedGrant, +) { + state_tx.send_if_modified(|current| { + if current.generation != expected_generation + || current.id != expected_id + || current.fence_token != expected_fence_token + { + return false; + } + let changed = current.lease_expires_at_ms != grant.lease_expires_at_ms + || current.renewal != RenewalState::Succeeded; + current.lease_expires_at_ms = grant.lease_expires_at_ms; + current.renewal = RenewalState::Succeeded; + changed + }); +} + impl Lock { /// Current lease id (live). pub fn id(&self) -> String { @@ -271,26 +515,56 @@ impl Lock { /// Extend the lease by `additional`. pub async fn extend(&self, additional: Duration) -> Result { + let expected = self.state.borrow().clone(); let mut c: WaymakerServiceClient = self.client.locks_client(); - let r = c + let response = c .extend_lease(Request::new(ExtendLeaseRequest { key: self.key.clone(), - id: self.id(), + id: expected.id.clone(), lease_timeout: additional.as_millis().min(u32::MAX as u128) as u32, })) - .await? - .into_inner(); - if !r.success { - return Err(Error::server(r.result_code, r.message)); - } - let lease = r.lease.ok_or_else(|| Error::server("internal", "missing lease in response"))?; + .await + .map_err(Error::from) + .map(|response| response.into_inner()) + .and_then(|response| { + validate_renewal_response( + response, + &self.key, + &expected.id, + expected.fence_token, + self.scope, + unix_now_ms(), + ) + }); + let grant = match response { + Ok(grant) => grant, + Err(error) => { + let (failure, fatal) = renewal_failure(&error); + publish_renewal_failure( + &self.state_tx, + expected.generation, + &expected.id, + expected.fence_token, + failure, + fatal, + ); + return Err(error); + } + }; + publish_renewal_success( + &self.state_tx, + expected.generation, + &expected.id, + expected.fence_token, + grant.clone(), + ); Ok(Lease { - id: lease.id, - key: lease.key, - acquired_at_ms: lease.created_at, - lease_expires_at_ms: lease.lease_expires_at, - fence_token: lease.fence_token, - priority: lease.priority, + id: grant.id, + key: self.key.clone(), + acquired_at_ms: grant.acquired_at_ms, + lease_expires_at_ms: grant.lease_expires_at_ms, + fence_token: grant.fence_token, + priority: grant.priority, }) } @@ -329,6 +603,8 @@ impl Lock { // mint a new id, and renewing the stale one would silently let the // real lease expire. let state = self.state.clone(); + let state_tx = self.state_tx.clone(); + let scope = self.scope; let ttl_ms = (every.as_millis() * 2).min(u32::MAX as u128) as u32; let stop = Arc::new(Notify::new()); let stop_task = stop.clone(); @@ -352,11 +628,11 @@ impl Lock { _ = ticker.tick() => { // Read the live id; skip if the lock is already // lost (extending a dead lease just errors). - let (id, lost) = { + let current = { let s = state.borrow(); - (s.id.clone(), s.lost) + s.clone() }; - if lost { + if current.lost { continue; } let mut c = client.locks_client(); @@ -365,20 +641,68 @@ impl Lock { // forever — without this the holder could never // release and every other waiter would skip until a // process restart. - let _ = tokio::time::timeout( + let result = tokio::time::timeout( every, c.extend_lease(Request::new(ExtendLeaseRequest { key: key.clone(), - id, + id: current.id.clone(), lease_timeout: ttl_ms, })), ) .await; + match result { + Err(_) => publish_renewal_failure( + &state_tx, + current.generation, + ¤t.id, + current.fence_token, + RenewalFailure::Timeout, + false, + ), + Ok(Err(_)) => publish_renewal_failure( + &state_tx, + current.generation, + ¤t.id, + current.fence_token, + RenewalFailure::Transport, + false, + ), + Ok(Ok(response)) => match validate_renewal_response( + response.into_inner(), + &key, + ¤t.id, + current.fence_token, + scope, + unix_now_ms(), + ) { + Ok(grant) => publish_renewal_success( + &state_tx, + current.generation, + ¤t.id, + current.fence_token, + grant, + ), + Err(error) => { + let (failure, fatal) = renewal_failure(&error); + publish_renewal_failure( + &state_tx, + current.generation, + ¤t.id, + current.fence_token, + failure, + fatal, + ); + } + }, + } } } } }); - RenewalHandle { stop, handle: Some(handle) } + RenewalHandle { + stop, + handle: Some(handle), + } } } @@ -414,7 +738,7 @@ enum Drained { enum Rebound { /// Got a live stream back (transparent re-bind, or a clean re-win /// after a real loss). Carries the fresh `(stream, id, fence, exp)`. - Bound(tonic::Streaming, String, u64, i64), + Bound(Box>, ValidatedGrant), /// The acquire did not grant (key held — possibly by our own /// still-adopted lease on a new primary, possibly by someone else) /// or the RPC failed. Disambiguate via lease_status. @@ -423,7 +747,7 @@ enum Rebound { /// Whether we still demonstrably own the lease. enum Ownership { - Held(i64), + Held(ValidatedGrant), Lost, Unknown, } @@ -432,17 +756,9 @@ const HOLD_BASE_BACKOFF: Duration = Duration::from_millis(200); const HOLD_MAX_BACKOFF: Duration = Duration::from_secs(10); const HOLD_RPC_TIMEOUT: Duration = Duration::from_secs(10); -/// Publish `next` if it differs from the local mirror `cur`, notifying -/// watchers only on a real change. -fn publish(tx: &watch::Sender, cur: &mut LockState, next: LockState) { - if *cur != next { - *cur = next.clone(); - let _ = tx.send_replace(next); - } -} - /// Background task body: hold the event stream, re-bind it on drop, and /// surface fence-token changes / loss through `state_tx`. +#[allow(clippy::too_many_arguments)] async fn hold_loop( client: Client, key: String, @@ -450,14 +766,11 @@ async fn hold_loop( reacquire_req: LockRequest, mut stream: tonic::Streaming, state_tx: watch::Sender, - init: LockState, + scope: Scope, stop: Arc, ) { - // Local mirror of the published state — lets us read the current id - // and publish only on change without depending on `Sender::borrow`. - let mut cur = init; loop { - match drain_stream(&mut stream, &state_tx, &mut cur, &stop).await { + match drain_stream(&mut stream, &state_tx, &key, scope, &stop).await { Drained::Stopped => return, Drained::Disconnected => {} } @@ -469,44 +782,54 @@ async fn hold_loop( // Race every step against `stop` so unlock/drop halts us // promptly — and, crucially, before any re-acquire could // resurrect a released lock. - let rebound = match race_stop(reacquire(&client, read, &reacquire_req), &stop).await { - Some(r) => r, - None => return, - }; + let rebound = + match race_stop(reacquire(&client, read, &reacquire_req, &key, scope), &stop).await + { + Some(r) => r, + None => return, + }; match rebound { - Rebound::Bound(s, id, fence, exp) => { - publish( - &state_tx, - &mut cur, - LockState { id, fence_token: fence, lease_expires_at_ms: exp, lost: false }, - ); - stream = s; + Rebound::Bound(s, grant) => { + publish_grant(&state_tx, grant); + stream = *s; break; // resume draining the fresh stream } Rebound::NotBound => { - let cur_id = cur.id.clone(); - let owned = - match race_stop(confirm_ownership(&client, &key, &cur_id), &stop).await { - Some(o) => o, - None => return, - }; + let current = state_tx.borrow().clone(); + let owned = match race_stop( + confirm_ownership(&client, &key, ¤t.id, current.fence_token, scope), + &stop, + ) + .await + { + Some(o) => o, + None => return, + }; match owned { - Ownership::Held(exp) => { + Ownership::Held(grant) => { // Still ours (renewal keeps the lease alive); // we just could not get an event stream — keep // monitoring with backoff. - let next = LockState { lease_expires_at_ms: exp, ..cur.clone() }; - publish(&state_tx, &mut cur, next); + publish_grant(&state_tx, grant); } Ownership::Unknown => { /* transient — back off, retry */ } Ownership::Lost => { - let next = LockState { lost: true, ..cur.clone() }; - publish(&state_tx, &mut cur, next); + state_tx.send_if_modified(|current| { + if current.lost { + false + } else { + current.lost = true; + true + } + }); return; } } - if race_stop(tokio::time::sleep(backoff), &stop).await.is_none() { + if race_stop(tokio::time::sleep(backoff), &stop) + .await + .is_none() + { return; } backoff = (backoff * 2).min(HOLD_MAX_BACKOFF); @@ -522,7 +845,8 @@ async fn hold_loop( async fn drain_stream( stream: &mut tonic::Streaming, state_tx: &watch::Sender, - cur: &mut LockState, + expected_key: &str, + scope: Scope, stop: &Notify, ) -> Drained { loop { @@ -533,17 +857,37 @@ async fn drain_stream( Ok(Some(ev)) => { let et = ev.event_type; if et == LockEventType::Heartbeat as i32 { - let next = LockState { lease_expires_at_ms: ev.lease_expires_at, ..cur.clone() }; - publish(state_tx, cur, next); + let current = state_tx.borrow().clone(); + let key_matches = ev.key.is_empty() || ev.key == expected_key; + let id_matches = ev.id.is_empty() || ev.id == current.id; + let fence_matches = ev.fence_token == 0 + || ev.fence_token == current.fence_token; + if !ev.success + || !key_matches + || !id_matches + || !fence_matches + { + return Drained::Disconnected; + } + // Connectivity heartbeats historically carry expiry=0 + // and fence=0. They prove only that the event stream is + // live; never let that zero erase the last deadline + // proven by acquire/status/renewal. + if ev.lease_expires_at > unix_now_ms() { + replace_state( + state_tx, + LockState { + lease_expires_at_ms: ev.lease_expires_at, + ..current + }, + ); + } } else if et == LockEventType::Acquired as i32 { // Re-emitted after an idempotent re-bind. - let next = LockState { - id: ev.id, - fence_token: ev.fence_token, - lease_expires_at_ms: ev.lease_expires_at, - lost: false, - }; - publish(state_tx, cur, next); + match validate_grant_event(&ev, expected_key, scope, unix_now_ms()) { + Ok(grant) => publish_grant(state_tx, grant), + Err(_) => return Drained::Disconnected, + } } else if et == LockEventType::Expired as i32 || et == LockEventType::Failed as i32 { @@ -563,7 +907,13 @@ async fn drain_stream( /// lease still held this re-binds the stream and re-emits the held /// entry; if the lease was lost and the key is free it re-wins with a /// fresh, higher token. -async fn reacquire(client: &Client, read: bool, req: &LockRequest) -> Rebound { +async fn reacquire( + client: &Client, + read: bool, + req: &LockRequest, + expected_key: &str, + scope: Scope, +) -> Rebound { let req = req.clone(); let mut c = client.locks_client(); let attempt = async move { @@ -575,15 +925,15 @@ async fn reacquire(client: &Client, read: bool, req: &LockRequest) -> Rebound { loop { match stream.message().await? { Some(ev) if ev.event_type == LockEventType::Acquired as i32 => { + let grant = match validate_grant_event(&ev, expected_key, scope, unix_now_ms()) + { + Ok(grant) => grant, + Err(_) => return Ok(None), + }; return Ok::< - Option<(tonic::Streaming, String, u64, i64)>, + Option<(tonic::Streaming, ValidatedGrant)>, tonic::Status, - >(Some(( - stream, - ev.id, - ev.fence_token, - ev.lease_expires_at, - ))); + >(Some((stream, grant))); } Some(ev) if ev.event_type == LockEventType::Failed as i32 @@ -591,13 +941,13 @@ async fn reacquire(client: &Client, read: bool, req: &LockRequest) -> Rebound { { return Ok(None); // not granted (contended) } - Some(_) => continue, // Waiting / Heartbeat + Some(_) => continue, // Waiting / Heartbeat None => return Ok(None), // closed before grant } } }; match tokio::time::timeout(HOLD_RPC_TIMEOUT, attempt).await { - Ok(Ok(Some((s, id, fence, exp)))) => Rebound::Bound(s, id, fence, exp), + Ok(Ok(Some((stream, grant)))) => Rebound::Bound(Box::new(stream), grant), Ok(Ok(None)) => Rebound::NotBound, Ok(Err(_)) | Err(_) => Rebound::NotBound, } @@ -606,7 +956,13 @@ async fn reacquire(client: &Client, read: bool, req: &LockRequest) -> Rebound { /// Query whether `id` still holds `key`. `success = false` (or a /// missing lease) means we lost it; a transport / timeout error is /// inconclusive (`Unknown`) so we never falsely declare loss. -async fn confirm_ownership(client: &Client, key: &str, id: &str) -> Ownership { +async fn confirm_ownership( + client: &Client, + key: &str, + id: &str, + fence_token: u64, + scope: Scope, +) -> Ownership { let mut c = client.locks_client(); let call = c.lease_status(Request::new(LeaseStatusRequest { key: key.to_string(), @@ -616,7 +972,21 @@ async fn confirm_ownership(client: &Client, key: &str, id: &str) -> Ownership { Ok(Ok(resp)) => { let r = resp.into_inner(); match (r.success, r.lease) { - (true, Some(lease)) => Ownership::Held(lease.lease_expires_at), + (true, Some(lease)) => validate_generation( + &lease.key, + &lease.id, + lease.fence_token, + lease.lease_expires_at, + lease.created_at, + lease.priority, + key, + Some(id), + Some(fence_token), + scope, + unix_now_ms(), + ) + .map(Ownership::Held) + .unwrap_or(Ownership::Lost), _ => Ownership::Lost, } } @@ -668,24 +1038,19 @@ impl Client { /// `LockEventType::Expired` (the contended wait-TTL elapsed) which this client /// surfaces as `Error::Server { code: "expired", .. }`. (There is no /// `"timeout"` code — callers that loop-until-granted must match `"expired"`.) - pub async fn acquire_lock( - &self, - key: impl Into, - config: Config, - ) -> Result { - self.acquire_lock_inner(key.into(), config, /* read = */ false).await + pub async fn acquire_lock(&self, key: impl Into, config: Config) -> Result { + self.acquire_lock_inner(key.into(), config, /* read = */ false) + .await } /// Acquire a shared (read) lock. - pub async fn acquire_read_lock( - &self, - key: impl Into, - config: Config, - ) -> Result { - self.acquire_lock_inner(key.into(), config, /* read = */ true).await + pub async fn acquire_read_lock(&self, key: impl Into, config: Config) -> Result { + self.acquire_lock_inner(key.into(), config, /* read = */ true) + .await } async fn acquire_lock_inner(&self, key: String, config: Config, read: bool) -> Result { + let scope = config.scope; let req = config.into_request(key.clone()); // Template for transparent re-acquire after a stream drop: same // request_id (so a still-held lease is recovered idempotently @@ -716,10 +1081,13 @@ impl Client { }; match event.event_type { t if t == LockEventType::Acquired as i32 => { + let grant = validate_grant_event(&event, &key, scope, unix_now_ms())?; let init = LockState { - id: event.id, - fence_token: event.fence_token, - lease_expires_at_ms: event.lease_expires_at, + id: grant.id, + fence_token: grant.fence_token, + lease_expires_at_ms: grant.lease_expires_at_ms, + generation: 1, + renewal: RenewalState::NotStarted, lost: false, }; let (state_tx, state_rx) = watch::channel(init.clone()); @@ -733,14 +1101,16 @@ impl Client { read, reacquire_req, stream, - state_tx, - init, + state_tx.clone(), + scope, stop.clone(), )); return Ok(Lock { client: self.clone(), key, state: state_rx, + state_tx, + scope, stop, _hold: hold, }); @@ -773,7 +1143,9 @@ impl Client { if !r.success { return Err(Error::server(r.result_code, r.message)); } - let lease = r.lease.ok_or_else(|| Error::server("internal", "missing lease in response"))?; + let lease = r + .lease + .ok_or_else(|| Error::server("internal", "missing lease in response"))?; Ok(Lease { id: lease.id, key: lease.key, @@ -834,3 +1206,176 @@ impl Client { .collect()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::server::Lease as WireLease; + + const NOW_MS: i64 = 1_000; + + fn acquired(key: &str, id: &str, fence_token: u64) -> LockEvent { + LockEvent { + success: true, + event_type: LockEventType::Acquired as i32, + message: String::new(), + id: id.to_string(), + key: key.to_string(), + lease_expires_at: NOW_MS + 10_000, + acquired_at: NOW_MS - 10, + waiting_expires_at: 0, + fence_token, + } + } + + fn renewal(key: &str, id: &str, fence_token: u64, expires_at: i64) -> ExtendLeaseResponse { + ExtendLeaseResponse { + success: true, + result_code: String::new(), + message: String::new(), + lease: Some(WireLease { + id: id.to_string(), + key: key.to_string(), + acquired: true, + priority: 7, + created_at: NOW_MS - 10, + lease_expires_at: expires_at, + waiting_expires_at: 0, + fence_token, + }), + } + } + + fn state(id: &str, fence_token: u64, generation: u64) -> LockState { + LockState { + id: id.to_string(), + fence_token, + lease_expires_at_ms: NOW_MS + 5_000, + generation, + renewal: RenewalState::NotStarted, + lost: false, + } + } + + #[test] + fn acquired_event_validation_fails_closed() { + let mut event = acquired("key", "lease-1", 9); + assert!(validate_grant_event(&event, "key", Scope::Quorum, NOW_MS).is_ok()); + + event.success = false; + assert!(validate_grant_event(&event, "key", Scope::Quorum, NOW_MS).is_err()); + event.success = true; + event.key = "other".into(); + assert!(validate_grant_event(&event, "key", Scope::Quorum, NOW_MS).is_err()); + event.key = "key".into(); + event.id.clear(); + assert!(validate_grant_event(&event, "key", Scope::Quorum, NOW_MS).is_err()); + event.id = "lease-1".into(); + event.fence_token = 0; + assert!(validate_grant_event(&event, "key", Scope::Quorum, NOW_MS).is_err()); + assert!(validate_grant_event(&event, "key", Scope::Ephemeral, NOW_MS).is_ok()); + event.lease_expires_at = NOW_MS; + assert!(validate_grant_event(&event, "key", Scope::Ephemeral, NOW_MS).is_err()); + } + + #[test] + fn renewal_must_preserve_the_acquired_generation() { + let valid = validate_renewal_response( + renewal("key", "lease-1", 9, NOW_MS + 20_000), + "key", + "lease-1", + 9, + Scope::Quorum, + NOW_MS, + ) + .unwrap(); + assert_eq!(valid.id, "lease-1"); + assert_eq!(valid.fence_token, 9); + assert_eq!(valid.priority, 7); + + for response in [ + renewal("other", "lease-1", 9, NOW_MS + 20_000), + renewal("key", "lease-2", 9, NOW_MS + 20_000), + renewal("key", "lease-1", 10, NOW_MS + 20_000), + renewal("key", "lease-1", 9, NOW_MS), + ] { + assert!(validate_renewal_response( + response, + "key", + "lease-1", + 9, + Scope::Quorum, + NOW_MS, + ) + .is_err()); + } + + let mut missing = renewal("key", "lease-1", 9, NOW_MS + 20_000); + missing.lease = None; + assert!( + validate_renewal_response(missing, "key", "lease-1", 9, Scope::Quorum, NOW_MS,) + .is_err() + ); + } + + #[test] + fn rebind_preserves_generation_but_rewin_increments_it() { + let (state_tx, _state_rx) = watch::channel(state("lease-1", 9, 4)); + publish_grant( + &state_tx, + ValidatedGrant { + id: "lease-1".into(), + fence_token: 9, + lease_expires_at_ms: NOW_MS + 20_000, + acquired_at_ms: NOW_MS - 10, + priority: 0, + }, + ); + assert_eq!(state_tx.borrow().generation, 4); + + publish_grant( + &state_tx, + ValidatedGrant { + id: "lease-2".into(), + fence_token: 10, + lease_expires_at_ms: NOW_MS + 30_000, + acquired_at_ms: NOW_MS, + priority: 0, + }, + ); + let current = state_tx.borrow(); + assert_eq!(current.generation, 5); + assert_eq!(current.id, "lease-2"); + assert_eq!(current.fence_token, 10); + } + + #[test] + fn renewal_outcomes_update_only_the_expected_generation() { + let (state_tx, _state_rx) = watch::channel(state("lease-1", 9, 4)); + publish_renewal_success( + &state_tx, + 4, + "lease-1", + 9, + ValidatedGrant { + id: "lease-1".into(), + fence_token: 9, + lease_expires_at_ms: NOW_MS + 20_000, + acquired_at_ms: NOW_MS - 10, + priority: 0, + }, + ); + assert_eq!(state_tx.borrow().renewal, RenewalState::Succeeded); + assert_eq!(state_tx.borrow().lease_expires_at_ms, NOW_MS + 20_000); + + publish_renewal_failure(&state_tx, 3, "lease-1", 9, RenewalFailure::Malformed, true); + assert!(!state_tx.borrow().lost); + + publish_renewal_failure(&state_tx, 4, "lease-1", 9, RenewalFailure::Malformed, true); + assert!(state_tx.borrow().lost); + assert_eq!( + state_tx.borrow().renewal, + RenewalState::Failed(RenewalFailure::Malformed) + ); + } +} diff --git a/rust/src/stream/consumer.rs b/rust/src/stream/consumer.rs index 3b16192..a29ea8a 100644 --- a/rust/src/stream/consumer.rs +++ b/rust/src/stream/consumer.rs @@ -281,9 +281,10 @@ impl tokio_stream::Stream for Messages { } /// Consumer delivery start policy. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum DeliverPolicy { /// Start from the first message in the stream. + #[default] All, /// Start from messages published after subscription (emulated /// via `ByStartTime(now)` at consumer creation since waymaker @@ -297,12 +298,6 @@ pub enum DeliverPolicy { ByStartTime(i64), } -impl Default for DeliverPolicy { - fn default() -> Self { - Self::All - } -} - impl DeliverPolicy { fn into_pb(self) -> crate::streams_server::DeliveryPolicyPb { use crate::streams_server::DeliveryPolicyType as T; diff --git a/rust/src/stream/mod.rs b/rust/src/stream/mod.rs index 6a62dc1..87ab60a 100644 --- a/rust/src/stream/mod.rs +++ b/rust/src/stream/mod.rs @@ -197,19 +197,14 @@ pub struct PublishAck { } /// Retention policy. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum RetentionPolicy { + #[default] Limits, WorkQueue, Interest, } -impl Default for RetentionPolicy { - fn default() -> Self { - Self::Limits - } -} - /// Stream configuration. `Default` produces an unbounded Limits /// stream with no subject filter, suitable for /// `StreamConfig { name: ..., ..Default::default() }`. diff --git a/ts/package-lock.json b/ts/package-lock.json index 2f5ee83..99ba49b 100644 --- a/ts/package-lock.json +++ b/ts/package-lock.json @@ -1,12 +1,12 @@ { "name": "@waymaker/client", - "version": "0.1.27", + "version": "0.1.31", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@waymaker/client", - "version": "0.1.27", + "version": "0.1.31", "license": "MIT OR Apache-2.0", "dependencies": { "@grpc/grpc-js": "^1.12.0", diff --git a/ts/package.json b/ts/package.json index 3859200..705a8ae 100644 --- a/ts/package.json +++ b/ts/package.json @@ -1,6 +1,6 @@ { "name": "@waymaker/client", - "version": "0.1.27", + "version": "0.1.31", "description": "Official TypeScript client for waymaker — locks, streams, KV, collections, sketches, cache, object store", "repository": "https://git.awesomike.com/pub/waymaker-client", "license": "MIT OR Apache-2.0",