newton-core 0.7.5

//! Storage backend configuration.
//!
//! Selects and constructs the [`UserStorageBackend`] used for content-addressed
//! objects: a local filesystem tree in development and S3 in production. It
//! also carries limits enforced by the authenticated storage endpoint.

use std::sync::Arc;

use eyre::Result;
use serde::{Deserialize, Serialize};

use super::{
    backend::{LocalBackend, UserStorageBackend},
    local_cache::LocalObjectCache,
    resolver::PersistedImmutableDataResolver,
};

/// Default cap on a single uploaded object (16 MiB).
const fn default_max_object_bytes() -> usize {
    16 * 1024 * 1024
}

/// Default cap on the number of objects one API key may own.
const fn default_max_objects_per_owner() -> usize {
    256
}

/// Default cap on the local object cache (20 GiB).
const fn default_local_cache_max_bytes() -> u64 {
    20 * 1024 * 1024 * 1024
}

/// Default local object cache directory, relative to the working directory like the
/// other `./data` defaults.
fn default_local_cache_dir() -> std::path::PathBuf {
    std::path::PathBuf::from("./data/object-cache")
}

/// Local disk tier for content-addressed reads: verified objects are kept under `dir`
/// and the backend is read only on a miss. Used only by
/// [`StorageConfig::build_with_local_cache`] (the operator, and so the unified binary),
/// and only in front of a remote (S3) backend: a local backend is already files on local
/// disk. Both fields default.
///
/// TOML: `[storage.local_cache]` in operator.toml. Env:
/// `OPERATOR__STORAGE__LOCAL_CACHE__DIR`, `OPERATOR__STORAGE__LOCAL_CACHE__MAX_BYTES`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default)]
pub struct LocalCacheConfig {
    /// Directory for cached objects (kept under `<dir>/objects/`). Must be writable only by
    /// the service user: an object is re-verified only on its first read after start-up.
    pub dir: std::path::PathBuf,
    /// Least-recently-used objects are deleted past this many bytes.
    pub max_bytes: u64,
}

impl Default for LocalCacheConfig {
    fn default() -> Self {
        Self {
            dir: default_local_cache_dir(),
            max_bytes: default_local_cache_max_bytes(),
        }
    }
}

/// Which storage backend to construct.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum StorageBackendConfig {
    /// Local filesystem tree rooted at `dir` (dev/test).
    Local {
        /// Root directory for stored objects.
        dir: std::path::PathBuf,
    },
    /// S3 bucket in `region` (prod). Requires the `persisted-immutable-data-s3` feature.
    S3 {
        /// Bucket name.
        bucket: String,
        /// AWS region.
        #[serde(default = "default_s3_region")]
        region: String,
    },
}

fn default_s3_region() -> String {
    "us-west-2".to_string()
}

impl Default for StorageBackendConfig {
    fn default() -> Self {
        Self::Local {
            dir: std::path::PathBuf::from("./data/persisted-immutable-data"),
        }
    }
}

/// Storage configuration: backend selection + endpoint limits.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default)]
pub struct StorageConfig {
    /// Backend to construct.
    pub backend: StorageBackendConfig,
    /// Maximum size of a single uploaded object, in bytes.
    #[serde(default = "default_max_object_bytes")]
    pub max_object_bytes: usize,
    /// Maximum number of objects a single owner (API key) may store.
    #[serde(default = "default_max_objects_per_owner")]
    pub max_objects_per_owner: usize,
    /// Local disk tier in front of the backend for content-addressed reads.
    pub local_cache: LocalCacheConfig,
}

impl Default for StorageConfig {
    fn default() -> Self {
        Self {
            backend: StorageBackendConfig::default(),
            max_object_bytes: default_max_object_bytes(),
            max_objects_per_owner: default_max_objects_per_owner(),
            local_cache: LocalCacheConfig::default(),
        }
    }
}

