use std::collections::HashMap;
use rushdb_types::ids::{TenantId, TenantScope};
use rushdb_types::keys;
use crate::{StorageError, Store};
const CACHE_ENTRIES: usize = 65_536;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TenantState {
Active,
Purged,
}
impl TenantState {
fn to_byte(self) -> u8 {
match self {
Self::Active => 0,
Self::Purged => 1,
}
}
fn from_byte(byte: u8) -> Result<Self, StorageError> {
match byte {
0 => Ok(Self::Active),
1 => Ok(Self::Purged),
_ => Err(StorageError::CorruptLog),
}
}
}
#[must_use]
pub fn encode_by_scope(tenant: TenantId, state: TenantState) -> Vec<u8> {
let mut out = tenant.to_be_bytes().to_vec();
out.push(state.to_byte());
out
}
pub fn decode_by_scope(raw: &[u8]) -> Result<(TenantId, TenantState), StorageError> {
let (ordinal, state) = raw.split_at_checked(4).ok_or(StorageError::CorruptLog)?;
let ordinal: [u8; 4] = ordinal.try_into().map_err(|_| StorageError::CorruptLog)?;
let [state] = state else {
return Err(StorageError::CorruptLog);
};
Ok((
TenantId(u32::from_be_bytes(ordinal)),
TenantState::from_byte(*state)?,
))
}
#[must_use]
pub fn encode_by_ord(scope: &TenantScope, state: TenantState) -> Vec<u8> {
let mut out = scope.0.to_vec();
out.push(state.to_byte());
out
}
pub fn decode_by_ord(raw: &[u8]) -> Result<(TenantScope, TenantState), StorageError> {
let (digest, state) = raw.split_at_checked(32).ok_or(StorageError::CorruptLog)?;
let digest: [u8; 32] = digest.try_into().map_err(|_| StorageError::CorruptLog)?;
let [state] = state else {
return Err(StorageError::CorruptLog);
};
Ok((TenantScope(digest), TenantState::from_byte(*state)?))
}
pub fn decode_seq(raw: Option<&[u8]>) -> Result<u32, StorageError> {
raw.map_or(Ok(TenantId::FIRST_ASSIGNED), |raw| {
raw.try_into()
.map(u32::from_be_bytes)
.map_err(|_| StorageError::CorruptLog)
})
}
#[derive(Default)]
pub(crate) struct TenantCache {
entries: HashMap<TenantScope, TenantId>,
}
impl Store {
pub fn resolve_tenant(&self, scope: &TenantScope) -> Result<Option<TenantId>, StorageError> {
if let Some(tenant) = self
.health
.guard(self.tenant_cache.lock())?
.entries
.get(scope)
{
return Ok(Some(*tenant));
}
let Some(raw) = self.try_get(&keys::tenant_by_scope_key(scope))? else {
return Ok(None);
};
let (tenant, state) = decode_by_scope(&raw)?;
if state != TenantState::Active {
return Ok(None);
}
let mut cache = self.health.guard(self.tenant_cache.lock())?;
if cache.entries.len() >= CACHE_ENTRIES {
cache.entries.clear();
}
cache.entries.insert(*scope, tenant);
Ok(Some(tenant))
}
pub fn tenant_scope(
&self,
tenant: TenantId,
) -> Result<Option<(TenantScope, TenantState)>, StorageError> {
self.try_get(&keys::tenant_by_ord_key(tenant))?
.map(|raw| decode_by_ord(&raw))
.transpose()
}
pub fn active_tenants(&self) -> Result<Vec<(TenantId, TenantScope)>, StorageError> {
let prefix = keys::tenant_by_ord_prefix();
let mut out = Vec::new();
for (key, raw) in self.try_range_prefix(&prefix)? {
let ordinal: [u8; 4] = key
.get(prefix.len()..)
.and_then(|bytes| bytes.try_into().ok())
.ok_or(StorageError::CorruptLog)?;
let (scope, state) = decode_by_ord(&raw)?;
if state == TenantState::Active {
out.push((TenantId(u32::from_be_bytes(ordinal)), scope));
}
}
Ok(out)
}
pub fn forget_tenant(&self, scope: &TenantScope) {
self.health
.cleanup(self.tenant_cache.lock())
.entries
.remove(scope);
}
}