perf(control): prove ownership with one etcd read and keep paged reads stable (3/7)

crates/nokv-client/src/artifact.rs

@@ -215,6 +215,24 @@ impl From<&PathMetadata> for ArtifactReadAuthority {
     }
 }
 
+/// Whether two metadata reads name the same immutable artifact revision.
+///
+/// A live paged read stays consistent while the path keeps its generation and
+/// revision: the bytes are immutable. `workspace_revision` is deliberately not
+/// compared, because every publication to any other path of the workspace
+/// advances it; comparing it would make concurrent writers elsewhere fail an
+/// otherwise unaffected read. Callers that need the exact workspace revision
+/// use `read_artifact_at_authority`.
+fn same_artifact_revision(left: &PathMetadata, right: &PathMetadata) -> bool {
+    left.path == right.path
+        && left.workspace_incarnation_id == right.workspace_incarnation_id
+        && left.generation == right.generation
+        && left.artifact_revision_id == right.artifact_revision_id
+        && left.dependency_count == right.dependency_count
+        && left.dependency_depth == right.dependency_depth
+        && left.descriptor == right.descriptor
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 enum ExpectedArtifactReadFence {
     Generation(u64),
@@ -1816,7 +1834,7 @@ where
                     "restore source range plan omitted path metadata".to_owned(),
                 )
             })?;
-            if &page_metadata != expected_metadata {
+            if !same_artifact_revision(&page_metadata, expected_metadata) {
                 return Err(ClientError::ArtifactReadFenceChanged);
             }
             append_range_page(&mut rows, result.blocks)?;
@@ -1875,7 +1893,7 @@ where
             })?;
             if metadata
                 .as_ref()
-                .is_some_and(|expected| expected != &page_metadata)
+                .is_some_and(|expected| !same_artifact_revision(expected, &page_metadata))
             {
                 return Err(ClientError::ArtifactReadFenceChanged);
             }
@@ -2852,6 +2870,7 @@ mod tests {
         abort_applies: usize,
         change_fence_on_second_page_once: bool,
         change_fence_on_every_range: bool,
+        advance_workspace_on_every_range: bool,
         fence_changed: bool,
         append_begin_conflicts: usize,
         advance_stage_before_response_once: bool,
@@ -2878,6 +2897,7 @@ mod tests {
                     abort_applies: 0,
                     change_fence_on_second_page_once: false,
                     change_fence_on_every_range: false,
+                    advance_workspace_on_every_range: false,
                     fence_changed: false,
                     append_begin_conflicts: 0,
                     advance_stage_before_response_once: false,
@@ -2894,6 +2914,10 @@ mod tests {
             self.state.lock().unwrap().change_fence_on_second_page_once = true;
         }
 
+        fn advance_workspace_on_every_range(&self) {
+            self.state.lock().unwrap().advance_workspace_on_every_range = true;
+        }
+
         fn change_fence_on_every_range(&self) {
             self.state.lock().unwrap().change_fence_on_every_range = true;
         }
@@ -3339,6 +3363,17 @@ mod tests {
                         metadata.generation = metadata.generation.saturating_add(1);
                     }
                 }
+                if state.advance_workspace_on_every_range {
+                    // Another path of the same workspace was published.
+                    if let Some(metadata) =
+                        state.path.as_mut().and_then(|path| path.metadata.as_mut())
+                    {
+                        metadata.workspace_revision = metadata.workspace_revision.saturating_add(1);
+                    }
+                    if let Some(metadata) = stored.metadata.as_mut() {
+                        metadata.workspace_revision = metadata.workspace_revision.saturating_add(1);
+                    }
+                }
                 let range_end = range
                     .offset
                     .checked_add(range.length)
@@ -4532,6 +4567,35 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn full_read_survives_publications_to_other_paths_of_the_workspace() {
+        let events = Arc::new(Mutex::new(Vec::new()));
+        let transport = ScriptedArtifactTransport::new(Arc::clone(&events), false);
+        let inspector = transport.clone();
+        let client = client(transport);
+        let store = RecordingStore::new(events, None);
+        let bytes = vec![0x5a; MAX_ARTIFACT_READ_PLAN_ROWS + 1];
+        client
+            .publish_artifact(&store, publish_options(1), &bytes)
+            .unwrap();
+        inspector.advance_workspace_on_every_range();
+
+        let read = client
+            .read_artifact(&store, None, target(), WorkspaceReadView::Live)
+            .unwrap();
+        assert_eq!(read.bytes, bytes);
+        // One metadata read and one plan read per page, without a restart.
+        assert_eq!(
+            inspector
+                .state()
+                .attempts
+                .iter()
+                .filter(|(label, _)| *label == "get_path")
+                .count(),
+            3
+        );
+    }
+
     #[test]
     fn full_read_pages_manifest_rows_and_restarts_on_a_metadata_fence_change() {
         let events = Arc::new(Mutex::new(Vec::new()));

crates/nokv-control/src/etcd.rs

@@ -2,7 +2,8 @@ use std::collections::BTreeMap;
 use std::future::Future;
 
 use etcd_client::{
-    Client, Compare, CompareOp, GetOptions, GetResponse, PutOptions, Txn, TxnOp, TxnOpResponse,
+    Client, Compare, CompareOp, GetOptions, GetResponse, KeyValue, PutOptions, Txn, TxnOp,
+    TxnOpResponse,
 };
 use tokio::runtime::{Builder, Runtime};
 
@@ -593,6 +594,52 @@ impl ControlStore for EtcdControlStore {
         })
     }
 
