use std::{collections::BTreeMap, error::Error, fmt};
use ic_memory::{BootstrapAdmission, BootstrapAdmissionError, StableKey};
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum MemoryBootstrapAdmissionError {
UnsupportedKey(String),
InvalidDeclaration(String),
IncompleteRoles {
namespace: String,
store: Option<String>,
},
NamespaceRemoved(String),
HistoricalJournal(BootstrapAdmissionError),
}
impl fmt::Display for MemoryBootstrapAdmissionError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::UnsupportedKey(key) => write!(f, "unsupported IcyDB allocation key: {key}"),
Self::InvalidDeclaration(key) => {
write!(
f,
"IcyDB allocation requires a namespace-authorized logical request: {key}"
)
}
Self::IncompleteRoles { namespace, store } => {
write!(f, "incomplete IcyDB allocation roles for {namespace}")?;
if let Some(store) = store {
write!(f, " store {store}")?;
}
Ok(())
}
Self::NamespaceRemoved(namespace) => {
write!(f, "historical IcyDB namespace is not declared: {namespace}")
}
Self::HistoricalJournal(error) => error.fmt(f),
}
}
}
impl Error for MemoryBootstrapAdmissionError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::HistoricalJournal(error) => Some(error),
_ => None,
}
}
}
pub fn prepare_memory_bootstrap(
admission: &mut BootstrapAdmission<'_>,
) -> Result<(), MemoryBootstrapAdmissionError> {
let mut groups = BTreeMap::<(&str, Option<&str>), u8>::new();
for registration in admission.declarations().registered_declarations() {
let key = registration.declaration().stable_key();
if allocation_role(key)?.is_some() {
return Err(MemoryBootstrapAdmissionError::InvalidDeclaration(
key.as_str().to_owned(),
));
}
}
for request in admission.declarations().requests() {
let Some(role) = allocation_role(request.stable_key())? else {
continue;
};
if request.authority().strip_prefix("icydb.") != Some(role.namespace) {
return Err(MemoryBootstrapAdmissionError::InvalidDeclaration(
request.stable_key().as_str().to_owned(),
));
}
*groups.entry((role.namespace, role.store)).or_default() |= role.bit;
}
for (&(namespace, store), &roles) in &groups {
if !groups.contains_key(&(namespace, None)) {
return Err(incomplete_roles(namespace, None));
}
let required = if store.is_some() { 0b1111 } else { 0b111 };
if roles != required {
return Err(incomplete_roles(namespace, store));
}
}
let mut journals = Vec::new();
for record in admission.recovered_allocations() {
let Some(role) = allocation_role(record.stable_key)? else {
continue;
};
if !groups.contains_key(&(role.namespace, None)) {
return Err(MemoryBootstrapAdmissionError::NamespaceRemoved(
role.namespace.to_owned(),
));
}
if role.store.is_some()
&& role.bit == 0b1000
&& !groups.contains_key(&(role.namespace, role.store))
{
journals.push((
format!("icydb.{}", role.namespace),
record.stable_key.clone(),
));
}
}
for (authority, key) in journals {
admission
.include_historical(&authority, key.as_str())
.map_err(MemoryBootstrapAdmissionError::HistoricalJournal)?;
}
Ok(())
}
struct AllocationRole<'a> {
namespace: &'a str,
store: Option<&'a str>,
bit: u8,
}
fn allocation_role(
key: &StableKey,
) -> Result<Option<AllocationRole<'_>>, MemoryBootstrapAdmissionError> {
let mut parts = key.as_str().split('.');
if parts.next() != Some("icydb") {
return Ok(None);
}
let namespace = parts.next();
let role = match (
parts.next(),
parts.next(),
parts.next(),
parts.next(),
parts.next(),
) {
(Some("commit"), Some("control"), Some("v1"), None, None) => Some((None, 0b001)),
(Some("startup"), Some("control"), Some("v1"), None, None) => Some((None, 0b010)),
(Some("integrity"), Some("progress"), Some("v1"), None, None) => Some((None, 0b100)),
(Some("store"), Some(store), Some(role), Some("v1"), None) => match role {
"data" => Some((Some(store), 0b0001)),
"index" => Some((Some(store), 0b0010)),
"schema" => Some((Some(store), 0b0100)),
"journal" => Some((Some(store), 0b1000)),
_ => None,
},
_ => None,
};
match (namespace, role) {
(Some(namespace), Some((store, bit))) => Ok(Some(AllocationRole {
namespace,
store,
bit,
})),
_ => Err(MemoryBootstrapAdmissionError::UnsupportedKey(
key.as_str().to_owned(),
)),
}
}
fn incomplete_roles(namespace: &str, store: Option<&str>) -> MemoryBootstrapAdmissionError {
MemoryBootstrapAdmissionError::IncompleteRoles {
namespace: namespace.to_owned(),
store: store.map(str::to_owned),
}
}