persist: fresh command timestamps during retries (#38060)

Prior to this PR persist commands that accepted a heartbeat timestamp
parameter computed that parameter once and used the same value
throughout all the retries, if any. When the number of retry attempts is
low this is usually not a problem but on heavily utilized clusters
retries can go on for multiple minutes. In the extreme case by the time
the command succeeds, it commits a timestamp that is too stale.

Two instances of this problem were the root cause of incidents 1177 and
1134 were the registration of a read handle was retrying for more than
15 minutes, and by the time it succeeded the handle had already lost its
lease.

This PR removes the heartbeat argument from the affected methods
altogether and instead uses internal fresh calls to persist's `NowFn`
configuration on every retry. By doing this the committed heartbeat
timestamp is only stale by the network latency of the final consensus op
that succeeded.

Signed-off-by: Petros Angelatos <petrosagg@gmail.com>

Author: petrosagg

src/persist-client/src/internal/machine.rs

@@ -222,12 +222,12 @@ where
         }
     }
 
+    /// Registers a leased reader, returning its initial state.
     pub async fn register_leased_reader(
         &self,
         reader_id: &LeasedReaderId,
         purpose: &str,
         lease_duration: Duration,
-        heartbeat_timestamp_ms: u64,
         use_critical_since: bool,
     ) -> (LeasedReaderState<T>, RoutineMaintenance) {
         let metrics = Arc::clone(&self.applier.metrics);
@@ -239,7 +239,9 @@ where
                     purpose,
                     seqno,
                     lease_duration,
-                    heartbeat_timestamp_ms,
+                    // NOTE: Sample the clock here rather than hoisting it out of the closure so
+                    // that a fresh value is used on every retry of this command.
+                    (cfg.now)(),
                     use_critical_since,
                 )
             })
@@ -318,20 +320,19 @@ where
         (reqs, maintenance)
     }
 
