extend WAL retention if some other replicas are dead

generall · generall · commit 67e26d672eb7 · 2025-12-28T01:47:06.000+01:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -283,7 +283,7 @@ tonic-reflection = "0.11.0"
 tracing = { version = "0.1", features = ["async-await"] }
 uuid = { version = "1.19", features = ["v4", "serde"] }
 validator = { version = "0.20.0", features = ["derive"] }
-wal = { git = "https://github.com/qdrant/wal.git", rev = "b03a156b0ffdbe8109b2c3c409ce3be2ec6d229e" }
+wal = { git = "https://github.com/qdrant/wal.git", branch = "allow-set-retention" }
 zerocopy = { version = "0.8.31", features = ["derive"] }
 atomic_refcell = "0.1.13"
 byteorder = "1.5.0"
diff --git a/lib/collection/src/shards/forward_proxy_shard.rs b/lib/collection/src/shards/forward_proxy_shard.rs
@@ -395,6 +395,14 @@ impl ForwardProxyShard {
             .estimate_cardinality(filter, hw_measurement_acc)
             .await
     }
+
+    pub async fn set_extended_wal_retention(&self) {
+        self.wrapped_shard.set_extended_wal_retention().await;
+    }
+
+    pub async fn set_normal_wal_retention(&self) {
+        self.wrapped_shard.set_normal_wal_retention().await;
+    }
 }
 
 #[async_trait]
diff --git a/lib/collection/src/shards/local_shard/mod.rs b/lib/collection/src/shards/local_shard/mod.rs
@@ -17,6 +17,7 @@ mod snapshot_tests;
 mod drop;
 #[cfg(feature = "testing")]
 pub mod testing;
+mod wal_ops;
 
 use std::collections::{BTreeSet, HashMap};
 use std::ops::Deref;
diff --git a/lib/collection/src/shards/local_shard/wal_ops.rs b/lib/collection/src/shards/local_shard/wal_ops.rs
@@ -0,0 +1,14 @@
+use crate::shards::local_shard::LocalShard;
+
+impl LocalShard {
+    /// WAL is keeping more data, even if truncated.
+    /// Useful if we expect to read old WAL records soon.
+    pub async fn set_extended_wal_retention(&self) {
+        self.wal.set_extended_retention().await;
+    }
+
+    /// WAL is keeping normal amount of data after truncation.
+    pub async fn set_normal_wal_retention(&self) {
+        self.wal.set_normal_retention().await;
+    }
+}
diff --git a/lib/collection/src/shards/proxy_shard.rs b/lib/collection/src/shards/proxy_shard.rs
@@ -181,6 +181,14 @@ impl ProxyShard {
             .estimate_cardinality(filter, hw_measurement_acc)
             .await
     }
+
+    pub async fn set_extended_wal_retention(&self) {
+        self.wrapped_shard.set_extended_wal_retention().await;
+    }
+
+    pub async fn set_normal_wal_retention(&self) {
+        self.wrapped_shard.set_normal_wal_retention().await;
+    }
 }
 
 #[async_trait]
diff --git a/lib/collection/src/shards/queue_proxy_shard.rs b/lib/collection/src/shards/queue_proxy_shard.rs
@@ -269,6 +269,18 @@ impl QueueProxyShard {
             .estimate_cardinality(filter, hw_measurement_acc)
             .await
     }
+
+    pub async fn set_extended_wal_retention(&self) {
+        if let Some(inner) = &self.inner {
+            inner.wrapped_shard.set_extended_wal_retention().await;
+        }
+    }
+
+    pub async fn set_normal_wal_retention(&self) {
+        if let Some(inner) = &self.inner {
+            inner.wrapped_shard.set_normal_wal_retention().await;
+        }
+    }
 }
 
 #[async_trait]