impl StorageConfig {
    /// Construct the configured backend.
    ///
    /// S3 selection errors unless the crate was built with `persisted-immutable-data-s3`.
    pub async fn build_backend(&self) -> Result<Arc<dyn UserStorageBackend>> {
        match &self.backend {
            StorageBackendConfig::Local { dir } => Ok(Arc::new(LocalBackend::new(dir))),
            StorageBackendConfig::S3 { bucket, region } => {
                #[cfg(feature = "persisted-immutable-data-s3")]
                {
                    let backend = super::backend::S3Backend::new(bucket.clone(), region.clone()).await?;
                    Ok(Arc::new(backend))
                }
                #[cfg(not(feature = "persisted-immutable-data-s3"))]
                {
                    let _ = (bucket, region);
                    eyre::bail!("S3 storage backend requires the `persisted-immutable-data-s3` feature")
                }
            }
        }
    }

    /// Build the backend and a resolver over it, with no local cache. For processes that
    /// read few objects (the challenger's data provider, tools, tests).
    ///
    /// # Errors
    ///
    /// Fails if the backend cannot be constructed.
    pub async fn build(&self) -> Result<ObjectStorage> {
        let backend = self.build_backend().await?;
        let resolver = PersistedImmutableDataResolver::new(backend.clone());
        Ok(ObjectStorage {
            config: self.clone(),
            backend,
            resolver,
        })
    }

    /// Like [`Self::build`], but an S3 backend is fronted by the local disk cache in
    /// `[storage.local_cache]`. The operator builds its storage this way, once per
    /// process, and the unified binary hands it to the data providers and the gateway, so
    /// they all read one store through one cache. A local backend is read directly.
    ///
    /// # Errors
    ///
    /// Fails if the backend cannot be constructed or the local cache directory cannot be
    /// created or listed.
    pub async fn build_with_local_cache(&self) -> Result<ObjectStorage> {
        let mut storage = self.build().await?;
        storage.resolver = self.front_with_local_cache(storage.resolver).await?;
        Ok(storage)
    }

    /// Put the local cache in front of `resolver` when the backend is S3.
    async fn front_with_local_cache(
        &self,
        resolver: PersistedImmutableDataResolver,
    ) -> Result<PersistedImmutableDataResolver> {
        match self.backend {
            StorageBackendConfig::Local { .. } => Ok(resolver),
            StorageBackendConfig::S3 { .. } => {
                let local = LocalObjectCache::shared(&self.local_cache.dir, self.local_cache.max_bytes).await?;
                Ok(resolver.with_local_cache(local))
            }
        }
    }
}

/// Built content-addressed storage: the config it came from, the backend and the resolver
/// (with its local cache). Cheap to clone; clones share everything.
#[derive(Clone)]
pub struct ObjectStorage {
    /// Configuration, including the upload limits the gateway enforces.
    pub config: StorageConfig,
    /// Raw backend, for markers, listing and deletion the resolver does not expose.
    pub backend: Arc<dyn UserStorageBackend>,
    /// Content-addressed reads and writes, served from the local cache when present.
    pub resolver: PersistedImmutableDataResolver,
}