+    /// Appends `batch` if the shard upper matches its lower.
     pub async fn compare_and_append(
         &self,
         batch: &HollowBatch<T>,
         writer_id: &WriterId,
         debug_info: &HandleDebugState,
-        heartbeat_timestamp_ms: u64,
     ) -> CompareAndAppendRes<T> {
         let idempotency_token = IdempotencyToken::new();
         loop {
             let res = self
                 .compare_and_append_idempotent(
                     batch,
                     writer_id,
-                    heartbeat_timestamp_ms,
                     &idempotency_token,
                     debug_info,
                     None,
@@ -375,7 +376,6 @@ where
         &self,
         batch: &HollowBatch<T>,
         writer_id: &WriterId,
-        heartbeat_timestamp_ms: u64,
         idempotency_token: &IdempotencyToken,
         debug_info: &HandleDebugState,
         // Only exposed for testing. In prod, this always starts as None, but
@@ -488,7 +488,9 @@ where
                     state.compare_and_append(
                         batch,
                         writer_id,
-                        heartbeat_timestamp_ms,
+                        // NOTE: Sample the clock here rather than hoisting it out of the closure
+                        // so that a fresh value is used on every retry of this command.
+                        (cfg.now)(),
                         lease_duration_ms,
                         idempotency_token,
                         debug_info,
@@ -627,22 +629,18 @@ where
         }
     }
 
+    /// Downgrades the reader's since capability, also heartbeating its lease.
     pub async fn downgrade_since(
         &self,
         reader_id: &LeasedReaderId,
         outstanding_seqno: SeqNo,
         new_since: &Antichain<T>,
-        heartbeat_timestamp_ms: u64,
     ) -> (SeqNo, Since<T>, RoutineMaintenance) {
         let metrics = Arc::clone(&self.applier.metrics);
-        self.apply_unbatched_idempotent_cmd(&metrics.cmds.downgrade_since, |seqno, _cfg, state| {
-            state.downgrade_since(
-                reader_id,
-                seqno,
-                outstanding_seqno,
-                new_since,
-                heartbeat_timestamp_ms,
-            )
+        self.apply_unbatched_idempotent_cmd(&metrics.cmds.downgrade_since, |seqno, cfg, state| {
+            // NOTE: Sample the clock here rather than hoisting it out of the closure so that a
+            // fresh value is used on every retry of this command.
+            state.downgrade_since(reader_id, seqno, outstanding_seqno, new_since, (cfg.now)())
         })
         .await
     }
@@ -1500,12 +1498,7 @@ pub mod datadriven {
         let reader_id = args.expect("reader_id");
         let (_, since, routine) = datadriven
             .machine
-            .downgrade_since(
-                &reader_id,
-                seqno,
-                &since,
-                (datadriven.machine.applier.cfg.now)(),
-            )
+            .downgrade_since(&reader_id, seqno, &since)
             .await;
         datadriven.routine.push(routine);
         Ok(format!(
@@ -2195,7 +2188,6 @@ pub mod datadriven {
                 &reader_id,
                 "tests",
                 READER_LEASE_DURATION.get(&datadriven.client.cfg),
-                (datadriven.client.cfg.now)(),
                 false,
             )
             .await;
@@ -2313,7 +2305,6 @@ pub mod datadriven {
             .expect("unknown batch")
             .clone();
         let token = args.optional("token").unwrap_or_else(IdempotencyToken::new);
-        let now = (datadriven.client.cfg.now)();
 
         let (id, maintenance) = datadriven
             .machine
@@ -2330,7 +2321,6 @@ pub mod datadriven {
                 .compare_and_append_idempotent(
                     &batch.batch,
                     &writer_id,
-                    now,
                     &token,
                     &HandleDebugState::default(),
                     indeterminate,
@@ -2494,7 +2484,6 @@ pub mod tests {
                     &batch.into_hollow_batch(),
                     &write.writer_id,
                     &HandleDebugState::default(),
-                    (write.cfg.now)(),
                 )
                 .await
                 .unwrap();

src/persist-client/src/lib.rs

@@ -357,13 +357,11 @@ impl PersistClient {
         let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
 
         let reader_id = LeasedReaderId::new();
-        let heartbeat_ts = (self.cfg.now)();
         let (reader_state, maintenance) = machine
             .register_leased_reader(
                 &reader_id,
                 &diagnostics.handle_purpose,
                 READER_LEASE_DURATION.get(&self.cfg),
-                heartbeat_ts,
                 use_critical_since,
             )
             .await;

src/persist-client/src/read.rs

@@ -722,7 +722,6 @@ where
             }
 
             let before_heartbeat = Instant::now();
-            let heartbeat_ms = (machine.applier.cfg.now)();
             let current_seqno = machine.seqno();
             let result = leased_seqnos.modify(|s| {
                 if s.expired {
@@ -737,7 +736,7 @@ where
             let actual_since = match result {
                 Ok(held_seqno) => {
                     let (seqno, actual_since, maintenance) = machine
-                        .downgrade_since(&reader_id, held_seqno, &held_since, heartbeat_ms)
+                        .downgrade_since(&reader_id, held_seqno, &held_since)
                         .await;
                     leased_seqnos.modify(|s| {
                         s.applied_since.clone_from(&actual_since.0);
@@ -990,13 +989,11 @@ where
         let new_reader_id = LeasedReaderId::new();
         let machine = self.machine.clone();
         let gc = self.gc.clone();
-        let heartbeat_ts = (self.cfg.now)();
         let (reader_state, maintenance) = machine
             .register_leased_reader(
                 &new_reader_id,
                 purpose,
                 READER_LEASE_DURATION.get(&self.cfg),
-                heartbeat_ts,
                 false,
             )
             .await;

src/persist-client/src/write.rs

@@ -301,15 +301,9 @@ where
             let desc = Description::new(lower.clone(), target.clone(), since);
             let batch = HollowBatch::empty(desc);
 
-            let heartbeat_timestamp = (self.cfg.now)();
             let res = self
                 .machine
-                .compare_and_append(
-                    &batch,
-                    &self.writer_id,
-                    &self.debug_state,
-                    heartbeat_timestamp,
-                )
+                .compare_and_append(&batch, &self.writer_id, &self.debug_state)
                 .await;
 
             use CompareAndAppendRes::*;
@@ -749,15 +743,9 @@ where
                 }
             }
 
-            let heartbeat_timestamp = (self.cfg.now)();
             let res = self
                 .machine
-                .compare_and_append(
-                    &combined_batch,
-                    &self.writer_id,
-                    &self.debug_state,
-                    heartbeat_timestamp,
-                )
+                .compare_and_append(&combined_batch, &self.writer_id, &self.debug_state)
                 .await;
 
             match res {