diff --git a/lib/collection/src/shards/replica_set/mod.rs b/lib/collection/src/shards/replica_set/mod.rs
@@ -694,6 +694,8 @@ impl ShardReplicaSet {
 
         if self.this_peer_id() == peer_id {
             self.on_local_state_updated(state).await?;
+        } else {
+            self.on_remote_state_updated(peer_id, state).await;
         }
 
         self.update_locally_disabled(peer_id);
@@ -764,6 +766,8 @@ impl ShardReplicaSet {
 
         if self.this_peer_id() == peer_id {
             self.on_local_state_updated(state).await?;
+        } else {
+            self.on_remote_state_updated(peer_id, state).await;
         }
 
         self.update_locally_disabled(peer_id);
@@ -783,11 +787,42 @@ impl ShardReplicaSet {
             } else {
                 local_shard.take_newest_clocks_snapshot().await?;
             }
+            // Reset WAL retention to normal whenever local shard changes state
+            local_shard.set_normal_wal_retention().await;
         }
 
         Ok(())
     }
 
+    /// Called when a peer state is changed (except local peer).
+    ///
+    async fn on_remote_state_updated(&self, _peer_id: PeerId, _new_state: ReplicaState) {
+        let mut is_any_remote_dead = false;
+        let mut is_local_active = false;
+        let this_peer_id = self.this_peer_id();
+
+        for (peer_id, peer_state) in self.replica_state.read().peers().iter() {
+            if *peer_id == this_peer_id {
+                if peer_state.is_active() {
+                    is_local_active = true;
+                }
+            } else if peer_state.requires_recovery() {
+                is_any_remote_dead = true;
+            }
+        }
+
+        {
+            let local_opt = self.local.read().await;
+            if let Some(local_shard) = local_opt.as_ref() {
+                if is_local_active && is_any_remote_dead {
+                    local_shard.set_extended_wal_retention().await;
+                } else {
+                    local_shard.set_normal_wal_retention().await;
+                }
+            }
+        }
+    }
+
     pub async fn remove_peer(&self, peer_id: PeerId) -> CollectionResult<()> {
         if self.this_peer_id() == peer_id {
             self.remove_local().await?;
diff --git a/lib/collection/src/shards/replica_set/replica_set_state.rs b/lib/collection/src/shards/replica_set/replica_set_state.rs
@@ -207,6 +207,22 @@ impl ReplicaState {
         }
     }
 
+    /// Check if the replica state requires automatic recovery to be scheduled.
+    pub fn requires_recovery(self) -> bool {
+        match self {
+            ReplicaState::Dead => true,
+            ReplicaState::Active
+            | ReplicaState::Partial
+            | ReplicaState::Initializing
+            | ReplicaState::Listener
+            | ReplicaState::PartialSnapshot
+            | ReplicaState::Recovery
+            | ReplicaState::Resharding
+            | ReplicaState::ReshardingScaleDown
+            | ReplicaState::ActiveRead => false,
+        }
+    }
+
     /// Check whether the replica state is partial or partial-like.
     ///
     /// In other words: is the state related to shard transfers?
diff --git a/lib/collection/src/shards/shard.rs b/lib/collection/src/shards/shard.rs
@@ -368,6 +368,27 @@ impl Shard {
         }
     }
 
+    pub async fn set_extended_wal_retention(&self) {
+        match self {
+            Shard::Local(local) => local.set_extended_wal_retention().await,
+            Shard::Proxy(proxy) => proxy.set_extended_wal_retention().await,
+            Shard::ForwardProxy(forward_proxy) => forward_proxy.set_extended_wal_retention().await,
+            Shard::QueueProxy(queue_proxy) => queue_proxy.set_extended_wal_retention().await,
+            Shard::Dummy(_) => {} // Do nothing for dummy shard
+        }
+    }
+
+    /// WAL is keeping normal amount of data after truncation.
+    pub async fn set_normal_wal_retention(&self) {
+        match self {
+            Shard::Local(local) => local.set_normal_wal_retention().await,
+            Shard::Proxy(proxy) => proxy.set_normal_wal_retention().await,
+            Shard::ForwardProxy(forward_proxy) => forward_proxy.set_normal_wal_retention().await,
+            Shard::QueueProxy(queue_proxy) => queue_proxy.set_normal_wal_retention().await,
+            Shard::Dummy(_) => {} // Do nothing for dummy shard
+        }
+    }
+
     pub async fn estimate_cardinality(
         &self,
         filter: Option<&Filter>,
diff --git a/lib/collection/src/wal_delta.rs b/lib/collection/src/wal_delta.rs
@@ -170,6 +170,14 @@ impl RecoverableWal {
         }
         Ok(())
     }
