use std::collections::BTreeMap;
use std::fmt;
#[cfg(feature = "durable-value-store")]
use std::path::Path;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use crate::cluster::ClusterEpoch;
#[cfg(feature = "durable-value-store")]
use crate::grid::durable_store::DurableValueStore;
use crate::grid::elasticity::RegionId;
use crate::grid::hardening::{ReplicatedValueRecord, ReplicatedValueStore, ValueStoreError};
use crate::grid::persistence_policy::{
PersistencePolicy, PersistencePolicyError, PersistenceRegionPlacement,
};
use crate::grid::EffectiveReplicationMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RecoveryMode {
FullRecoveryOnly,
PartialAllowed,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecoveryPolicy {
pub mode: RecoveryMode,
pub validation_timeout: Duration,
pub data_load_timeout: Duration,
pub auto_remove_stale_data: bool,
}
impl RecoveryPolicy {
pub fn full_recovery_only() -> Self {
Self {
mode: RecoveryMode::FullRecoveryOnly,
validation_timeout: Duration::from_secs(30),
data_load_timeout: Duration::from_secs(30),
auto_remove_stale_data: false,
}
}
pub fn partial_allowed() -> Self {
Self {
mode: RecoveryMode::PartialAllowed,
validation_timeout: Duration::from_secs(30),
data_load_timeout: Duration::from_secs(30),
auto_remove_stale_data: false,
}
}
pub fn with_validation_timeout(mut self, timeout: Duration) -> Self {
self.validation_timeout = timeout;
self
}
pub fn with_data_load_timeout(mut self, timeout: Duration) -> Self {
self.data_load_timeout = timeout;
self
}
pub fn with_auto_remove_stale_data(mut self, auto_remove: bool) -> Self {
self.auto_remove_stale_data = auto_remove;
self
}
}
impl Default for RecoveryPolicy {
fn default() -> Self {
Self::full_recovery_only()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecoveryNamespace {
pub namespace: String,
pub placement: PersistenceRegionPlacement,
pub replication_map: EffectiveReplicationMap,
pub key_prefix: Option<String>,
}
impl RecoveryNamespace {
pub fn new(
namespace: impl Into<String>,
placement: PersistenceRegionPlacement,
replication_map: EffectiveReplicationMap,
) -> Self {
Self {
namespace: namespace.into(),
placement,
replication_map,
key_prefix: None,
}
}
pub fn with_key_prefix(mut self, key_prefix: impl Into<String>) -> Self {
self.key_prefix = Some(key_prefix.into());
self
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecoveredNamespace {
pub namespace: String,
pub persistent: bool,
pub records: BTreeMap<String, ReplicatedValueRecord>,
pub stale_keys: Vec<String>,
pub partial: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecoveryReport {
pub namespaces: BTreeMap<String, RecoveredNamespace>,
pub recovered_record_total: u64,
pub stale_fenced_total: u64,
pub non_persistent_skipped_total: u64,
pub partial_recovery_total: u64,
pub timeout_total: u64,
pub auto_remove_stale_data: bool,
}
impl RecoveryReport {
pub fn record(&self, namespace: &str, key: &str) -> Option<&ReplicatedValueRecord> {
self.namespaces.get(namespace)?.records.get(key)
}
pub fn namespace_persistent(&self, namespace: &str) -> bool {
self.namespaces
.get(namespace)
.map(|report| report.persistent)
.unwrap_or(false)
}
}
pub fn recover_namespaces<S>(
store: &S,
policy: &PersistencePolicy,
local_region: &RegionId,
authority_epoch: ClusterEpoch,
recovery_policy: &RecoveryPolicy,
namespaces: impl IntoIterator<Item = RecoveryNamespace>,
) -> Result<RecoveryReport, RecoveryError>
where
S: ReplicatedValueStore,
{
if recovery_policy.validation_timeout.is_zero() {
return recovery_timeout_or_partial(recovery_policy, RecoveryReport::default());
}
let mut report = RecoveryReport {
auto_remove_stale_data: recovery_policy.auto_remove_stale_data,
..RecoveryReport::default()
};
for request in namespaces {
let resolved = policy
.resolve_for_region(&request.namespace, local_region, &request.placement)
.map_err(RecoveryError::policy)?;
if !resolved.persists() {
report.non_persistent_skipped_total =
report.non_persistent_skipped_total.saturating_add(1);
report.namespaces.insert(
request.namespace.clone(),
RecoveredNamespace {
namespace: request.namespace,
persistent: false,
..RecoveredNamespace::default()
},
);
continue;
}
let scanned = store
.scan_owned(&request.replication_map)
.map_err(RecoveryError::store)?;
if recovery_policy.data_load_timeout.is_zero() && !scanned.is_empty() {
report.timeout_total = report.timeout_total.saturating_add(1);
return recovery_timeout_or_partial(recovery_policy, report);
}
let mut namespace_report = RecoveredNamespace {
namespace: request.namespace.clone(),
persistent: true,
..RecoveredNamespace::default()
};
for (key, record) in scanned {
if let Some(prefix) = &request.key_prefix {
if !key.starts_with(prefix) {
continue;
}
}
if record.epoch < authority_epoch {
namespace_report.stale_keys.push(key);
report.stale_fenced_total = report.stale_fenced_total.saturating_add(1);
continue;
}
namespace_report.records.insert(key, record);
report.recovered_record_total = report.recovered_record_total.saturating_add(1);
}
report
.namespaces
.insert(request.namespace.clone(), namespace_report);
}
Ok(report)
}
#[cfg(feature = "durable-value-store")]
pub fn open_durable_value_store_for_recovery(
path: impl AsRef<Path>,
max_total_bytes: u64,
) -> Result<DurableValueStore, RecoveryError> {
DurableValueStore::open_with_budget(path, max_total_bytes).map_err(RecoveryError::store)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecoveryError {
kind: RecoveryErrorKind,
message: String,
}
impl RecoveryError {
fn policy(error: PersistencePolicyError) -> Self {
Self {
kind: RecoveryErrorKind::Policy,
message: error.to_string(),
}
}
fn store(error: ValueStoreError) -> Self {
Self {
kind: RecoveryErrorKind::Store,
message: error.to_string(),
}
}
fn timeout(message: impl Into<String>) -> Self {
Self {
kind: RecoveryErrorKind::Timeout,
message: message.into(),
}
}
pub fn kind(&self) -> RecoveryErrorKind {
self.kind
}
}
impl fmt::Display for RecoveryError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for RecoveryError {}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RecoveryErrorKind {
Policy,
Store,
Timeout,
}
fn recovery_timeout_or_partial(
policy: &RecoveryPolicy,
mut report: RecoveryReport,
) -> Result<RecoveryReport, RecoveryError> {
match policy.mode {
RecoveryMode::FullRecoveryOnly => Err(RecoveryError::timeout(
"full recovery timed out before persistent namespaces were safely loaded",
)),
RecoveryMode::PartialAllowed => {
report.partial_recovery_total = report.partial_recovery_total.saturating_add(1);
for namespace in report.namespaces.values_mut() {
namespace.partial = true;
}
Ok(report)
}
}
}