use super::identity::validate_digest;
use super::result::{
EvaluationRecordV1, EvaluationResultSink, EvaluationStoreError, EvaluationWriteOutcomeV1,
};
use super::ExecutionTargetV1;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;
pub const EVALUATION_RESULT_STORE_SCHEMA_V1: &str = "a3s.code.evaluation-result-store.v1";
pub const EVALUATION_RESULT_STORE_DEFAULT_MAX_RECORDS: usize = 4096;
pub const EVALUATION_RESULT_STORE_MAX_BYTES: usize = 64 * 1024 * 1024;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct EvaluationResultStoreFileV1 {
schema: String,
records: Vec<EvaluationRecordV1>,
}
#[derive(Debug, Clone)]
pub struct FileEvaluationResultStore {
root: PathBuf,
max_records: usize,
process_lock: Arc<Mutex<()>>,
}
impl FileEvaluationResultStore {
pub fn new(root: impl Into<PathBuf>) -> Self {
Self {
root: root.into(),
max_records: EVALUATION_RESULT_STORE_DEFAULT_MAX_RECORDS,
process_lock: Arc::new(Mutex::new(())),
}
}
pub fn with_max_records(
root: impl Into<PathBuf>,
max_records: usize,
) -> Result<Self, EvaluationStoreError> {
if max_records == 0 {
return Err(EvaluationStoreError::InvalidField("max_records"));
}
Ok(Self {
root: root.into(),
max_records,
process_lock: Arc::new(Mutex::new(())),
})
}
pub async fn open(root: impl Into<PathBuf>) -> Result<Self, EvaluationStoreError> {
let store = Self::new(root);
store.validate_store().await?;
Ok(store)
}
pub async fn open_with_max_records(
root: impl Into<PathBuf>,
max_records: usize,
) -> Result<Self, EvaluationStoreError> {
let store = Self::with_max_records(root, max_records)?;
store.validate_store().await?;
Ok(store)
}
pub fn root(&self) -> &Path {
&self.root
}
pub fn max_records(&self) -> usize {
self.max_records
}
fn data_path(&self) -> PathBuf {
self.root.join("evaluation-results.json")
}
fn lock_path(&self) -> PathBuf {
self.root.join(".evaluation-results.lock")
}
pub async fn validate_store(&self) -> Result<usize, EvaluationStoreError> {
Ok(self.read_records().await?.len())
}
pub async fn get_checked(
&self,
record_digest: &str,
) -> Result<Option<EvaluationRecordV1>, EvaluationStoreError> {
validate_digest(record_digest)
.map_err(|_| EvaluationStoreError::InvalidField("record_digest"))?;
let records = self.read_records().await?;
Ok(records
.into_iter()
.find(|record| record.record_digest == record_digest))
}
pub async fn list_for_target_checked(
&self,
target: &ExecutionTargetV1,
) -> Result<Vec<EvaluationRecordV1>, EvaluationStoreError> {
target
.validate()
.map_err(|_| EvaluationStoreError::InvalidField("target"))?;
let records = self.read_records().await?;
Ok(records
.into_iter()
.filter(|record| record.result.target == *target)
.collect())
}
async fn acquire_file_lock(&self) -> Result<std::fs::File, EvaluationStoreError> {
tokio::fs::create_dir_all(&self.root)
.await
.map_err(|error| storage_error("create evaluation result directory", error))?;
let path = self.lock_path();
tokio::task::spawn_blocking(move || {
use fs2::FileExt;
let file = std::fs::OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(&path)
.map_err(|error| storage_error("open evaluation result lock", error))?;
file.lock_exclusive()
.map_err(|error| storage_error("lock evaluation result store", error))?;
Ok(file)
})
.await
.map_err(|error| EvaluationStoreError::Storage(format!("lock task failed: {error}")))?
}
async fn read_records(&self) -> Result<Vec<EvaluationRecordV1>, EvaluationStoreError> {
read_records_from_path(&self.data_path(), self.max_records).await
}
async fn mutate(
&self,
record: EvaluationRecordV1,
) -> Result<EvaluationWriteOutcomeV1, EvaluationStoreError> {
record.validate()?;
let _process_guard = self.process_lock.lock().await;
let _file_guard = self.acquire_file_lock().await?;
let mut records = self.read_records().await?;
if let Some(existing) = records
.iter()
.find(|existing| existing.record_digest == record.record_digest)
{
if existing == &record {
return Ok(EvaluationWriteOutcomeV1 {
written: false,
replayed: true,
});
}
return Err(EvaluationStoreError::Conflict);
}
if records.iter().any(|existing| {
existing.result.target == record.result.target
&& existing.result.evaluator_id == record.result.evaluator_id
&& existing.result.auxiliary_run_id == record.result.auxiliary_run_id
}) {
return Err(EvaluationStoreError::Conflict);
}
records.push(record);
while records.len() > self.max_records {
records.remove(0);
}
self.publish_records(&records).await?;
Ok(EvaluationWriteOutcomeV1 {
written: true,
replayed: false,
})
}
async fn publish_records(
&self,
records: &[EvaluationRecordV1],
) -> Result<(), EvaluationStoreError> {
let bytes = serde_json::to_vec(&EvaluationResultStoreFileV1 {
schema: EVALUATION_RESULT_STORE_SCHEMA_V1.to_string(),
records: records.to_vec(),
})
.map_err(|error| EvaluationStoreError::Serialization(error.to_string()))?;
if bytes.len() > EVALUATION_RESULT_STORE_MAX_BYTES {
return Err(EvaluationStoreError::SizeLimit);
}
tokio::fs::create_dir_all(&self.root)
.await
.map_err(|error| storage_error("create evaluation result directory", error))?;
let path = self.data_path();
let temp = path.with_extension(format!("tmp-{}", uuid::Uuid::new_v4()));
let result = async {
let mut file = tokio::fs::File::create(&temp)
.await
.map_err(|error| storage_error("create evaluation result generation", error))?;
file.write_all(&bytes)
.await
.map_err(|error| storage_error("write evaluation result generation", error))?;
file.sync_all()
.await
.map_err(|error| storage_error("sync evaluation result generation", error))?;
drop(file);
let temp_copy = temp.clone();
let path_copy = path.clone();
tokio::task::spawn_blocking(move || {
tempfile::TempPath::try_from_path(temp_copy)
.map_err(|error| storage_error("prepare evaluation result replacement", error))?
.persist(path_copy)
.map_err(|error| {
storage_error("publish evaluation result generation", error.error)
})
})
.await
.map_err(|error| {
EvaluationStoreError::Storage(format!("publish task failed: {error}"))
})??;
Ok::<(), EvaluationStoreError>(())
}
.await;
if result.is_err() {
let _ = tokio::fs::remove_file(&temp).await;
}
result
}
}
#[async_trait]
impl EvaluationResultSink for FileEvaluationResultStore {
async fn write(
&self,
record: EvaluationRecordV1,
) -> Result<EvaluationWriteOutcomeV1, EvaluationStoreError> {
self.mutate(record).await
}
async fn get(&self, record_digest: &str) -> Option<EvaluationRecordV1> {
self.get_checked(record_digest).await.ok().flatten()
}
async fn list_for_target(&self, target: &ExecutionTargetV1) -> Vec<EvaluationRecordV1> {
self.list_for_target_checked(target)
.await
.unwrap_or_default()
}
async fn get_checked(
&self,
record_digest: &str,
) -> Result<Option<EvaluationRecordV1>, EvaluationStoreError> {
FileEvaluationResultStore::get_checked(self, record_digest).await
}
async fn list_for_target_checked(
&self,
target: &ExecutionTargetV1,
) -> Result<Vec<EvaluationRecordV1>, EvaluationStoreError> {
FileEvaluationResultStore::list_for_target_checked(self, target).await
}
}
async fn read_records_from_path(
path: &Path,
max_records: usize,
) -> Result<Vec<EvaluationRecordV1>, EvaluationStoreError> {
let bytes =
match crate::bounded_io::read_file_bounded_async(path, EVALUATION_RESULT_STORE_MAX_BYTES)
.await
{
Ok(bytes) => bytes,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(error) if error.kind() == std::io::ErrorKind::InvalidData => {
return Err(EvaluationStoreError::SizeLimit)
}
Err(error) => return Err(storage_error("read evaluation result store", error)),
};
let file: EvaluationResultStoreFileV1 = serde_json::from_slice(&bytes)
.map_err(|error| EvaluationStoreError::Corrupt(format!("decode generation: {error}")))?;
if file.schema != EVALUATION_RESULT_STORE_SCHEMA_V1 {
return Err(EvaluationStoreError::UnsupportedSchema);
}
if file.records.is_empty() {
return Ok(file.records);
}
if file.records.len() > max_records {
return Err(EvaluationStoreError::Corrupt(format!(
"{} records exceed configured retention limit {}",
file.records.len(),
max_records
)));
}
let mut digests = HashSet::with_capacity(file.records.len());
let mut identities = HashSet::with_capacity(file.records.len());
for record in &file.records {
record.validate()?;
if !digests.insert(record.record_digest.clone()) {
return Err(EvaluationStoreError::Corrupt(
"duplicate record digest".to_string(),
));
}
let identity = (
record.result.target.clone(),
record.result.evaluator_id.clone(),
record.result.auxiliary_run_id.clone(),
);
if !identities.insert(identity) {
return Err(EvaluationStoreError::Corrupt(
"duplicate evaluator identity".to_string(),
));
}
}
Ok(file.records)
}
fn storage_error(operation: &str, error: impl std::fmt::Display) -> EvaluationStoreError {
EvaluationStoreError::Storage(format!("{operation}: {error}"))
}