+
+    pub async fn set_extended_retention(&self) {
+        self.wal.lock().await.set_extended_retention();
+    }
+
+    pub async fn set_normal_retention(&self) {
+        self.wal.lock().await.set_normal_retention();
+    }
 }
 
 /// Resolve the WAL delta for the given `recovery_point`
diff --git a/lib/shard/src/wal.rs b/lib/shard/src/wal.rs
@@ -26,6 +26,10 @@ pub struct SerdeWal<R> {
 
 const FIRST_INDEX_FILE: &str = "first-index";
 
+/// When increased retention is used, how many times more segments to retain.
+/// (this is used to extend recoverable history and allow WAL shard transfers)
+const INCREASED_RETENTION_FACTOR: usize = 10;
+
 impl<R: DeserializeOwned + Serialize> SerdeWal<R> {
     pub fn new(dir: &Path, wal_options: WalOptions) -> Result<SerdeWal<R>> {
         let wal = Wal::with_options(dir, &wal_options)
@@ -210,6 +214,17 @@ impl<R: DeserializeOwned + Serialize> SerdeWal<R> {
     pub fn segment_capacity(&self) -> usize {
         self.options.segment_capacity
     }
+
+    pub fn set_extended_retention(&mut self) {
+        let normal_retention = self.options.retain_closed.get();
+        self.wal
+            .set_retention(normal_retention * INCREASED_RETENTION_FACTOR);
+    }
+
+    pub fn set_normal_retention(&mut self) {
+        let normal_retention = self.options.retain_closed.get();
+        self.wal.set_retention(normal_retention);
+    }
 }
 
 #[derive(Debug, Deserialize, Serialize)]

Original file line number	Diff line number	Diff line change
`@@ -395,6 +395,14 @@ impl ForwardProxyShard {`
`395`	`395`	`.estimate_cardinality(filter, hw_measurement_acc)`
`396`	`396`	`.await`
`397`	`397`	`}`
	`398`	`+`
	`399`	`+ pub async fn set_extended_wal_retention(&self) {`
	`400`	`+ self.wrapped_shard.set_extended_wal_retention().await;`
	`401`	`+ }`
	`402`	`+`
	`403`	`+ pub async fn set_normal_wal_retention(&self) {`
	`404`	`+ self.wrapped_shard.set_normal_wal_retention().await;`
	`405`	`+ }`
`398`	`406`	`}`
`399`	`407`
`400`	`408`	`#[async_trait]`
Original file line number	Diff line number	Diff line change
`@@ -181,6 +181,14 @@ impl ProxyShard {`
`181`	`181`	`.estimate_cardinality(filter, hw_measurement_acc)`
`182`	`182`	`.await`
`183`	`183`	`}`
	`184`	`+`
	`185`	`+ pub async fn set_extended_wal_retention(&self) {`
	`186`	`+ self.wrapped_shard.set_extended_wal_retention().await;`
	`187`	`+ }`
	`188`	`+`
	`189`	`+ pub async fn set_normal_wal_retention(&self) {`
	`190`	`+ self.wrapped_shard.set_normal_wal_retention().await;`
	`191`	`+ }`
`184`	`192`	`}`
`185`	`193`
`186`	`194`	`#[async_trait]`