use super::embed_repair::EmbedHealth;
use crate::json_rmw;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::time::Duration;
use uuid::Uuid;
pub const CHUNK_MARKER_FILE: &str = "chunk_backfill.json";
const MARKER_VERSION: u32 = 1;
pub const MARKER_RESCAN_AFTER: Duration = Duration::from_secs(24 * 3600);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ChunkState {
Backfilling,
Complete,
CompleteWithFailures,
Unknown,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ChunkProgress {
pub total: usize,
pub done: usize,
pub pending: usize,
pub failed: usize,
pub state: ChunkState,
}
impl EmbedHealth {
pub fn chunk_progress(&self) -> ChunkProgress {
let needing: HashSet<Uuid> = self
.missing_vector_ids
.iter()
.chain(&self.short_chunk_ids)
.copied()
.collect();
let ledgered: HashSet<Uuid> = self.recorded_failures.iter().map(|f| f.drawer_id).collect();
let failed = needing.iter().filter(|id| ledgered.contains(id)).count();
let pending = needing.len() - failed;
let state = if self.chunk_check.unavailable_reason().is_some() {
ChunkState::Unknown
} else if pending > 0 {
ChunkState::Backfilling
} else if failed > 0 {
ChunkState::CompleteWithFailures
} else {
ChunkState::Complete
};
ChunkProgress {
total: self.drawer_count,
done: self.drawer_count.saturating_sub(needing.len()),
pending,
failed,
state,
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ChunkMarker {
pub version: u32,
pub window: Option<usize>,
pub total: usize,
pub pending: usize,
pub failed: usize,
pub state: Option<ChunkState>,
pub updated_at: Option<DateTime<Utc>>,
pub last_error: Option<String>,
}
impl ChunkMarker {
pub fn from_progress(
progress: &ChunkProgress,
window: Option<usize>,
last_error: Option<String>,
now: DateTime<Utc>,
) -> Self {
Self {
version: MARKER_VERSION,
window,
total: progress.total,
pending: progress.pending,
failed: progress.failed,
state: Some(progress.state),
updated_at: Some(now),
last_error,
}
}
pub fn read(dir: &Path) -> Option<Self> {
let bytes = std::fs::read(marker_path(dir)).ok()?;
match serde_json::from_slice::<Self>(&bytes) {
Ok(marker) if marker.version == MARKER_VERSION => Some(marker),
Ok(_) => None,
Err(e) => {
tracing::debug!(dir = %dir.display(), "#9735: unparsable chunk marker: {e}");
None
}
}
}
pub fn write(&self, dir: &Path) -> anyhow::Result<()> {
let path = marker_path(dir);
let replace = |path: &Path| {
json_rmw::update(path, |m: &mut ChunkMarker| {
*m = self.clone();
Ok::<(), anyhow::Error>(())
})
};
match replace(&path) {
Ok(()) => Ok(()),
Err(_) if path.is_file() && Self::read(dir).is_none() => {
std::fs::remove_file(&path)?;
replace(&path)
}
Err(e) => Err(e),
}
}
pub fn needs_scan(marker: Option<&Self>, window: Option<usize>, now: DateTime<Utc>) -> bool {
let Some(m) = marker else {
return true;
};
let complete = matches!(
m.state,
Some(ChunkState::Complete | ChunkState::CompleteWithFailures)
);
let stale = m.updated_at.is_none_or(|at| {
now.signed_duration_since(at)
.to_std()
.map_or(true, |age| age >= MARKER_RESCAN_AFTER)
});
m.window != window || !complete || m.pending > 0 || stale
}
}
fn marker_path(dir: &Path) -> PathBuf {
dir.join(CHUNK_MARKER_FILE)
}