impl std::fmt::Debug for ObjectStorage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ObjectStorage")
            .field("config", &self.config)
            .finish_non_exhaustive()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn defaults_are_local() {
        let cfg = StorageConfig::default();
        assert!(matches!(cfg.backend, StorageBackendConfig::Local { .. }));
        assert_eq!(cfg.max_object_bytes, 16 * 1024 * 1024);
        assert_eq!(cfg.max_objects_per_owner, 256);
    }

    #[test]
    fn local_config_roundtrips() {
        let json = serde_json::json!({
            "max_object_bytes": 1024,
            "max_objects_per_owner": 8,
            "backend": { "kind": "local", "dir": "/var/newton/persisted-immutable-data" }
        });
        let cfg: StorageConfig = serde_json::from_value(json).unwrap();
        assert_eq!(cfg.max_object_bytes, 1024);
        assert_eq!(cfg.max_objects_per_owner, 8);
        assert_eq!(
            cfg.backend,
            StorageBackendConfig::Local {
                dir: "/var/newton/persisted-immutable-data".into()
            }
        );
    }

    #[test]
    fn s3_config_parses_with_default_region() {
        let json = serde_json::json!({
            "backend": { "kind": "s3", "bucket": "newton-persisted-immutable-data-prod" }
        });
        let cfg: StorageConfig = serde_json::from_value(json).unwrap();
        assert_eq!(
            cfg.backend,
            StorageBackendConfig::S3 {
                bucket: "newton-persisted-immutable-data-prod".into(),
                region: "us-west-2".into(),
            }
        );
    }

    #[tokio::test]
    async fn build_local_backend_works() {
        let tmp = tempfile::tempdir().unwrap();
        let cfg = StorageConfig {
            backend: StorageBackendConfig::Local {
                dir: tmp.path().to_path_buf(),
            },
            ..Default::default()
        };
        let backend = cfg.build_backend().await.unwrap();
        backend.put("k", b"v").await.unwrap();
        assert_eq!(backend.get("k").await.unwrap(), b"v");
    }

    #[tokio::test]
    async fn s3_backend_is_fronted_by_the_local_cache() {
        let backend_dir = tempfile::tempdir().unwrap();
        let cache_root = tempfile::tempdir().unwrap();
        let cache_dir = cache_root.path().join("objects-cache");
        let cfg = StorageConfig {
            backend: StorageBackendConfig::S3 {
                bucket: "bucket".into(),
                region: default_s3_region(),
            },
            local_cache: LocalCacheConfig {
                dir: cache_dir.clone(),
                ..Default::default()
            },
            ..Default::default()
        };
        // A local backend stands in for S3; the config decides whether to cache.
        let resolver = cfg
            .front_with_local_cache(PersistedImmutableDataResolver::new(Arc::new(LocalBackend::new(
                backend_dir.path(),
            ))))
            .await
            .unwrap();
        let cid = resolver.put(b"rego").await.unwrap();
        assert_eq!(resolver.get(&cid).await.unwrap(), b"rego");
        assert!(cache_dir.join("objects").join(&cid).is_file());
    }

    #[tokio::test]
    async fn local_backend_is_read_directly() {
        let backend_dir = tempfile::tempdir().unwrap();
        let cache_root = tempfile::tempdir().unwrap();
        let cache_dir = cache_root.path().join("objects-cache");
        let cfg = StorageConfig {
            backend: StorageBackendConfig::Local {
                dir: backend_dir.path().to_path_buf(),
            },
            local_cache: LocalCacheConfig {
                dir: cache_dir.clone(),
                ..Default::default()
            },
            ..Default::default()
        };
        let resolver = cfg.build_with_local_cache().await.unwrap().resolver;
        let cid = resolver.put(b"rego").await.unwrap();
        assert_eq!(resolver.get(&cid).await.unwrap(), b"rego");
        assert!(!cache_dir.exists());
    }

    #[test]
    fn local_cache_config_defaults_each_field() {
        let unset: StorageConfig = serde_json::from_value(serde_json::json!({
            "backend": { "kind": "local", "dir": "/tmp/x" }
        }))
        .unwrap();
        assert_eq!(unset.local_cache.dir, default_local_cache_dir());
        assert_eq!(unset.local_cache.max_bytes, 20 * 1024 * 1024 * 1024);

        let dir_only: StorageConfig = serde_json::from_value(serde_json::json!({
            "backend": { "kind": "local", "dir": "/tmp/x" },
            "local_cache": { "dir": "/var/cache/newton" }
        }))
        .unwrap();
        assert_eq!(dir_only.local_cache.dir, std::path::PathBuf::from("/var/cache/newton"));
        assert_eq!(dir_only.local_cache.max_bytes, default_local_cache_max_bytes());
    }
}