use std::future::{Future, Ready, ready};
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Arc;
use anyhow::Result;
use futures::future::Either;
use kvbm_config::{PolicyType, TierOffloadConfig};
use crate::{BlockId, SequenceHash};
use kvbm_logical::blocks::{BlockMetadata, BlockRegistry, ImmutableBlock};
use super::pending::{PendingCheck, PendingTracker};
use crate::object::{ObjectBlockOps, ObjectLockManager};
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub type PolicyFuture<'a> = Either<Ready<Result<bool>>, BoxFuture<'a, Result<bool>>>;
pub type PolicyBatchFuture<'a> = Either<Ready<Result<Vec<bool>>>, BoxFuture<'a, Result<Vec<bool>>>>;
#[inline]
pub fn sync_result(result: Result<bool>) -> PolicyFuture<'static> {
Either::Left(ready(result))
}
#[inline]
pub fn sync_batch_result(result: Result<Vec<bool>>) -> PolicyBatchFuture<'static> {
Either::Left(ready(result))
}
#[inline]
pub fn async_result<'a, F>(future: F) -> PolicyFuture<'a>
where
F: Future<Output = Result<bool>> + Send + 'a,
{
Either::Right(Box::pin(future))
}
#[inline]
pub fn async_batch_result<'a, F>(future: F) -> PolicyBatchFuture<'a>
where
F: Future<Output = Result<Vec<bool>>> + Send + 'a,
{
Either::Right(Box::pin(future))
}
pub trait PresenceChecker: Send + Sync {
fn check_presence(
&self,
keys: Vec<SequenceHash>,
) -> BoxFuture<'static, Vec<(SequenceHash, bool)>>;
}
pub struct S3PresenceChecker {
object_ops: Arc<dyn ObjectBlockOps>,
}
impl S3PresenceChecker {
pub fn new(object_ops: Arc<dyn ObjectBlockOps>) -> Self {
Self { object_ops }
}
}
impl PresenceChecker for S3PresenceChecker {
fn check_presence(
&self,
keys: Vec<SequenceHash>,
) -> BoxFuture<'static, Vec<(SequenceHash, bool)>> {
let future = self.object_ops.has_blocks(keys);
Box::pin(async move {
let results = future.await;
results
.into_iter()
.map(|(hash, size_opt)| (hash, size_opt.is_some()))
.collect()
})
}
}
#[derive(Debug)]
pub struct EvalContext<T: BlockMetadata> {
pub block_id: BlockId,
pub sequence_hash: SequenceHash,
pub block: Option<ImmutableBlock<T>>,
}
impl<T: BlockMetadata> EvalContext<T> {
pub fn new(block: ImmutableBlock<T>) -> Self {
Self {
block_id: block.block_id(),
sequence_hash: block.sequence_hash(),
block: Some(block),
}
}
pub fn from_weak(block_id: BlockId, sequence_hash: SequenceHash) -> Self {
Self {
block_id,
sequence_hash,
block: None,
}
}
pub fn from_external(block_id: BlockId, sequence_hash: SequenceHash) -> Self {
Self {
block_id,
sequence_hash,
block: None,
}
}
}
pub trait OffloadPolicy<T: BlockMetadata>: Send + Sync {
fn name(&self) -> &str;
fn evaluate<'a>(&'a self, ctx: &'a EvalContext<T>) -> PolicyFuture<'a>;
fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<T>]) -> PolicyBatchFuture<'a> {
let contexts_clone: Vec<_> = contexts.iter().collect();
async_batch_result(async move {
let mut results = Vec::with_capacity(contexts_clone.len());
for ctx in contexts_clone {
let result = match self.evaluate(ctx) {
Either::Left(ready) => ready.await,
Either::Right(boxed) => boxed.await,
};
results.push(result?);
}
Ok(results)
})
}
}
pub struct PresenceFilter<Src: BlockMetadata, Dst: BlockMetadata> {
registry: Arc<BlockRegistry>,
pending_tracker: Option<Arc<PendingTracker>>,
_marker: PhantomData<(Src, Dst)>,
}
impl<Src: BlockMetadata, Dst: BlockMetadata> PresenceFilter<Src, Dst> {
pub fn new(registry: Arc<BlockRegistry>) -> Self {
Self {
registry,
pending_tracker: None,
_marker: PhantomData,
}
}
pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
self.pending_tracker = Some(tracker);
self
}
pub fn pending_tracker(&self) -> Option<&Arc<PendingTracker>> {
self.pending_tracker.as_ref()
}
}
impl<Src: BlockMetadata, Dst: BlockMetadata> OffloadPolicy<Src> for PresenceFilter<Src, Dst> {
fn name(&self) -> &str {
"PresenceFilter"
}
fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
let presence = self.registry.check_presence::<Dst>(&[ctx.sequence_hash]);
if presence[0].1 {
return sync_result(Ok(false)); }
if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
return sync_result(Ok(false)); }
sync_result(Ok(true)) }
fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
if contexts.is_empty() {
return sync_batch_result(Ok(Vec::new()));
}
let hashes: Vec<SequenceHash> = contexts.iter().map(|ctx| ctx.sequence_hash).collect();
let presence = self.registry.check_presence::<Dst>(&hashes);
let results: Vec<bool> = presence
.into_iter()
.map(|(hash, present)| {
if present {
return false;
}
if self.pending_tracker.is_hash_pending(&hash) {
return false;
}
true
})
.collect();
sync_batch_result(Ok(results))
}
}
pub struct PresenceAndLFUFilter<Src: BlockMetadata, Dst: BlockMetadata> {
registry: Arc<BlockRegistry>,
min_lfu_count: u32,
pending_tracker: Option<Arc<PendingTracker>>,
_marker: PhantomData<(Src, Dst)>,
}
impl<Src: BlockMetadata, Dst: BlockMetadata> PresenceAndLFUFilter<Src, Dst> {
pub fn new(registry: Arc<BlockRegistry>, min_lfu_count: u32) -> Self {
Self {
registry,
min_lfu_count,
pending_tracker: None,
_marker: PhantomData,
}
}
pub fn with_default_threshold(registry: Arc<BlockRegistry>) -> Self {
Self::new(registry, 8)
}
pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
self.pending_tracker = Some(tracker);
self
}
}
impl<Src: BlockMetadata, Dst: BlockMetadata> OffloadPolicy<Src> for PresenceAndLFUFilter<Src, Dst> {
fn name(&self) -> &str {
"PresenceAndLFUFilter"
}
fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
let presence = self.registry.check_presence::<Dst>(&[ctx.sequence_hash]);
if presence[0].1 {
return sync_result(Ok(false));
}
if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
return sync_result(Ok(false));
}
if let Some(tracker) = self.registry.frequency_tracker() {
let count = tracker.count(ctx.sequence_hash.as_u128());
return sync_result(Ok(count > self.min_lfu_count));
}
sync_result(Ok(true))
}
fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
if contexts.is_empty() {
return sync_batch_result(Ok(Vec::new()));
}
let hashes: Vec<SequenceHash> = contexts.iter().map(|ctx| ctx.sequence_hash).collect();
let presence = self.registry.check_presence::<Dst>(&hashes);
let freq_tracker = self.registry.frequency_tracker();
let min_lfu = self.min_lfu_count;
let results: Vec<bool> = presence
.into_iter()
.zip(contexts.iter())
.map(|((hash, present), ctx)| {
if present {
return false;
}
if self.pending_tracker.is_hash_pending(&hash) {
return false;
}
if let Some(ref t) = freq_tracker {
let count = t.count(ctx.sequence_hash.as_u128());
count > min_lfu
} else {
true }
})
.collect();
sync_batch_result(Ok(results))
}
}
pub struct ObjectPresenceFilter<Src: BlockMetadata> {
presence_checker: Arc<dyn PresenceChecker>,
pending_tracker: Option<Arc<PendingTracker>>,
_marker: PhantomData<Src>,
}
impl<Src: BlockMetadata> ObjectPresenceFilter<Src> {
pub fn new(presence_checker: Arc<dyn PresenceChecker>) -> Self {
Self {
presence_checker,
pending_tracker: None,
_marker: PhantomData,
}
}
pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
self.pending_tracker = Some(tracker);
self
}
pub fn pending_tracker(&self) -> Option<&Arc<PendingTracker>> {
self.pending_tracker.as_ref()
}
}
impl<Src: BlockMetadata> OffloadPolicy<Src> for ObjectPresenceFilter<Src> {
fn name(&self) -> &str {
"ObjectPresenceFilter"
}
fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
return sync_result(Ok(false)); }
let checker = self.presence_checker.clone();
let hash = ctx.sequence_hash;
async_result(async move {
let results = checker.check_presence(vec![hash]).await;
let exists = results
.into_iter()
.next()
.map(|(_, exists)| exists)
.unwrap_or(false);
Ok(!exists) })
}
fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
if contexts.is_empty() {
return sync_batch_result(Ok(Vec::new()));
}
let mut pending_status: Vec<bool> = Vec::with_capacity(contexts.len());
let mut hashes_to_check: Vec<SequenceHash> = Vec::new();
let mut hash_indices: Vec<usize> = Vec::new();
for (i, ctx) in contexts.iter().enumerate() {
if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
pending_status.push(true); } else {
pending_status.push(false);
hashes_to_check.push(ctx.sequence_hash);
hash_indices.push(i);
}
}
if hashes_to_check.is_empty() {
return sync_batch_result(Ok(vec![false; contexts.len()]));
}
let checker = self.presence_checker.clone();
let num_contexts = contexts.len();
async_batch_result(async move {
let presence_results = checker.check_presence(hashes_to_check).await;
let mut results = vec![false; num_contexts];
for (check_idx, original_idx) in hash_indices.into_iter().enumerate() {
if let Some((_, exists)) = presence_results.get(check_idx) {
results[original_idx] = !*exists;
}
}
Ok(results)
})
}
}
pub struct ObjectLockPresenceFilter<Src: BlockMetadata> {
lock_manager: Arc<dyn ObjectLockManager>,
pending_tracker: Option<Arc<PendingTracker>>,
_marker: PhantomData<Src>,
}
impl<Src: BlockMetadata> ObjectLockPresenceFilter<Src> {
pub fn new(lock_manager: Arc<dyn ObjectLockManager>) -> Self {
Self {
lock_manager,
pending_tracker: None,
_marker: PhantomData,
}
}
pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
self.pending_tracker = Some(tracker);
self
}
pub fn pending_tracker(&self) -> Option<&Arc<PendingTracker>> {
self.pending_tracker.as_ref()
}
pub fn lock_manager(&self) -> &Arc<dyn ObjectLockManager> {
&self.lock_manager
}
}
impl<Src: BlockMetadata> OffloadPolicy<Src> for ObjectLockPresenceFilter<Src> {
fn name(&self) -> &str {
"ObjectLockPresenceFilter"
}
fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
return sync_result(Ok(false)); }
let lock_manager = self.lock_manager.clone();
let hash = ctx.sequence_hash;
async_result(async move {
match lock_manager.has_meta(hash).await {
Ok(true) => {
tracing::debug!(?hash, "Block already offloaded (meta exists)");
return Ok(false); }
Ok(false) => {
}
Err(e) => {
tracing::warn!(?hash, error = %e, "Error checking meta file");
return Ok(false); }
}
match lock_manager.try_acquire_lock(hash).await {
Ok(true) => {
tracing::debug!(?hash, "Lock acquired");
Ok(true) }
Ok(false) => {
tracing::debug!(?hash, "Lock held by another instance");
Ok(false) }
Err(e) => {
tracing::warn!(?hash, error = %e, "Error acquiring lock");
Ok(false) }
}
})
}
fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
if contexts.is_empty() {
return sync_batch_result(Ok(Vec::new()));
}
let mut pending_mask: Vec<bool> = Vec::with_capacity(contexts.len());
let mut to_check: Vec<(usize, SequenceHash)> = Vec::new();
for (i, ctx) in contexts.iter().enumerate() {
if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
pending_mask.push(true);
} else {
pending_mask.push(false);
to_check.push((i, ctx.sequence_hash));
}
}
if to_check.is_empty() {
return sync_batch_result(Ok(vec![false; contexts.len()]));
}
let lock_manager = self.lock_manager.clone();
let num_contexts = contexts.len();
async_batch_result(async move {
let mut results = vec![false; num_contexts];
for (original_idx, hash) in to_check {
let has_meta = match lock_manager.has_meta(hash).await {
Ok(has) => has,
Err(e) => {
tracing::warn!(?hash, error = %e, "Error checking meta file");
continue; }
};
if has_meta {
tracing::debug!(?hash, "Block already offloaded (meta exists)");
continue; }
match lock_manager.try_acquire_lock(hash).await {
Ok(true) => {
tracing::debug!(?hash, "Lock acquired");
results[original_idx] = true; }
Ok(false) => {
tracing::debug!(?hash, "Lock held by another instance");
}
Err(e) => {
tracing::warn!(?hash, error = %e, "Error acquiring lock");
}
}
}
Ok(results)
})
}
}
pub struct AllOfPolicy<T: BlockMetadata> {
policies: Vec<Arc<dyn OffloadPolicy<T>>>,
}
impl<T: BlockMetadata> AllOfPolicy<T> {
pub fn new(policies: Vec<Arc<dyn OffloadPolicy<T>>>) -> Self {
Self { policies }
}
pub fn with(mut self, policy: Arc<dyn OffloadPolicy<T>>) -> Self {
self.policies.push(policy);
self
}
}
impl<T: BlockMetadata> OffloadPolicy<T> for AllOfPolicy<T> {
fn name(&self) -> &str {
"AllOfPolicy"
}
fn evaluate<'a>(&'a self, ctx: &'a EvalContext<T>) -> PolicyFuture<'a> {
let policies = &self.policies;
async_result(async move {
for policy in policies {
let result = match policy.evaluate(ctx) {
Either::Left(ready) => ready.await,
Either::Right(boxed) => boxed.await,
};
if !result? {
return Ok(false);
}
}
Ok(true)
})
}
}
pub struct AnyOfPolicy<T: BlockMetadata> {
policies: Vec<Arc<dyn OffloadPolicy<T>>>,
}
impl<T: BlockMetadata> AnyOfPolicy<T> {
pub fn new(policies: Vec<Arc<dyn OffloadPolicy<T>>>) -> Self {
Self { policies }
}
pub fn with(mut self, policy: Arc<dyn OffloadPolicy<T>>) -> Self {
self.policies.push(policy);
self
}
}
impl<T: BlockMetadata> OffloadPolicy<T> for AnyOfPolicy<T> {
fn name(&self) -> &str {
"AnyOfPolicy"
}
fn evaluate<'a>(&'a self, ctx: &'a EvalContext<T>) -> PolicyFuture<'a> {
if self.policies.is_empty() {
return sync_result(Ok(true)); }
let policies = &self.policies;
async_result(async move {
for policy in policies {
let result = match policy.evaluate(ctx) {
Either::Left(ready) => ready.await,
Either::Right(boxed) => boxed.await,
};
if result? {
return Ok(true);
}
}
Ok(false)
})
}
}
pub struct PassAllPolicy<T: BlockMetadata> {
_marker: PhantomData<T>,
}
impl<T: BlockMetadata> PassAllPolicy<T> {
pub fn new() -> Self {
Self {
_marker: PhantomData,
}
}
}
impl<T: BlockMetadata> Default for PassAllPolicy<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: BlockMetadata> OffloadPolicy<T> for PassAllPolicy<T> {
fn name(&self) -> &str {
"PassAllPolicy"
}
fn evaluate<'a>(&'a self, _ctx: &'a EvalContext<T>) -> PolicyFuture<'a> {
sync_result(Ok(true))
}
fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<T>]) -> PolicyBatchFuture<'a> {
sync_batch_result(Ok(vec![true; contexts.len()]))
}
}
pub fn create_policy_from_config<Src, Dst>(
config: &TierOffloadConfig,
registry: Arc<BlockRegistry>,
pending_tracker: Option<Arc<PendingTracker>>,
) -> Arc<dyn OffloadPolicy<Src>>
where
Src: BlockMetadata + 'static,
Dst: BlockMetadata + 'static,
{
if config.policies.is_empty() {
return Arc::new(PassAllPolicy::<Src>::new());
}
let policies: Vec<Arc<dyn OffloadPolicy<Src>>> = config
.policies
.iter()
.map(|policy_type| -> Arc<dyn OffloadPolicy<Src>> {
match policy_type {
PolicyType::PassAll => Arc::new(PassAllPolicy::<Src>::new()),
PolicyType::Presence => {
let mut filter = PresenceFilter::<Src, Dst>::new(registry.clone());
if let Some(tracker) = &pending_tracker {
filter = filter.with_pending_tracker(tracker.clone());
}
Arc::new(filter)
}
PolicyType::PresenceLfu => {
let mut filter = PresenceAndLFUFilter::<Src, Dst>::new(
registry.clone(),
config.presence_lfu.min_lfu_count,
);
if let Some(tracker) = &pending_tracker {
filter = filter.with_pending_tracker(tracker.clone());
}
Arc::new(filter)
}
}
})
.collect();
if policies.len() == 1 {
policies.into_iter().next().unwrap()
} else {
Arc::new(AllOfPolicy::new(policies))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pass_all_policy() {
let _policy: PassAllPolicy<()> = PassAllPolicy::new();
}
#[test]
fn test_all_of_policy_creation() {
let policies: Vec<Arc<dyn OffloadPolicy<()>>> = vec![Arc::new(PassAllPolicy::new())];
let composite = AllOfPolicy::new(policies);
assert_eq!(composite.name(), "AllOfPolicy");
}
#[test]
fn test_any_of_policy_creation() {
let policies: Vec<Arc<dyn OffloadPolicy<()>>> = vec![Arc::new(PassAllPolicy::new())];
let composite = AnyOfPolicy::new(policies);
assert_eq!(composite.name(), "AnyOfPolicy");
}
#[tokio::test]
async fn test_sync_result_zero_alloc() {
let future = sync_result(Ok(true));
assert!(matches!(future, Either::Left(_)));
let result = match future {
Either::Left(ready) => ready.await,
Either::Right(_) => unreachable!(),
};
assert!(result.unwrap());
}
#[tokio::test]
async fn test_async_result_boxes() {
let future = async_result(async { Ok(false) });
assert!(matches!(future, Either::Right(_)));
let result = match future {
Either::Left(_) => unreachable!(),
Either::Right(boxed) => boxed.await,
};
assert!(!result.unwrap());
}
#[test]
fn test_pending_tracker_wiring() {
use super::PendingTracker;
let tracker = Arc::new(PendingTracker::new());
let registry = Arc::new(BlockRegistry::new());
let filter: PresenceFilter<(), ()> =
PresenceFilter::new(registry).with_pending_tracker(tracker.clone());
assert!(filter.pending_tracker().is_some());
assert!(Arc::ptr_eq(filter.pending_tracker().unwrap(), &tracker));
}
#[test]
fn test_pending_tracker_wiring_lfu() {
use super::PendingTracker;
let tracker = Arc::new(PendingTracker::new());
let registry = Arc::new(BlockRegistry::new());
let filter: PresenceAndLFUFilter<(), ()> =
PresenceAndLFUFilter::new(registry, 8).with_pending_tracker(tracker);
assert_eq!(filter.name(), "PresenceAndLFUFilter");
}
}