use std::time::Duration;
use serde::{Deserialize, Serialize};
use tracing::{debug, info};
use chronix_core::ShardId;
use crate::objstore::backend::ObjectStoreBackend;
use crate::objstore::error::{ObjStoreError, Result};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TieringConfig {
pub cold_after: Duration,
pub remote_url: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ArchiveObject {
measurement: String,
shard_id: ShardId,
object_name: String,
}
impl ArchiveObject {
pub fn new(
measurement: impl Into<String>,
shard_id: ShardId,
object_name: impl Into<String>,
) -> Result<Self> {
let measurement = measurement.into();
let object_name = object_name.into();
for (what, value) in [("measurement", &measurement), ("object name", &object_name)] {
if value.is_empty()
|| value.contains('/')
|| value.contains('\\')
|| value.contains("..")
|| value == "."
|| value.contains('=')
{
return Err(ObjStoreError::InvalidConfig {
detail: format!(
"invalid archive {what} {value:?}: must not be empty or contain \
path separators, '..' components or '='"
),
});
}
}
Ok(Self {
measurement,
shard_id,
object_name,
})
}
#[must_use]
pub fn measurement(&self) -> &str {
&self.measurement
}
#[must_use]
pub const fn shard_id(&self) -> ShardId {
self.shard_id
}
#[must_use]
pub fn key(&self) -> String {
format!(
"measurement={}/shard={}/{}",
self.measurement, self.shard_id.0, self.object_name
)
}
#[must_use]
pub fn measurement_prefix(measurement: &str) -> String {
format!("measurement={measurement}/")
}
}
impl std::fmt::Display for ArchiveObject {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.key())
}
}
pub struct TieringEngine {
config: TieringConfig,
remote: ObjectStoreBackend,
}
impl TieringEngine {
pub const fn new(config: TieringConfig, remote: ObjectStoreBackend) -> Self {
Self { config, remote }
}
#[must_use]
pub const fn config(&self) -> &TieringConfig {
&self.config
}
#[must_use]
pub fn is_eligible(&self, max_timestamp_ns: i64, now_ns: i64) -> bool {
let cold_after_ns = i64::try_from(self.config.cold_after.as_nanos()).unwrap_or(i64::MAX);
let threshold = now_ns.saturating_sub(cold_after_ns);
max_timestamp_ns < threshold
}
pub async fn upload_archive(&self, object: &ArchiveObject, data: &[u8]) -> Result<String> {
let key = object.key();
info!(
measurement = %object.measurement,
shard = object.shard_id.0,
object = %key,
size = data.len(),
"uploading cold archive object"
);
self.remote
.put(&key, bytes::Bytes::copy_from_slice(data))
.await?;
metrics::counter!("chronix_objstore_tiering_objects_total").increment(1);
metrics::counter!("chronix_objstore_tiering_bytes_total").increment(data.len() as u64);
match self.remote.object_size(&key).await? {
Some(remote_size) if remote_size == data.len() => {}
found => {
return Err(ObjStoreError::IntegrityCheckFailed {
object: key,
expected: data.len(),
found,
});
}
}
debug!(object = %key, "cold archive object verified");
Ok(key)
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use super::*;
use std::sync::Arc;
use crate::objstore::backend::ObjectStoreBackend;
fn engine() -> TieringEngine {
let remote =
ObjectStoreBackend::from_store(Arc::new(object_store::memory::InMemory::new()), None);
TieringEngine::new(
TieringConfig {
cold_after: Duration::from_secs(3600),
remote_url: "memory://test".to_string(),
},
remote,
)
}
#[test]
fn is_eligible_works() {
let engine = engine();
let now_ns: i64 = 10_000_000_000_000;
assert!(engine.is_eligible(5_000_000_000_000, now_ns));
assert!(!engine.is_eligible(7_000_000_000_000, now_ns));
assert!(!engine.is_eligible(6_400_000_000_000, now_ns));
}
#[test]
fn an_archive_object_names_its_measurement_and_shard() {
let object = ArchiveObject::new("power", ShardId(7), "part-a.parquet").unwrap();
assert_eq!(object.key(), "measurement=power/shard=7/part-a.parquet");
assert_eq!(object.measurement(), "power");
assert_eq!(object.shard_id(), ShardId(7));
assert_eq!(
ArchiveObject::measurement_prefix("power"),
"measurement=power/"
);
}
#[test]
fn an_archive_object_refuses_a_traversing_name() {
for bad in ["", "..", "a/b", "a\\b", "../etc", "shard=1"] {
assert!(
ArchiveObject::new(bad, ShardId(0), "part.parquet").is_err(),
"measurement {bad:?} must be refused"
);
assert!(
ArchiveObject::new("power", ShardId(0), bad).is_err(),
"object name {bad:?} must be refused"
);
}
}
#[tokio::test]
async fn upload_archive_verifies_the_object_landed() {
let engine = engine();
let object = ArchiveObject::new("power", ShardId(1), "part-a.parquet").unwrap();
let key = engine
.upload_archive(&object, b"parquet-bytes")
.await
.unwrap();
assert_eq!(key, "measurement=power/shard=1/part-a.parquet");
let back = engine.remote.get(&key).await.unwrap();
assert_eq!(&back[..], b"parquet-bytes");
assert_eq!(engine.remote.object_size(&key).await.unwrap(), Some(13));
}
#[tokio::test]
async fn a_missing_object_fails_verification() {
let engine = engine();
assert_eq!(
engine
.remote
.object_size("measurement=power/shard=1/absent.parquet")
.await
.unwrap(),
None,
"nothing was uploaded, so the HEAD must report absence rather than erroring"
);
}
#[test]
fn tiering_config_serde_roundtrip() {
let config = TieringConfig {
cold_after: Duration::from_secs(7200),
remote_url: "s3://my-bucket/cold".to_string(),
};
let json = serde_json::to_string(&config).unwrap();
let deserialized: TieringConfig = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.cold_after, config.cold_after);
assert_eq!(deserialized.remote_url, config.remote_url);
}
}