feat(ecstore): add remote snapshot lease RPCs (#5389)

* feat(ecstore): add local snapshot leases

* feat(ecstore): add remote snapshot lease RPCs

* fix(rpc): keep snapshot lease checks CI-compatible

crates/ecstore/src/api/mod.rs

@@ -305,7 +305,8 @@ pub mod disk {
         CheckPartsResp, DeleteOptions, Disk, DiskAPI, DiskInfo, DiskInfoOptions, DiskLocation, DiskOption, DiskStore,
         FileInfoVersions, FileReader, FileWriter, HEALING_MARKER_PATH, NsScannerOpenRequest, OldCurrentSize,
         PartTransactionAction, RUSTFS_META_BUCKET, ReadMultipleReq, ReadMultipleResp, ReadOptions, RenameDataResp,
-        STORAGE_FORMAT_FILE, UpdateMetadataOpts, VolumeInfo, WalkDirOptions, new_disk, validate_batch_read_version_item_count,
+        STORAGE_FORMAT_FILE, SnapshotLeaseToken, UpdateMetadataOpts, VolumeInfo, WalkDirOptions, new_disk,
+        validate_batch_read_version_item_count,
     };
     pub use bytes::Bytes;
     pub use endpoint::Endpoint;

crates/ecstore/src/cluster/rpc/remote_disk.rs

@@ -25,7 +25,7 @@ use crate::disk::error::{Error, Result};
 use crate::disk::{
     BatchReadVersionReq, BatchReadVersionResp, CheckPartsResp, DeleteOptions, DiskAPI, DiskInfo, DiskInfoOptions, DiskLocation,
     DiskOption, FileInfoVersions, FileReader, FileWriter, PartTransactionAction, ReadMultipleReq, ReadMultipleResp, ReadOptions,
-    RenameDataResp, UpdateMetadataOpts, VolumeInfo, WalkDirOptions, batch_read_version_one_by_one,
+    RenameDataResp, SnapshotLeaseToken, UpdateMetadataOpts, VolumeInfo, WalkDirOptions, batch_read_version_one_by_one,
     disk_store::{
         DEFAULT_RUSTFS_DRIVE_ACTIVE_MONITORING, ENV_RUSTFS_DRIVE_ACTIVE_MONITORING, SKIP_IF_SUCCESS_BEFORE,
         get_drive_active_check_interval, get_drive_active_check_timeout, get_drive_disk_info_timeout, get_drive_list_dir_timeout,
@@ -50,8 +50,9 @@ use rustfs_protos::proto_gen::node_service::{
     DeleteVersionRequest, DeleteVersionsRequest, DeleteVolumeRequest, DiskInfoRequest, ListDirRequest, ListVolumesRequest,
     MakeVolumeRequest, MakeVolumesRequest, PreparePartTransactionRequest, ReadAllRequest, ReadMetadataRequest,
     ReadMultipleRequest, ReadMultipleResponse, ReadPartsRequest, ReadVersionRequest, ReadXlRequest, RenameDataRequest,
-    RenameFileRequest, SettlePartTransactionRequest, StatVolumeRequest, UpdateMetadataRequest, VerifyFileRequest,
-    WriteAllRequest, WriteMetadataRequest, node_service_client::NodeServiceClient,
+    RenameFileRequest, SettlePartTransactionRequest, SnapshotLeaseReleaseRequest, SnapshotLeaseRenewRequest,
+    SnapshotLeaseRequest, SnapshotLeaseResponse, StatVolumeRequest, UpdateMetadataRequest, VerifyFileRequest, WriteAllRequest,
+    WriteMetadataRequest, node_service_client::NodeServiceClient,
 };
 use serde::{Serialize, de::DeserializeOwned};
 use std::{
@@ -100,6 +101,18 @@ const LOG_COMPONENT_ECSTORE: &str = "ecstore";
 const LOG_SUBSYSTEM_REMOTE_DISK: &str = "remote_disk";
 const EVENT_REMOTE_DISK_HEALTH: &str = "remote_disk_health";
 const EVENT_REMOTE_DISK_RPC: &str = "remote_disk_rpc";
+const SNAPSHOT_LEASE_PROTOCOL_VERSION: u32 = 1;
+pub const REMOTE_SNAPSHOT_LEASE_TTL: Duration = Duration::from_secs(60);
+
+fn snapshot_lease_token_from_response(response: SnapshotLeaseResponse) -> Result<SnapshotLeaseToken> {
+    if !response.success {
+        return Err(response.error.unwrap_or_default().into());
+    }
+    if response.protocol_version != SNAPSHOT_LEASE_PROTOCOL_VERSION {
+        return Err(Error::other("remote snapshot lease protocol is incompatible"));
+    }
+    SnapshotLeaseToken::from_slice(&response.token)
+}
 
 /// Bind a mutating disk RPC to its canonical body: the digest lands in the request metadata, and
 /// the signing interceptor folds it (plus a replay-protected nonce) into the v2 signature scope
@@ -1784,6 +1797,81 @@ impl DiskAPI for RemoteDisk {
         .await
     }
 
+    async fn acquire_snapshot_lease(&self, volume: &str, path: &str) -> Result<SnapshotLeaseToken> {
+        self.execute_with_timeout(
+            || async {
+                let mut client = self
+                    .get_client()
+                    .await
+                    .map_err(|err| Error::other(format!("can not get client, err: {err}")))?;
+                let mut request = Request::new(SnapshotLeaseRequest {
+                    disk: self.endpoint.to_string(),
+                    volume: volume.to_string(),
+                    path: path.to_string(),
+                    ttl_ms: u64::try_from(REMOTE_SNAPSHOT_LEASE_TTL.as_millis())
+                        .map_err(|_| Error::other("snapshot lease TTL cannot be represented"))?,
+                });
+                let canonical_body = rustfs_protos::canonical_snapshot_lease_request_body(request.get_ref());
+                attach_mutation_body_digest(&mut request, canonical_body, "acquire_snapshot_lease")?;
+                let response = client.acquire_snapshot_lease(request).await?.into_inner();
+                snapshot_lease_token_from_response(response)
+            },
+            get_max_timeout_duration(),
+        )
+        .await
+    }
+
+    async fn renew_snapshot_lease(&self, volume: &str, path: &str, token: SnapshotLeaseToken) -> Result<SnapshotLeaseToken> {
+        self.execute_with_timeout(
+            || async {
+                let mut client = self
+                    .get_client()
+                    .await
+                    .map_err(|err| Error::other(format!("can not get client, err: {err}")))?;
+                let mut request = Request::new(SnapshotLeaseRenewRequest {
+                    disk: self.endpoint.to_string(),
+                    volume: volume.to_string(),
+                    path: path.to_string(),
+                    token: token.as_bytes().to_vec().into(),
+                    ttl_ms: u64::try_from(REMOTE_SNAPSHOT_LEASE_TTL.as_millis())
+                        .map_err(|_| Error::other("snapshot lease TTL cannot be represented"))?,
+                });
+                let canonical_body = rustfs_protos::canonical_snapshot_lease_renew_request_body(request.get_ref());
+                attach_mutation_body_digest(&mut request, canonical_body, "renew_snapshot_lease")?;
+                let response = client.renew_snapshot_lease(request).await?.into_inner();
+                snapshot_lease_token_from_response(response)
+            },
+            get_max_timeout_duration(),
+        )
+        .await
+    }
+
+    async fn release_snapshot_lease(&self, volume: &str, path: &str, token: SnapshotLeaseToken) -> Result<()> {
+        self.execute_with_timeout(
+            || async {
+                let mut client = self
+                    .get_client()
+                    .await
+                    .map_err(|err| Error::other(format!("can not get client, err: {err}")))?;
+                let mut request = Request::new(SnapshotLeaseReleaseRequest {
+                    disk: self.endpoint.to_string(),
+                    volume: volume.to_string(),
+                    path: path.to_string(),
+                    token: token.as_bytes().to_vec().into(),
+                });
+                let canonical_body = rustfs_protos::canonical_snapshot_lease_release_request_body(request.get_ref());
+                attach_mutation_body_digest(&mut request, canonical_body, "release_snapshot_lease")?;
+                let response = client.release_snapshot_lease(request).await?.into_inner();
+                if !response.success {
+                    return Err(response.error.unwrap_or_default().into());
+                }
+                Ok(())
+            },
+            get_max_timeout_duration(),
+        )
+        .await
+    }
+
     #[tracing::instrument(level = "trace", skip_all)]
     async fn write_metadata(&self, _org_volume: &str, volume: &str, path: &str, fi: FileInfo) -> Result<()> {
         trace!(
@@ -2930,6 +3018,34 @@ mod tests {
 
     static INIT: Once = Once::new();
 
+    #[test]
+    fn snapshot_lease_response_requires_current_protocol_and_valid_token() {
+        let token = SnapshotLeaseToken::new();
+        let response = SnapshotLeaseResponse {
+            success: true,
+            token: token.as_bytes().to_vec().into(),
+            protocol_version: SNAPSHOT_LEASE_PROTOCOL_VERSION,
+            error: None,
+        };
+        assert_eq!(snapshot_lease_token_from_response(response).unwrap(), token);
+
+        let incompatible = SnapshotLeaseResponse {
+            success: true,
+            token: token.as_bytes().to_vec().into(),
+            protocol_version: SNAPSHOT_LEASE_PROTOCOL_VERSION + 1,
+            error: None,
+        };
+        assert!(snapshot_lease_token_from_response(incompatible).is_err());
+
+        let malformed = SnapshotLeaseResponse {
+            success: true,
+            token: Bytes::from_static(b"not-a-uuid"),
+            protocol_version: SNAPSHOT_LEASE_PROTOCOL_VERSION,
+            error: None,
+        };
+        assert!(snapshot_lease_token_from_response(malformed).is_err());
+    }
+
     #[test]
     fn list_volumes_decode_rejects_a_malformed_entry() {
         let valid = serde_json::to_string(&VolumeInfo {

crates/ecstore/src/disk/disk_store.rs

@@ -1365,6 +1365,14 @@ impl DiskAPI for LocalDiskWrapper {
         .await
     }
 
+    async fn renew_snapshot_lease(&self, volume: &str, path: &str, token: SnapshotLeaseToken) -> Result<SnapshotLeaseToken> {
+        self.track_disk_health(
+            || async { self.disk.renew_snapshot_lease(volume, path, token).await },
+            get_max_timeout_duration(),
+        )
+        .await
+    }
+
     async fn delete_data_dir(&self, volume: &str, path: &str, opts: DeleteOptions) -> Result<DataDirDeleteStatus> {
         self.track_disk_health(
             || async { self.disk.delete_data_dir(volume, path, opts).await },

crates/ecstore/src/disk/local.rs

@@ -7908,6 +7908,23 @@ impl DiskAPI for LocalDisk {
         }
     }
 
+    async fn renew_snapshot_lease(&self, volume: &str, path: &str, token: SnapshotLeaseToken) -> Result<SnapshotLeaseToken> {
+        let key = SnapshotLeaseKey {
+            volume: volume.to_string(),
+            path: path.to_string(),
+        };
+        let mut registry = self.snapshot_leases.lock().await;
+        let Some(entry) = registry.entries.get_mut(&key) else {
+            return Err(DiskError::FileNotFound);
+        };
+        if entry.deleting || !entry.tokens.remove(&token) {
+            return Err(DiskError::FileNotFound);
+        }
+        let renewed = SnapshotLeaseToken::new();
+        entry.tokens.insert(renewed);
+        Ok(renewed)
+    }
+
     async fn delete_data_dir(&self, volume: &str, path: &str, opts: DeleteOptions) -> Result<DataDirDeleteStatus> {
         let key = SnapshotLeaseKey {
             volume: volume.to_string(),
@@ -14957,6 +14974,13 @@ mod test {
             .acquire_snapshot_lease(volume, &data_dir)
             .await
             .expect("second lease should be acquired");
+        let renewed = disk
+            .renew_snapshot_lease(volume, &data_dir, first)
+            .await
+            .expect("first lease should renew atomically");
+        disk.release_snapshot_lease(volume, &data_dir, first)
+            .await
+            .expect("the superseded token should be idempotent");
         let status = disk
             .delete_data_dir(
                 volume,
@@ -14976,9 +15000,9 @@ mod test {
             Bytes::from_static(b"later")
         );
 
-        disk.release_snapshot_lease(volume, &data_dir, first)
+        disk.release_snapshot_lease(volume, &data_dir, renewed)
             .await
-            .expect("first lease release should succeed");
+            .expect("renewed lease release should succeed");
         assert!(
             disk.read_all(volume, &first_part).await.is_ok(),
             "one remaining lease must keep the data directory"

crates/ecstore/src/disk/mod.rs

@@ -79,6 +79,18 @@ impl SnapshotLeaseToken {
     pub fn new() -> Self {
         Self(Uuid::new_v4())
     }
+
+    pub fn from_slice(bytes: &[u8]) -> Result<Self> {
+        let uuid = Uuid::from_slice(bytes).map_err(|_| Error::other("invalid snapshot lease token"))?;
+        if uuid.is_nil() {
+            return Err(Error::other("invalid snapshot lease token"));
+        }
+        Ok(Self(uuid))
+    }
+
+    pub fn as_bytes(&self) -> &[u8; 16] {
+        self.0.as_bytes()
+    }
 }
 
 impl Default for SnapshotLeaseToken {
@@ -284,6 +296,13 @@ impl DiskAPI for Disk {
         }
     }
 
+    async fn renew_snapshot_lease(&self, volume: &str, path: &str, token: SnapshotLeaseToken) -> Result<SnapshotLeaseToken> {
+        match self {
+            Disk::Local(local_disk) => local_disk.renew_snapshot_lease(volume, path, token).await,
+            Disk::Remote(remote_disk) => remote_disk.renew_snapshot_lease(volume, path, token).await,
+        }
+    }
+
     async fn delete_data_dir(&self, volume: &str, path: &str, opts: DeleteOptions) -> Result<DataDirDeleteStatus> {
         match self {
             Disk::Local(local_disk) => local_disk.delete_data_dir(volume, path, opts).await,
@@ -694,6 +713,9 @@ pub trait DiskAPI: Debug + Send + Sync + 'static {
     async fn release_snapshot_lease(&self, _volume: &str, _path: &str, _token: SnapshotLeaseToken) -> Result<()> {
         Err(Error::other("snapshot leases are not supported by this disk"))
     }
+    async fn renew_snapshot_lease(&self, _volume: &str, _path: &str, _token: SnapshotLeaseToken) -> Result<SnapshotLeaseToken> {
+        Err(Error::other("snapshot leases are not supported by this disk"))
+    }
     async fn delete_data_dir(&self, volume: &str, path: &str, opts: DeleteOptions) -> Result<DataDirDeleteStatus> {
         self.delete(volume, path, opts).await?;
         Ok(DataDirDeleteStatus::Deleted)