use async_trait::async_trait;
use bytes::Bytes;
use futures::stream::{self, StreamExt};
use super::nar_refs::{referrer_of, NarRefIndex, NarRefKey, NarRefScan};
use super::nar_stream::{self, NarSource, NarStream, NAR_CHUNK_BYTES};
use super::{NarResidency, StorageBackend};
use crate::StoreError;
const CHUNK_MARKER_SEQ: i32 = -1;
fn encode_marker(chunks: u64) -> [u8; 8] {
chunks.to_le_bytes()
}
fn decode_marker(raw: &[u8]) -> Result<u64, StoreError> {
<[u8; 8]>::try_from(raw)
.map(u64::from_le_bytes)
.map_err(|_| StoreError::NarInfo(format!("corrupt NAR chunk marker: {} bytes", raw.len())))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PgTable {
Narinfo,
Nar,
NarRef,
}
impl PgTable {
#[must_use]
pub const fn table_name(self) -> &'static str {
match self {
PgTable::Narinfo => "sui_cache_narinfo",
PgTable::Nar => "sui_cache_nar",
PgTable::NarRef => "sui_cache_nar_ref",
}
}
}
#[async_trait]
pub trait PgCacheConn: Send + Sync {
async fn select(&self, table: PgTable, key: &str) -> Result<Option<Vec<u8>>, StoreError>;
async fn upsert(&self, table: PgTable, key: &str, value: &[u8]) -> Result<(), StoreError>;
async fn delete(&self, table: PgTable, key: &str) -> Result<(), StoreError>;
async fn keys(&self, table: PgTable) -> Result<Vec<String>, StoreError>;
async fn keys_with_prefix(
&self,
table: PgTable,
prefix: &str,
) -> Result<Vec<String>, StoreError>;
async fn clear(&self, table: PgTable) -> Result<u64, StoreError>;
async fn upsert_nar_chunk(&self, key: &str, seq: i32, value: &[u8]) -> Result<(), StoreError>;
async fn select_nar_chunk(&self, key: &str, seq: i32) -> Result<Option<Vec<u8>>, StoreError>;
async fn delete_nar_chunks(&self, key: &str) -> Result<(), StoreError>;
async fn clear_nar_chunks(&self) -> Result<u64, StoreError>;
async fn select_nar_window(
&self,
key: &str,
offset: i64,
len: i32,
) -> Result<Option<(Vec<u8>, i64)>, StoreError>;
async fn ensure_schema(&self) -> Result<(), StoreError> {
Ok(())
}
}
pub struct PgStorageBackend<C: PgCacheConn> {
conn: std::sync::Arc<C>,
}
impl<C: PgCacheConn> PgStorageBackend<C> {
pub fn new(conn: C) -> Self {
Self { conn: std::sync::Arc::new(conn) }
}
pub fn conn(&self) -> &C {
&self.conn
}
async fn healing<T, F, Fut>(&self, op: F) -> Result<T, StoreError>
where
F: Fn() -> Fut,
Fut: std::future::Future<Output = Result<T, StoreError>>,
{
match op().await {
Err(StoreError::SchemaMissing(detail)) => {
tracing::warn!(
detail = %detail,
"pg L2: schema absent — re-running idempotent DDL and retrying once \
(a durable tier came back on an empty volume?)",
);
self.conn.ensure_schema().await?;
op().await
}
other => other,
}
}
}
impl<C: PgCacheConn + 'static> PgStorageBackend<C> {
fn chunked_stream(&self, path: &str, chunks: u64) -> NarStream {
let conn = std::sync::Arc::clone(&self.conn);
let key = path.to_string();
stream::unfold((conn, key, 0u64), move |(conn, key, seq)| async move {
if seq >= chunks {
return None;
}
match conn.select_nar_chunk(&key, seq as i32).await {
Ok(Some(v)) => Some((Ok(Bytes::from(v)), (conn, key, seq + 1))),
Ok(None) => {
let e = StoreError::NarInfo(format!(
"NAR {key}: chunk {seq} of {chunks} is missing though the \
completeness marker claims a whole value",
));
Some((Err(e), (conn, key, chunks)))
}
Err(e) => Some((Err(e), (conn, key, chunks))),
}
})
.boxed()
}
fn legacy_window_stream(&self, path: &str, first: Vec<u8>, total: i64) -> NarStream {
let conn = std::sync::Arc::clone(&self.conn);
let key = path.to_string();
let next_offset = 1 + first.len() as i64;
let head = stream::once(async move { Ok(Bytes::from(first)) });
let tail = stream::unfold((conn, key, next_offset), move |(conn, key, off)| async move {
if off > total {
return None;
}
match conn.select_nar_window(&key, off, chunk_len()).await {
Ok(Some((v, _))) if !v.is_empty() => {
let n = v.len() as i64;
Some((Ok(Bytes::from(v)), (conn, key, off + n)))
}
Ok(_) => None,
Err(e) => Some((Err(e), (conn, key, total + 1))),
}
});
head.chain(tail).boxed()
}
}
#[async_trait]
impl<C: PgCacheConn + 'static> StorageBackend for PgStorageBackend<C> {
async fn get_narinfo(&self, hash: &str) -> Result<Option<String>, StoreError> {
match self.healing(|| self.conn.select(PgTable::Narinfo, hash)).await? {
Some(bytes) => {
let text = String::from_utf8(bytes).map_err(|e| {
StoreError::NarInfo(format!("invalid utf-8 in pg narinfo {hash}: {e}"))
})?;
Ok(Some(text))
}
None => Ok(None),
}
}
async fn put_narinfo_record(&self, hash: &str, content: &str) -> Result<(), StoreError> {
self.healing(|| self.conn.upsert(PgTable::Narinfo, hash, content.as_bytes())).await
}
async fn delete_narinfo_record(&self, hash: &str) -> Result<(), StoreError> {
self.healing(|| self.conn.delete(PgTable::Narinfo, hash)).await
}
async fn delete_nar_record(&self, nar_path: &str) -> Result<(), StoreError> {
self.healing(|| self.conn.delete(PgTable::Nar, nar_path)).await?;
self.healing(|| self.conn.delete_nar_chunks(nar_path)).await
}
fn nar_ref_index(&self) -> &dyn NarRefIndex {
self
}
async fn get_nar(&self, path: &str) -> Result<Option<Vec<u8>>, StoreError> {
match self.get_nar_stream(path).await? {
Some(s) => Ok(Some(nar_stream::collect_nar(s, None).await?)),
None => Ok(None),
}
}
async fn put_nar(&self, path: &str, data: &[u8]) -> Result<(), StoreError> {
self.put_nar_stream(path, &nar_stream::BytesNarSource::from(data)).await
}
fn nar_residency(&self) -> NarResidency {
NarResidency::Streaming
}
async fn get_nar_stream(&self, path: &str) -> Result<Option<NarStream>, StoreError> {
if let Some(raw) = self
.healing(|| self.conn.select_nar_chunk(path, CHUNK_MARKER_SEQ))
.await?
{
let chunks = decode_marker(&raw)?;
return Ok(Some(self.chunked_stream(path, chunks)));
}
match self.healing(|| self.conn.select_nar_window(path, 1, chunk_len())).await? {
Some((first, total)) => Ok(Some(self.legacy_window_stream(path, first, total))),
None => Ok(None),
}
}
async fn put_nar_stream(&self, path: &str, src: &dyn NarSource) -> Result<(), StoreError> {
self.healing(|| self.conn.delete_nar_chunks(path)).await?;
self.healing(|| self.conn.delete(PgTable::Nar, path)).await?;
let mut stream = src.open().await?;
let mut seq: i32 = 0;
while let Some(chunk) = stream.next().await {
let chunk: Bytes = chunk?;
self.healing(|| self.conn.upsert_nar_chunk(path, seq, &chunk)).await?;
seq = seq.checked_add(1).ok_or_else(|| {
StoreError::NarInfo(format!("NAR {path} exceeds the addressable chunk count"))
})?;
}
let marker = encode_marker(seq as u64);
self.healing(|| self.conn.upsert_nar_chunk(path, CHUNK_MARKER_SEQ, &marker)).await
}
async fn list_narinfos(&self) -> Result<Vec<String>, StoreError> {
self.healing(|| self.conn.keys(PgTable::Narinfo)).await
}
async fn wipe_all(&self) -> Result<usize, StoreError> {
let narinfos = self.healing(|| self.conn.clear(PgTable::Narinfo)).await? as usize;
self.healing(|| self.conn.clear(PgTable::Nar)).await?;
self.healing(|| self.conn.clear_nar_chunks()).await?;
self.healing(|| self.conn.clear(PgTable::NarRef)).await?;
Ok(narinfos)
}
}
#[async_trait]
impl<C: PgCacheConn + 'static> NarRefIndex for PgStorageBackend<C> {
async fn record(&self, nar_path: &str, hash: &str) -> Result<(), StoreError> {
let key = NarRefKey { nar_path, hash }.to_string();
self.healing(|| self.conn.upsert(PgTable::NarRef, &key, b"")).await
}
async fn forget(&self, nar_path: &str, hash: &str) -> Result<(), StoreError> {
let key = NarRefKey { nar_path, hash }.to_string();
self.healing(|| self.conn.delete(PgTable::NarRef, &key)).await
}
async fn referrers(&self, nar_path: &str) -> Result<Vec<String>, StoreError> {
let scan = NarRefScan { nar_path };
let prefix = scan.to_string();
let keys = self
.healing(|| self.conn.keys_with_prefix(PgTable::NarRef, &prefix))
.await?;
let mut hashes: Vec<String> = keys
.iter()
.filter_map(|k| referrer_of(&scan, k))
.map(str::to_string)
.collect();
hashes.sort();
hashes.dedup();
Ok(hashes)
}
}
fn chunk_len() -> i32 {
i32::try_from(NAR_CHUNK_BYTES).unwrap_or(i32::MAX)
}
#[cfg(feature = "postgres")]
mod sqlx_conn {
use super::{StoreError, PgCacheConn, PgStorageBackend, PgTable};
use async_trait::async_trait;
use sqlx::postgres::{PgPool, PgPoolOptions};
use sqlx::Row;
const UNDEFINED_TABLE: &str = "42P01";
fn to_store_err(e: sqlx::Error) -> StoreError {
if let sqlx::Error::Database(db) = &e {
if db.code().as_deref() == Some(UNDEFINED_TABLE) {
return StoreError::SchemaMissing(format!("postgres: {e}"));
}
}
StoreError::Io(std::io::Error::other(format!("postgres: {e}")))
}
const NAR_CHUNK_DDL: &str = "CREATE TABLE IF NOT EXISTS sui_cache_nar_chunk (\
key TEXT NOT NULL, seq INTEGER NOT NULL, value BYTEA NOT NULL, \
PRIMARY KEY (key, seq))";
const NAR_CHUNK_SELECT: &str =
"SELECT value FROM sui_cache_nar_chunk WHERE key = $1 AND seq = $2";
const NAR_CHUNK_UPSERT: &str = "INSERT INTO sui_cache_nar_chunk (key, seq, value) \
VALUES ($1, $2, $3) ON CONFLICT (key, seq) DO UPDATE SET value = EXCLUDED.value";
const NAR_CHUNK_DELETE_KEY: &str = "DELETE FROM sui_cache_nar_chunk WHERE key = $1";
const NAR_CHUNK_CLEAR: &str = "DELETE FROM sui_cache_nar_chunk";
const NAR_LEGACY_WINDOW: &str = "SELECT substr(value, $2, $3) AS chunk, \
octet_length(value) AS total FROM sui_cache_nar WHERE key = $1";
impl PgTable {
const fn ddl(self) -> &'static str {
match self {
PgTable::Narinfo => {
"CREATE TABLE IF NOT EXISTS sui_cache_narinfo (key TEXT PRIMARY KEY, value BYTEA NOT NULL)"
}
PgTable::Nar => {
"CREATE TABLE IF NOT EXISTS sui_cache_nar (key TEXT PRIMARY KEY, value BYTEA NOT NULL)"
}
PgTable::NarRef => {
"CREATE TABLE IF NOT EXISTS sui_cache_nar_ref (key TEXT PRIMARY KEY, value BYTEA NOT NULL)"
}
}
}
const fn select_sql(self) -> &'static str {
match self {
PgTable::Narinfo => "SELECT value FROM sui_cache_narinfo WHERE key = $1",
PgTable::Nar => "SELECT value FROM sui_cache_nar WHERE key = $1",
PgTable::NarRef => "SELECT value FROM sui_cache_nar_ref WHERE key = $1",
}
}
const fn upsert_sql(self) -> &'static str {
match self {
PgTable::Narinfo => {
"INSERT INTO sui_cache_narinfo (key, value) VALUES ($1, $2) \
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value"
}
PgTable::Nar => {
"INSERT INTO sui_cache_nar (key, value) VALUES ($1, $2) \
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value"
}
PgTable::NarRef => {
"INSERT INTO sui_cache_nar_ref (key, value) VALUES ($1, $2) \
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value"
}
}
}
const fn delete_sql(self) -> &'static str {
match self {
PgTable::Narinfo => "DELETE FROM sui_cache_narinfo WHERE key = $1",
PgTable::Nar => "DELETE FROM sui_cache_nar WHERE key = $1",
PgTable::NarRef => "DELETE FROM sui_cache_nar_ref WHERE key = $1",
}
}
const fn clear_sql(self) -> &'static str {
match self {
PgTable::Narinfo => "DELETE FROM sui_cache_narinfo",
PgTable::Nar => "DELETE FROM sui_cache_nar",
PgTable::NarRef => "DELETE FROM sui_cache_nar_ref",
}
}
const fn keys_sql(self) -> &'static str {
match self {
PgTable::Narinfo => "SELECT key FROM sui_cache_narinfo",
PgTable::Nar => "SELECT key FROM sui_cache_nar",
PgTable::NarRef => "SELECT key FROM sui_cache_nar_ref",
}
}
const fn keys_with_prefix_sql(self) -> &'static str {
match self {
PgTable::Narinfo => {
"SELECT key FROM sui_cache_narinfo WHERE starts_with(key, $1)"
}
PgTable::Nar => "SELECT key FROM sui_cache_nar WHERE starts_with(key, $1)",
PgTable::NarRef => {
"SELECT key FROM sui_cache_nar_ref WHERE starts_with(key, $1)"
}
}
}
}
pub struct SqlxPgCacheConn {
pool: PgPool,
}
impl SqlxPgCacheConn {
pub async fn connect(url: &str, max_conns: u32) -> Result<Self, StoreError> {
let pool = PgPoolOptions::new()
.max_connections(max_conns)
.after_connect(|conn, _meta| {
Box::pin(async move {
for t in [PgTable::Narinfo, PgTable::Nar, PgTable::NarRef] {
sqlx::query(t.ddl()).execute(&mut *conn).await?;
}
sqlx::query(NAR_CHUNK_DDL).execute(&mut *conn).await?;
Ok(())
})
})
.connect(url)
.await
.map_err(to_store_err)?;
let this = Self { pool };
this.create_tables().await?;
Ok(this)
}
async fn create_tables(&self) -> Result<(), StoreError> {
for t in [PgTable::Narinfo, PgTable::Nar, PgTable::NarRef] {
sqlx::query(t.ddl()).execute(&self.pool).await.map_err(to_store_err)?;
}
sqlx::query(NAR_CHUNK_DDL).execute(&self.pool).await.map_err(to_store_err)?;
Ok(())
}
}
#[async_trait]
impl PgCacheConn for SqlxPgCacheConn {
async fn select(&self, table: PgTable, key: &str) -> Result<Option<Vec<u8>>, StoreError> {
let row = sqlx::query(table.select_sql())
.bind(key)
.fetch_optional(&self.pool)
.await
.map_err(to_store_err)?;
match row {
Some(r) => {
let v: Vec<u8> = r.try_get("value").map_err(to_store_err)?;
Ok(Some(v))
}
None => Ok(None),
}
}
async fn upsert(&self, table: PgTable, key: &str, value: &[u8]) -> Result<(), StoreError> {
sqlx::query(table.upsert_sql())
.bind(key)
.bind(value)
.execute(&self.pool)
.await
.map_err(to_store_err)?;
Ok(())
}
async fn delete(&self, table: PgTable, key: &str) -> Result<(), StoreError> {
sqlx::query(table.delete_sql())
.bind(key)
.execute(&self.pool)
.await
.map_err(to_store_err)?;
Ok(())
}
async fn keys(&self, table: PgTable) -> Result<Vec<String>, StoreError> {
let rows = sqlx::query(table.keys_sql())
.fetch_all(&self.pool)
.await
.map_err(to_store_err)?;
rows.into_iter()
.map(|r| r.try_get::<String, _>("key").map_err(to_store_err))
.collect()
}
async fn keys_with_prefix(
&self,
table: PgTable,
prefix: &str,
) -> Result<Vec<String>, StoreError> {
let rows = sqlx::query(table.keys_with_prefix_sql())
.bind(prefix)
.fetch_all(&self.pool)
.await
.map_err(to_store_err)?;
rows.into_iter()
.map(|r| r.try_get::<String, _>("key").map_err(to_store_err))
.collect()
}
async fn clear(&self, table: PgTable) -> Result<u64, StoreError> {
let res = sqlx::query(table.clear_sql())
.execute(&self.pool)
.await
.map_err(to_store_err)?;
Ok(res.rows_affected())
}
async fn upsert_nar_chunk(
&self,
key: &str,
seq: i32,
value: &[u8],
) -> Result<(), StoreError> {
sqlx::query(NAR_CHUNK_UPSERT)
.bind(key)
.bind(seq)
.bind(value)
.execute(&self.pool)
.await
.map_err(to_store_err)?;
Ok(())
}
async fn select_nar_chunk(
&self,
key: &str,
seq: i32,
) -> Result<Option<Vec<u8>>, StoreError> {
let row = sqlx::query(NAR_CHUNK_SELECT)
.bind(key)
.bind(seq)
.fetch_optional(&self.pool)
.await
.map_err(to_store_err)?;
match row {
Some(r) => Ok(Some(r.try_get::<Vec<u8>, _>("value").map_err(to_store_err)?)),
None => Ok(None),
}
}
async fn delete_nar_chunks(&self, key: &str) -> Result<(), StoreError> {
sqlx::query(NAR_CHUNK_DELETE_KEY)
.bind(key)
.execute(&self.pool)
.await
.map_err(to_store_err)?;
Ok(())
}
async fn clear_nar_chunks(&self) -> Result<u64, StoreError> {
let res = sqlx::query(NAR_CHUNK_CLEAR)
.execute(&self.pool)
.await
.map_err(to_store_err)?;
Ok(res.rows_affected())
}
async fn select_nar_window(
&self,
key: &str,
offset: i64,
len: i32,
) -> Result<Option<(Vec<u8>, i64)>, StoreError> {
let row = sqlx::query(NAR_LEGACY_WINDOW)
.bind(key)
.bind(offset)
.bind(len)
.fetch_optional(&self.pool)
.await
.map_err(to_store_err)?;
match row {
Some(r) => {
let chunk: Vec<u8> = r.try_get("chunk").map_err(to_store_err)?;
let total: i32 = r.try_get("total").map_err(to_store_err)?;
Ok(Some((chunk, i64::from(total))))
}
None => Ok(None),
}
}
async fn ensure_schema(&self) -> Result<(), StoreError> {
self.create_tables().await
}
}
impl PgStorageBackend<SqlxPgCacheConn> {
pub async fn connect(url: &str, max_conns: u32) -> Result<Self, StoreError> {
Ok(Self::new(SqlxPgCacheConn::connect(url, max_conns).await?))
}
}
}
#[cfg(feature = "postgres")]
pub use sqlx_conn::SqlxPgCacheConn;
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
use std::sync::Mutex;
#[derive(Default)]
struct MockPg {
narinfo: Mutex<HashMap<String, Vec<u8>>>,
nar: Mutex<HashMap<String, Vec<u8>>>,
nar_ref: Mutex<HashMap<String, Vec<u8>>>,
nar_chunk: Mutex<HashMap<(String, i32), Vec<u8>>>,
schema_missing: Mutex<bool>,
ensure_schema_calls: Mutex<usize>,
}
impl MockPg {
fn table(&self, t: PgTable) -> &Mutex<HashMap<String, Vec<u8>>> {
match t {
PgTable::Narinfo => &self.narinfo,
PgTable::Nar => &self.nar,
PgTable::NarRef => &self.nar_ref,
}
}
fn drop_schema(&self) {
*self.schema_missing.lock().unwrap() = true;
self.narinfo.lock().unwrap().clear();
self.nar.lock().unwrap().clear();
self.nar_ref.lock().unwrap().clear();
self.nar_chunk.lock().unwrap().clear();
}
fn chunk_rows(&self) -> usize {
self.nar_chunk.lock().unwrap().len()
}
fn seed_legacy_nar(&self, key: &str, value: &[u8]) {
self.nar.lock().unwrap().insert(key.to_string(), value.to_vec());
}
fn ddl_runs(&self) -> usize {
*self.ensure_schema_calls.lock().unwrap()
}
fn guard(&self) -> Result<(), StoreError> {
if *self.schema_missing.lock().unwrap() {
Err(StoreError::SchemaMissing(
"postgres: relation \"sui_cache_narinfo\" does not exist".to_string(),
))
} else {
Ok(())
}
}
}
#[async_trait]
impl PgCacheConn for MockPg {
async fn select(&self, table: PgTable, key: &str) -> Result<Option<Vec<u8>>, StoreError> {
self.guard()?;
Ok(self.table(table).lock().unwrap().get(key).cloned())
}
async fn upsert(&self, table: PgTable, key: &str, value: &[u8]) -> Result<(), StoreError> {
self.guard()?;
self.table(table).lock().unwrap().insert(key.to_string(), value.to_vec());
Ok(())
}
async fn delete(&self, table: PgTable, key: &str) -> Result<(), StoreError> {
self.guard()?;
self.table(table).lock().unwrap().remove(key);
Ok(())
}
async fn keys(&self, table: PgTable) -> Result<Vec<String>, StoreError> {
self.guard()?;
Ok(self.table(table).lock().unwrap().keys().cloned().collect())
}
async fn keys_with_prefix(
&self,
table: PgTable,
prefix: &str,
) -> Result<Vec<String>, StoreError> {
self.guard()?;
Ok(self
.table(table)
.lock()
.unwrap()
.keys()
.filter(|k| k.starts_with(prefix))
.cloned()
.collect())
}
async fn clear(&self, table: PgTable) -> Result<u64, StoreError> {
self.guard()?;
let mut m = self.table(table).lock().unwrap();
let n = m.len() as u64;
m.clear();
Ok(n)
}
async fn upsert_nar_chunk(&self, key: &str, seq: i32, value: &[u8]) -> Result<(), StoreError> {
self.guard()?;
self.nar_chunk
.lock()
.unwrap()
.insert((key.to_string(), seq), value.to_vec());
Ok(())
}
async fn select_nar_chunk(&self, key: &str, seq: i32) -> Result<Option<Vec<u8>>, StoreError> {
self.guard()?;
Ok(self.nar_chunk.lock().unwrap().get(&(key.to_string(), seq)).cloned())
}
async fn delete_nar_chunks(&self, key: &str) -> Result<(), StoreError> {
self.guard()?;
self.nar_chunk.lock().unwrap().retain(|(k, _), _| k != key);
Ok(())
}
async fn clear_nar_chunks(&self) -> Result<u64, StoreError> {
self.guard()?;
let mut m = self.nar_chunk.lock().unwrap();
let n = m.len() as u64;
m.clear();
Ok(n)
}
async fn select_nar_window(
&self,
key: &str,
offset: i64,
len: i32,
) -> Result<Option<(Vec<u8>, i64)>, StoreError> {
self.guard()?;
let map = self.nar.lock().unwrap();
let Some(v) = map.get(key) else { return Ok(None) };
let total = v.len() as i64;
let start = (offset - 1).clamp(0, total) as usize;
let end = (start + len.max(0) as usize).min(v.len());
Ok(Some((v[start..end].to_vec(), total)))
}
async fn ensure_schema(&self) -> Result<(), StoreError> {
*self.ensure_schema_calls.lock().unwrap() += 1;
*self.schema_missing.lock().unwrap() = false;
Ok(())
}
}
#[derive(Default)]
struct UnhealablePg {
attempts: Mutex<usize>,
}
#[async_trait]
impl PgCacheConn for UnhealablePg {
async fn select(&self, _t: PgTable, _k: &str) -> Result<Option<Vec<u8>>, StoreError> {
*self.attempts.lock().unwrap() += 1;
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn upsert(&self, _t: PgTable, _k: &str, _v: &[u8]) -> Result<(), StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn delete(&self, _t: PgTable, _k: &str) -> Result<(), StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn keys(&self, _t: PgTable) -> Result<Vec<String>, StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn keys_with_prefix(
&self,
_t: PgTable,
_p: &str,
) -> Result<Vec<String>, StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn clear(&self, _t: PgTable) -> Result<u64, StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn upsert_nar_chunk(&self, _k: &str, _s: i32, _v: &[u8]) -> Result<(), StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn select_nar_chunk(&self, _k: &str, _s: i32) -> Result<Option<Vec<u8>>, StoreError> {
*self.attempts.lock().unwrap() += 1;
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn delete_nar_chunks(&self, _k: &str) -> Result<(), StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn clear_nar_chunks(&self) -> Result<u64, StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
async fn select_nar_window(
&self,
_k: &str,
_o: i64,
_l: i32,
) -> Result<Option<(Vec<u8>, i64)>, StoreError> {
Err(StoreError::SchemaMissing("still gone".to_string()))
}
}
const NARINFO: &str = "StorePath: /nix/store/abc-hello\nURL: nar/abc.nar.xz\nCompression: xz\nNarHash: sha256:bbb\nNarSize: 200\nReferences: \n";
#[test]
fn table_names_are_distinct() {
assert_ne!(PgTable::Narinfo.table_name(), PgTable::Nar.table_name());
}
#[tokio::test]
async fn get_missing_narinfo_returns_none() {
let backend = PgStorageBackend::new(MockPg::default());
assert!(backend.get_narinfo("nope").await.unwrap().is_none());
}
#[tokio::test]
async fn put_then_get_narinfo_roundtrips() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("abc", NARINFO).await.unwrap();
assert_eq!(backend.get_narinfo("abc").await.unwrap().unwrap(), NARINFO);
}
#[tokio::test]
async fn put_then_get_nar_roundtrips() {
let backend = PgStorageBackend::new(MockPg::default());
let data = b"\x00\x01\x02 fake nar bytes";
backend.put_nar("nar/abc.nar.xz", data).await.unwrap();
assert_eq!(backend.get_nar("nar/abc.nar.xz").await.unwrap().unwrap(), data);
}
#[tokio::test]
async fn narinfo_and_nar_keyspaces_do_not_collide() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("dead", "the-narinfo").await.unwrap();
backend.put_nar("dead", b"the-nar").await.unwrap();
assert_eq!(backend.get_narinfo("dead").await.unwrap().unwrap(), "the-narinfo");
assert_eq!(backend.get_nar("dead").await.unwrap().unwrap(), b"the-nar");
}
fn narinfo_for(url: &str) -> String {
format!(
"StorePath: /nix/store/pkg\nURL: {url}\nCompression: xz\nFileHash: sha256:aaa\n\
FileSize: 100\nNarHash: sha256:bbb\nNarSize: 200\nReferences: \n"
)
}
#[tokio::test]
async fn delete_resolves_the_nar_from_the_narinfo_instead_of_guessing() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("storehash", &narinfo_for("nar/narhash.nar.xz")).await.unwrap();
backend.put_nar("nar/narhash.nar.xz", b"the real nar").await.unwrap();
backend.put_nar("nar/storehash.nar.zst", b"someone else's nar").await.unwrap();
backend.delete("storehash").await.unwrap();
assert!(backend.get_narinfo("storehash").await.unwrap().is_none());
assert!(
backend.get_nar("nar/narhash.nar.xz").await.unwrap().is_none(),
"the advertised NAR must be the one that goes",
);
assert_eq!(
backend.get_nar("nar/storehash.nar.zst").await.unwrap().unwrap(),
b"someone else's nar",
"a store-hash-shaped key this narinfo never named must be untouched",
);
}
#[tokio::test]
async fn deleting_one_of_two_paths_sharing_a_nar_leaves_the_nar() {
let backend = PgStorageBackend::new(MockPg::default());
let shared = "nar/sharednarhash.nar.xz";
backend.put_narinfo("pathA", &narinfo_for(shared)).await.unwrap();
backend.put_narinfo("pathB", &narinfo_for(shared)).await.unwrap();
backend.put_nar(shared, b"shared contents").await.unwrap();
backend.delete("pathA").await.unwrap();
assert!(backend.get_narinfo("pathA").await.unwrap().is_none());
assert!(
backend.get_nar(shared).await.unwrap().is_some(),
"pathB still advertises this NAR",
);
backend.delete("pathB").await.unwrap();
assert!(
backend.get_nar(shared).await.unwrap().is_none(),
"the last referrer gone means the NAR is reclaimable",
);
}
#[tokio::test]
async fn wipe_all_truncates_both_tables_incl_narhash_keyed_nar() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("storehash", NARINFO).await.unwrap();
backend.put_nar("nar/0narhashXXXXXXXXXXXXXXXXXXXXXXXXXX.nar", b"blob").await.unwrap();
backend.put_narinfo("other", NARINFO).await.unwrap();
let removed = backend.wipe_all().await.unwrap();
assert_eq!(removed, 2, "wipe_all should report the narinfo count");
assert!(backend.list_narinfos().await.unwrap().is_empty());
assert!(backend.get_narinfo("storehash").await.unwrap().is_none());
assert!(backend.get_narinfo("other").await.unwrap().is_none());
assert!(backend
.get_nar("nar/0narhashXXXXXXXXXXXXXXXXXXXXXXXXXX.nar")
.await
.unwrap()
.is_none());
}
#[tokio::test]
async fn delete_absent_is_idempotent() {
let backend = PgStorageBackend::new(MockPg::default());
backend.delete("ghost").await.unwrap();
}
#[tokio::test]
async fn list_narinfos_is_authoritative_and_full() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("aaa", "1").await.unwrap();
backend.put_narinfo("bbb", "2").await.unwrap();
backend.put_nar("nar/ccc.nar.xz", b"3").await.unwrap();
let mut hashes = backend.list_narinfos().await.unwrap();
hashes.sort();
assert_eq!(hashes, vec!["aaa".to_string(), "bbb".to_string()]);
}
#[tokio::test]
async fn overwrite_narinfo_takes_latest() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("h", "v1").await.unwrap();
backend.put_narinfo("h", "v2").await.unwrap();
assert_eq!(backend.get_narinfo("h").await.unwrap().unwrap(), "v2");
}
#[tokio::test]
async fn schema_vanishing_under_a_live_connection_self_heals_on_the_next_read() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("h", NARINFO).await.unwrap();
assert_eq!(backend.conn().ddl_runs(), 0, "no heal needed while healthy");
backend.conn().drop_schema();
let got = backend.get_narinfo("h").await.expect("must self-heal, not error");
assert!(got.is_none(), "the data really is gone — a miss, not a 500");
assert_eq!(backend.conn().ddl_runs(), 1, "the idempotent DDL must have re-run");
backend.put_narinfo("h2", NARINFO).await.unwrap();
assert_eq!(backend.get_narinfo("h2").await.unwrap().unwrap(), NARINFO);
}
#[tokio::test]
async fn ensure_schema_is_idempotent_run_twice_no_error() {
let conn = MockPg::default();
conn.ensure_schema().await.expect("first run");
conn.ensure_schema().await.expect("second run must be a harmless no-op");
conn.ensure_schema().await.expect("third run must be a harmless no-op");
assert_eq!(conn.ddl_runs(), 3);
conn.drop_schema();
conn.ensure_schema().await.expect("run against an absent schema");
conn.ensure_schema().await.expect("and again once it exists");
let backend = PgStorageBackend::new(conn);
backend.put_narinfo("x", NARINFO).await.expect("usable after repeated DDL");
}
#[tokio::test]
async fn every_verb_self_heals_not_just_reads() {
for_each_verb_self_heals().await;
}
async fn for_each_verb_self_heals() {
let b = PgStorageBackend::new(MockPg::default());
b.conn().drop_schema();
b.put_narinfo("h", NARINFO).await.expect("put_narinfo self-heals");
assert_eq!(b.get_narinfo("h").await.unwrap().unwrap(), NARINFO);
let b = PgStorageBackend::new(MockPg::default());
b.conn().drop_schema();
b.put_nar("nar/h.nar.xz", b"blob").await.expect("put_nar self-heals");
assert_eq!(b.get_nar("nar/h.nar.xz").await.unwrap().unwrap(), b"blob");
let b = PgStorageBackend::new(MockPg::default());
b.conn().drop_schema();
assert!(b.list_narinfos().await.expect("list self-heals").is_empty());
let b = PgStorageBackend::new(MockPg::default());
b.conn().drop_schema();
b.delete("h").await.expect("delete self-heals");
let b = PgStorageBackend::new(MockPg::default());
b.conn().drop_schema();
assert_eq!(b.wipe_all().await.expect("wipe self-heals"), 0);
}
#[tokio::test]
async fn an_unrepairable_schema_surfaces_after_exactly_one_retry() {
let backend = PgStorageBackend::new(UnhealablePg::default());
let err = backend.get_narinfo("h").await.unwrap_err();
assert!(matches!(err, StoreError::SchemaMissing(_)));
assert_eq!(
*backend.conn().attempts.lock().unwrap(),
2,
"exactly one retry after the heal attempt — never a spin",
);
}
#[tokio::test]
async fn a_non_schema_error_is_never_retried_as_a_schema_problem() {
struct BrokenPg;
#[async_trait]
impl PgCacheConn for BrokenPg {
async fn select(&self, _t: PgTable, _k: &str) -> Result<Option<Vec<u8>>, StoreError> {
Err(StoreError::Io(std::io::Error::other(
"postgres: expected to read 5 bytes, got 0 bytes at EOF",
)))
}
async fn upsert(&self, _t: PgTable, _k: &str, _v: &[u8]) -> Result<(), StoreError> {
unreachable!()
}
async fn delete(&self, _t: PgTable, _k: &str) -> Result<(), StoreError> {
unreachable!()
}
async fn keys(&self, _t: PgTable) -> Result<Vec<String>, StoreError> {
unreachable!()
}
async fn keys_with_prefix(
&self,
_t: PgTable,
_p: &str,
) -> Result<Vec<String>, StoreError> {
unreachable!()
}
async fn clear(&self, _t: PgTable) -> Result<u64, StoreError> {
unreachable!()
}
async fn upsert_nar_chunk(&self, _k: &str, _s: i32, _v: &[u8]) -> Result<(), StoreError> {
unreachable!()
}
async fn select_nar_chunk(&self, _k: &str, _s: i32) -> Result<Option<Vec<u8>>, StoreError> {
unreachable!()
}
async fn delete_nar_chunks(&self, _k: &str) -> Result<(), StoreError> {
unreachable!()
}
async fn clear_nar_chunks(&self) -> Result<u64, StoreError> {
unreachable!()
}
async fn select_nar_window(
&self,
_k: &str,
_o: i64,
_l: i32,
) -> Result<Option<(Vec<u8>, i64)>, StoreError> {
unreachable!()
}
async fn ensure_schema(&self) -> Result<(), StoreError> {
panic!("a non-schema error must never trigger the DDL path");
}
}
let backend = PgStorageBackend::new(BrokenPg);
assert!(matches!(backend.get_narinfo("h").await.unwrap_err(), StoreError::Io(_)));
}
#[tokio::test]
async fn invalid_utf8_narinfo_surfaces_typed_error() {
let mock = MockPg::default();
mock.narinfo.lock().unwrap().insert("bad".to_string(), vec![0xff, 0xfe, 0xfd]);
let backend = PgStorageBackend::new(mock);
let err = backend.get_narinfo("bad").await.unwrap_err();
assert!(matches!(err, StoreError::NarInfo(_)));
}
struct FaultyPg {
inner: MockPg,
ok_calls: Mutex<usize>,
fail_with: String,
}
impl FaultyPg {
fn always(msg: &str) -> Self {
Self {
inner: MockPg::default(),
ok_calls: Mutex::new(0),
fail_with: msg.to_string(),
}
}
fn after(n: usize, msg: &str) -> Self {
Self {
inner: MockPg::default(),
ok_calls: Mutex::new(n),
fail_with: msg.to_string(),
}
}
fn tick(&self) -> Result<(), StoreError> {
let mut left = self.ok_calls.lock().unwrap();
if *left == 0 {
return Err(StoreError::Io(std::io::Error::other(format!(
"postgres: {}",
self.fail_with
))));
}
*left -= 1;
Ok(())
}
}
const RELATION_MISSING: &str = "error returned from database: \
relation \"sui_cache_narinfo\" does not exist";
#[async_trait]
impl PgCacheConn for FaultyPg {
async fn select(&self, table: PgTable, key: &str) -> Result<Option<Vec<u8>>, StoreError> {
self.tick()?;
self.inner.select(table, key).await
}
async fn upsert(&self, table: PgTable, key: &str, value: &[u8]) -> Result<(), StoreError> {
self.tick()?;
self.inner.upsert(table, key, value).await
}
async fn delete(&self, table: PgTable, key: &str) -> Result<(), StoreError> {
self.tick()?;
self.inner.delete(table, key).await
}
async fn keys(&self, table: PgTable) -> Result<Vec<String>, StoreError> {
self.tick()?;
self.inner.keys(table).await
}
async fn keys_with_prefix(
&self,
table: PgTable,
prefix: &str,
) -> Result<Vec<String>, StoreError> {
self.tick()?;
self.inner.keys_with_prefix(table, prefix).await
}
async fn clear(&self, table: PgTable) -> Result<u64, StoreError> {
self.tick()?;
self.inner.clear(table).await
}
async fn upsert_nar_chunk(&self, key: &str, seq: i32, value: &[u8]) -> Result<(), StoreError> {
self.tick()?;
self.inner.upsert_nar_chunk(key, seq, value).await
}
async fn select_nar_chunk(&self, key: &str, seq: i32) -> Result<Option<Vec<u8>>, StoreError> {
self.tick()?;
self.inner.select_nar_chunk(key, seq).await
}
async fn delete_nar_chunks(&self, key: &str) -> Result<(), StoreError> {
self.tick()?;
self.inner.delete_nar_chunks(key).await
}
async fn clear_nar_chunks(&self) -> Result<u64, StoreError> {
self.tick()?;
self.inner.clear_nar_chunks().await
}
async fn select_nar_window(
&self,
key: &str,
offset: i64,
len: i32,
) -> Result<Option<(Vec<u8>, i64)>, StoreError> {
self.tick()?;
self.inner.select_nar_window(key, offset, len).await
}
}
#[tokio::test]
async fn backend_fault_is_an_error_never_a_silent_miss() {
let backend = PgStorageBackend::new(FaultyPg::always(RELATION_MISSING));
let err = backend.get_narinfo("abc").await.unwrap_err();
assert!(
err.to_string().contains("does not exist"),
"the underlying cause must survive to the caller, got: {err}"
);
let healthy = PgStorageBackend::new(MockPg::default());
assert!(healthy.get_narinfo("abc").await.unwrap().is_none());
}
#[tokio::test]
async fn backend_fault_on_write_is_an_error() {
let backend = PgStorageBackend::new(FaultyPg::always(RELATION_MISSING));
assert!(backend.put_narinfo("h", NARINFO).await.is_err());
assert!(backend.put_nar("nar/h.nar.xz", b"bytes").await.is_err());
}
#[tokio::test]
async fn every_read_path_propagates_a_backend_fault() {
let backend = PgStorageBackend::new(FaultyPg::always(RELATION_MISSING));
assert!(backend.get_narinfo("h").await.is_err(), "get_narinfo");
assert!(backend.get_nar("nar/h.nar.xz").await.is_err(), "get_nar");
assert!(backend.list_narinfos().await.is_err(), "list_narinfos");
assert!(backend.delete("h").await.is_err(), "delete");
}
#[tokio::test]
async fn concurrent_puts_do_not_lose_writes() {
use std::sync::Arc;
let backend = Arc::new(PgStorageBackend::new(MockPg::default()));
let mut set = tokio::task::JoinSet::new();
for i in 0..64 {
let b = Arc::clone(&backend);
set.spawn(async move { b.put_narinfo(&format!("k{i}"), &format!("v{i}")).await });
}
while let Some(r) = set.join_next().await {
r.expect("task panicked").expect("put failed");
}
assert_eq!(backend.list_narinfos().await.unwrap().len(), 64);
for i in 0..64 {
assert_eq!(
backend.get_narinfo(&format!("k{i}")).await.unwrap().unwrap(),
format!("v{i}"),
"key k{i} round-tripped wrong under concurrency"
);
}
}
#[tokio::test]
async fn repeated_identical_put_is_idempotent() {
let backend = PgStorageBackend::new(MockPg::default());
for _ in 0..10 {
backend.put_narinfo("same", NARINFO).await.unwrap();
}
assert_eq!(backend.list_narinfos().await.unwrap().len(), 1);
assert_eq!(backend.get_narinfo("same").await.unwrap().unwrap(), NARINFO);
}
const NAR_KEY: &str = "nar/deadbeef.nar.xz";
fn multi_chunk_nar() -> Vec<u8> {
(0..NAR_CHUNK_BYTES * 2 + 4096).map(|i| (i % 251) as u8).collect()
}
#[tokio::test]
async fn a_nar_is_stored_as_bounded_chunks_never_one_whole_row() {
let backend = PgStorageBackend::new(MockPg::default());
let nar = multi_chunk_nar();
backend.put_nar(NAR_KEY, &nar).await.unwrap();
assert_eq!(backend.conn().chunk_rows(), 4, "expected 3 chunks + a marker");
assert!(
backend.conn().nar.lock().unwrap().is_empty(),
"a streamed write must not also write the legacy whole-value row",
);
}
#[tokio::test]
async fn a_chunked_nar_round_trips_byte_identically() {
let backend = PgStorageBackend::new(MockPg::default());
let nar = multi_chunk_nar();
backend.put_nar(NAR_KEY, &nar).await.unwrap();
assert_eq!(backend.get_nar(NAR_KEY).await.unwrap().unwrap(), nar);
}
#[tokio::test]
async fn an_empty_nar_round_trips_as_empty_not_as_a_miss() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_nar(NAR_KEY, b"").await.unwrap();
assert_eq!(backend.get_nar(NAR_KEY).await.unwrap().unwrap(), Vec::<u8>::new());
}
#[tokio::test]
async fn a_write_killed_before_the_marker_reads_as_a_miss_not_a_truncated_nar() {
let backend = PgStorageBackend::new(MockPg::default());
backend.conn().upsert_nar_chunk(NAR_KEY, 0, b"first half").await.unwrap();
backend.conn().upsert_nar_chunk(NAR_KEY, 1, b"second half").await.unwrap();
assert!(
backend.get_nar(NAR_KEY).await.unwrap().is_none(),
"orphan chunks must not be servable",
);
}
#[tokio::test]
async fn a_gap_under_a_published_marker_is_an_error_never_a_short_read() {
let backend = PgStorageBackend::new(MockPg::default());
backend.conn().upsert_nar_chunk(NAR_KEY, 0, b"present").await.unwrap();
backend.conn().upsert_nar_chunk(NAR_KEY, CHUNK_MARKER_SEQ, &encode_marker(2)).await.unwrap();
let err = backend.get_nar(NAR_KEY).await.unwrap_err();
assert!(
err.to_string().contains("missing"),
"expected a typed corruption error, got: {err}",
);
}
#[tokio::test]
async fn a_corrupt_marker_surfaces_rather_than_being_coerced() {
let backend = PgStorageBackend::new(MockPg::default());
backend.conn().upsert_nar_chunk(NAR_KEY, CHUNK_MARKER_SEQ, b"nope").await.unwrap();
assert!(matches!(
backend.get_nar(NAR_KEY).await.unwrap_err(),
StoreError::NarInfo(_),
));
}
#[tokio::test]
async fn re_putting_a_shorter_nar_leaves_no_stale_tail() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_nar(NAR_KEY, &multi_chunk_nar()).await.unwrap();
backend.put_nar(NAR_KEY, b"short").await.unwrap();
assert_eq!(backend.get_nar(NAR_KEY).await.unwrap().unwrap(), b"short");
assert_eq!(backend.conn().chunk_rows(), 2, "1 chunk + a marker; the tail is gone");
}
#[tokio::test]
async fn a_legacy_whole_value_row_is_still_served_windowed() {
let backend = PgStorageBackend::new(MockPg::default());
let legacy = multi_chunk_nar();
backend.conn().seed_legacy_nar(NAR_KEY, &legacy);
assert_eq!(backend.conn().chunk_rows(), 0, "the fixture is pre-streaming by construction");
assert_eq!(backend.get_nar(NAR_KEY).await.unwrap().unwrap(), legacy);
}
#[tokio::test]
async fn a_chunked_write_shadows_a_legacy_row_for_the_same_key() {
let backend = PgStorageBackend::new(MockPg::default());
backend.conn().seed_legacy_nar(NAR_KEY, b"old whole-value bytes");
backend.put_nar(NAR_KEY, b"new chunked bytes").await.unwrap();
assert_eq!(backend.get_nar(NAR_KEY).await.unwrap().unwrap(), b"new chunked bytes");
assert!(
backend.conn().nar.lock().unwrap().is_empty(),
"the legacy row must be dropped, not shadowed",
);
}
#[tokio::test]
async fn delete_clears_both_storage_generations() {
let backend = PgStorageBackend::new(MockPg::default());
let key = "nar/xyz.nar.xz";
backend.put_narinfo("storehash", &narinfo_for(key)).await.unwrap();
backend.put_nar(key, b"chunked").await.unwrap();
backend.conn().seed_legacy_nar(key, b"legacy");
backend.delete("storehash").await.unwrap();
assert!(backend.get_nar(key).await.unwrap().is_none());
assert!(backend.conn().nar.lock().unwrap().is_empty(), "the legacy row must go too");
assert_eq!(backend.conn().chunk_rows(), 0);
}
#[tokio::test]
async fn wipe_all_reclaims_the_chunk_table_too() {
let backend = PgStorageBackend::new(MockPg::default());
backend.put_narinfo("h", NARINFO).await.unwrap();
backend.put_nar(NAR_KEY, &multi_chunk_nar()).await.unwrap();
assert!(backend.conn().chunk_rows() > 1);
assert_eq!(backend.wipe_all().await.unwrap(), 1);
assert_eq!(backend.conn().chunk_rows(), 0, "wipe must reach the chunk table");
assert!(backend.get_nar(NAR_KEY).await.unwrap().is_none());
}
#[tokio::test]
async fn the_chunked_write_path_self_heals_a_vanished_schema() {
let backend = PgStorageBackend::new(MockPg::default());
backend.conn().drop_schema();
backend.put_nar(NAR_KEY, b"bytes").await.expect("put_nar self-heals");
assert_eq!(backend.get_nar(NAR_KEY).await.unwrap().unwrap(), b"bytes");
}
#[tokio::test]
async fn a_backend_fault_partway_through_a_chunked_read_surfaces() {
let backend = PgStorageBackend::new(FaultyPg::after(6, RELATION_MISSING));
let err = backend.put_nar(NAR_KEY, &multi_chunk_nar()).await;
if err.is_ok() {
assert!(backend.get_nar(NAR_KEY).await.is_err(), "a mid-read fault must surface");
}
}
#[tokio::test]
async fn residency_is_streaming() {
let backend = PgStorageBackend::new(MockPg::default());
assert_eq!(backend.nar_residency(), NarResidency::Streaming);
}
}