use std::sync::Arc;
use dashmap::DashMap;
use crate::core::registry::IndexId;
use crate::service::persistence::{warmboot_sort_key, PersistedIndex};
pub fn select_warmboot_entries(
entries: Vec<PersistedIndex>,
max_n: Option<usize>,
) -> (Vec<PersistedIndex>, Vec<PersistedIndex>) {
let Some(n) = max_n else {
return (entries, Vec::new());
};
if n == 0 {
return (Vec::new(), entries);
}
if entries.len() <= n {
return (entries, Vec::new());
}
let mut sorted = entries;
sorted.sort_by(|a, b| {
let ka = warmboot_sort_key(a);
let kb = warmboot_sort_key(b);
kb.cmp(&ka).then_with(|| a.id.cmp(&b.id))
});
let cold = sorted.split_off(n);
(sorted, cold)
}
#[derive(Clone, Default)]
pub struct ColdIndexStore {
pub(crate) entries: Arc<DashMap<IndexId, PersistedIndex>>,
loading_gates: Arc<DashMap<IndexId, Arc<tokio::sync::Mutex<()>>>>,
failed_entries: Arc<DashMap<IndexId, ()>>,
}
impl ColdIndexStore {
pub fn new() -> Self {
Self::default()
}
pub fn register_cold_entries(&self, entries: Vec<PersistedIndex>) {
for entry in entries {
let id = IndexId::new(entry.id.clone());
self.entries.insert(id, entry);
}
}
pub fn contains(&self, id: &IndexId) -> bool {
self.entries.contains_key(id)
}
pub fn is_failed(&self, id: &IndexId) -> bool {
self.failed_entries.contains_key(id)
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn failed_len(&self) -> usize {
self.failed_entries.len()
}
pub fn count_matching<I>(&self, ids: I) -> usize
where
I: IntoIterator,
I::Item: AsRef<str>,
{
ids.into_iter()
.filter(|s| self.entries.contains_key(&IndexId::new(s.as_ref())))
.count()
}
pub fn mark_loaded(&self, id: &IndexId) {
self.entries.remove(id);
self.loading_gates.remove(id);
}
pub fn mark_failed(&self, id: &IndexId) {
self.entries.remove(id);
self.loading_gates.remove(id);
self.failed_entries.insert(id.clone(), ());
}
pub fn loading_gate(&self, id: &IndexId) -> Option<Arc<tokio::sync::Mutex<()>>> {
if !self.entries.contains_key(id) {
return None;
}
Some(
self.loading_gates
.entry(id.clone())
.or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
.clone(),
)
}
}