+    fn prove_owner(&self, lease: &LogicalShardLease) -> Result<LogicalShardRecord, ControlError> {
+        let mut client = self.client.clone();
+        let options = self.options.clone();
+        let lease = lease.clone();
+        self.block_on(async move {
+            let shard = &lease.logical_shard_id;
+            // One linearizable transaction reads the record and the
+            // lease-attached owner session at the same revision.
+            let response = client
+                .txn(Txn::new().and_then(vec![
+                    TxnOp::get(options.logical_shard_record_key(shard), None),
+                    TxnOp::get(options.logical_shard_recovery_key(shard), None),
+                    TxnOp::get(options.logical_shard_session_key(shard), None),
+                ]))
+                .await
+                .map_err(etcd_backend)?;
+            let mut responses = response.op_responses().into_iter();
+            let (
+                Some(TxnOpResponse::Get(record)),
+                Some(TxnOpResponse::Get(recovery)),
+                Some(TxnOpResponse::Get(session)),
+                None,
+            ) = (
+                responses.next(),
+                responses.next(),
+                responses.next(),
+                responses.next(),
+            )
+            else {
+                return Err(ControlError::Backend(
+                    "etcd ownership proof returned an unexpected response shape".to_owned(),
+                ));
+            };
+            let current = parse_logical_shard(&options, shard, &record, &recovery)?
+                .ok_or(ControlError::LogicalShardNotFound(*shard))?;
+            validate_record_lease(&current.record, &lease)?;
+            let session = parse_owner_session(&options, shard, session.kvs())?
+                .ok_or_else(|| ControlError::StaleLease(lease.clone()))?;
+            if session.lease != lease || session.attached_lease_id != lease_id_i64(lease.lease_id)?
+            {
+                return Err(ControlError::StaleLease(lease));
+            }
+            Ok(current.record)
+        })
+    }
+
     fn mark_serving(
         &self,
         lease: &LogicalShardLease,
@@ -957,6 +1004,23 @@ async fn fetch_logical_shard(
         .await
         .map_err(etcd_backend)?;
     let (record_response, recovery_response) = split_two_get_responses(response.op_responses())?;
+    parse_logical_shard(
+        options,
+        logical_shard_id,
+        &record_response,
+        &recovery_response,
+    )
+}
+
+/// Decode one logical shard from its routing and recovery reads taken at one
+/// etcd revision.
+fn parse_logical_shard(
+    options: &EtcdControlStoreOptions,
+    logical_shard_id: &LogicalShardId,
+    record_response: &GetResponse,
+    recovery_response: &GetResponse,
+) -> Result<Option<StoredLogicalShard>, ControlError> {
+    let record_key = options.logical_shard_record_key(logical_shard_id);
     let Some(kv) = record_response.kvs().first() else {
         return Ok(None);
     };
@@ -1136,8 +1200,17 @@ async fn fetch_owner_session(
     logical_shard_id: &LogicalShardId,
 ) -> Result<Option<StoredOwnerSession>, ControlError> {
     let key = options.logical_shard_session_key(logical_shard_id);
-    let response = client.get(key.clone(), None).await.map_err(etcd_backend)?;
-    let Some(kv) = response.kvs().first() else {
+    let response = client.get(key, None).await.map_err(etcd_backend)?;
+    parse_owner_session(options, logical_shard_id, response.kvs())
+}
+
+fn parse_owner_session(
+    options: &EtcdControlStoreOptions,
+    logical_shard_id: &LogicalShardId,
+    kvs: &[KeyValue],
+) -> Result<Option<StoredOwnerSession>, ControlError> {
+    let key = options.logical_shard_session_key(logical_shard_id);
+    let Some(kv) = kvs.first() else {
         return Ok(None);
     };
     let lease = decode_owner_session(kv.value())?;
@@ -1943,6 +2016,88 @@ mod tests {
             .unwrap();
     }
 
+    #[test]
+    #[ignore = "requires NOKV_TEST_ETCD_ENDPOINT"]
+    fn etcd_ownership_proof_reads_one_revision_and_fails_after_revocation() {
+        let endpoint = std::env::var("NOKV_TEST_ETCD_ENDPOINT")
+            .expect("NOKV_TEST_ETCD_ENDPOINT must name an isolated test etcd endpoint");
+        let nonce = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let prefix = format!("/nokv/test/ownership-proof-{}-{nonce}", std::process::id());
+        let options = EtcdControlStoreOptions::new([endpoint]).with_key_prefix(prefix.clone());
+        let store = EtcdControlStore::connect(options.clone()).unwrap();
+        let logical_shard_id = shard(11);
+        store.create_logical_shard(logical_shard_id).unwrap();
+        store
+            .create_root_placement(RootPlacement {
+                root_id: root(11),
+                logical_shard_id,
+                placement_generation: PlacementGeneration::new(1).unwrap(),
+                lifecycle: RootPlacementLifecycle::Provisioning,
+            })
+            .unwrap();
+        let lease = store
+            .acquire_owner(
+                &logical_shard_id,
+                NodeId::new("node-a").unwrap(),
+                "node-a:7000".to_owned(),
+            )
+            .unwrap();
+
+        let proved = store.prove_owner(&lease).unwrap();
+        assert_eq!(proved, store.renew_owner(&lease).unwrap());
+
+        // A proof only reads: it must not change any control key.
+        let mut client = store.client.clone();
+        let revision = |client: &mut Client| {
+            store
+                .block_on(async {
+                    let response = client
+                        .get(prefix.clone(), Some(GetOptions::new().with_prefix()))
+                        .await
+                        .map_err(etcd_backend)?;
+                    Ok::<i64, ControlError>(
+                        response
+                            .kvs()
+                            .iter()
+                            .map(|kv| kv.mod_revision())
+                            .max()
+                            .unwrap_or(0),
+                    )
+                })
+                .unwrap()
+        };
+        let before = revision(&mut client);
+        store.prove_owner(&lease).unwrap();
+        assert_eq!(revision(&mut client), before);
+
+        store
+            .block_on(async {
+                client
+                    .lease_revoke(lease_id_i64(lease.lease_id)?)
+                    .await
+                    .map_err(etcd_backend)?;
+                Ok::<(), ControlError>(())
+            })
+            .unwrap();
+        assert!(matches!(
+            store.prove_owner(&lease),
+            Err(ControlError::StaleLease(_))
+        ));
+
+        store
+            .block_on(async {
+                client
+                    .delete(prefix, Some(DeleteOptions::new().with_prefix()))
+                    .await
+                    .map_err(etcd_backend)?;
+                Ok::<(), ControlError>(())
+            })
+            .unwrap();
+    }
+
     #[test]
     #[ignore = "requires NOKV_TEST_ETCD_ENDPOINT"]
     fn etcd_expired_recovery_session_cleanup_rebinds_same_epoch() {

crates/nokv-control/src/store.rs

@@ -110,6 +110,17 @@ pub trait ControlStore: Send + Sync {
 
     fn renew_owner(&self, lease: &LogicalShardLease) -> Result<LogicalShardRecord, ControlError>;
 
+    /// Prove at one linearizable control read that `lease` still owns its
+    /// logical shard, without extending the lease.
+    ///
+    /// A successor cannot acquire the shard while this owner's lease-attached
+    /// session exists, so a successful proof orders every write the owner made
+    /// before it ahead of any successor. Keeping the lease alive remains the
+    /// job of periodic [`Self::renew_owner`] calls. The default renews.
+    fn prove_owner(&self, lease: &LogicalShardLease) -> Result<LogicalShardRecord, ControlError> {
+        self.renew_owner(lease)
+    }
+
     /// Publish a caller-serialized recovery frontier and make the exact owner
     /// generation routable.
     fn mark_serving(

crates/nokv-server/src/recovery_publisher.rs

@@ -134,8 +134,9 @@ pub enum RecoveryPublisherError {
     Poisoned,
 }
 
-/// Executor adapter that repairs an earlier pending upload before dispatch and
-/// publishes every local write before returning its response.
+/// Executor adapter that, in `Shared` mode, repairs an earlier pending upload
+/// before dispatch, and in every mode publishes (or proves owner liveness for)
+/// every local write before returning its response.
 pub struct RecoveryPublishingExecutor {
     inner: Arc<dyn WorkspaceRequestExecutor>,
     publisher: Arc<RecoveryPublisher>,
@@ -221,7 +222,7 @@ impl RecoveryPublisher {
         &self,
         required: RecoveryState,
     ) -> Result<LogicalShardRecord, RecoveryPublisherError> {
-        let record = self.control.renew_owner(&self.lease)?;
+        let record = self.control.prove_owner(&self.lease)?;
         self.validate_control_frontier(&record, required)?;
         if record.pending_recovery_upload.is_some() {
             return Err(RecoveryPublisherError::InvalidState(
@@ -240,7 +241,7 @@ impl RecoveryPublisher {
             // local and durable frontiers already match, prove that this exact
             // owner session is still live before allowing a response to leave
             // the shard.
-            let mut record = self.control.renew_owner(&self.lease)?;
+            let mut record = self.control.prove_owner(&self.lease)?;
             self.validate_control_frontier(&record, required)?;
             if record.pending_recovery_upload.is_none()
                 && record.durable_lsn == required.applied_recovery_lsn
@@ -489,9 +490,16 @@ impl RecoveryPublishingExecutor {
 
 impl WorkspaceRequestExecutor for RecoveryPublishingExecutor {
     fn execute(&self, request: &WorkspaceRpcRequest) -> Result<ExecutedRequest, RpcFailure> {
-        self.publisher
-            .publish_current()
-            .map_err(RecoveryPublisherError::rpc_failure)?;
+        // Shared publication first repairs an earlier interrupted upload. A
+        // local-only owner has nothing to repair, and the barrier below proves
+        // liveness after everything this request observed or wrote, so a
+        // second proof before dispatch would add a control round trip without
+        // fencing anything more.
+        if self.publisher.mode() == RecoveryPublicationMode::Shared {
+            self.publisher
+                .publish_current()
+                .map_err(RecoveryPublisherError::rpc_failure)?;
+        }
         let outcome = self.inner.execute(request);
         self.publisher
             .publish_current()