namespaced usage limits

MarinPostma · MarinPostma · commit ac2ad43ff79c · 2023-09-15T20:48:25.000+02:00
diff --git a/sqld/src/connection/config.rs b/sqld/src/connection/config.rs
@@ -7,6 +7,7 @@ use std::{fs, io};
 use crate::error::Error;
 use crate::Result;
 
+#[derive(Debug)]
 pub struct DatabaseConfigStore {
     config_path: PathBuf,
     tmp_config_path: PathBuf,
diff --git a/sqld/src/http/admin/mod.rs b/sqld/src/http/admin/mod.rs
@@ -11,31 +11,28 @@ use tokio_util::io::ReaderStream;
 use url::Url;
 use uuid::Uuid;
 
-use crate::connection::config::{DatabaseConfig, DatabaseConfigStore};
+use crate::connection::config::DatabaseConfig;
 use crate::error::LoadDumpError;
 use crate::namespace::{DumpStream, MakeNamespace, NamespaceStore, RestoreOption};
 
 pub mod stats;
 
 struct AppState<M: MakeNamespace> {
-    db_config_store: Arc<DatabaseConfigStore>,
     namespaces: NamespaceStore<M>,
 }
 
-pub async fn run<M, A>(
-    acceptor: A,
-    db_config_store: Arc<DatabaseConfigStore>,
-    namespaces: NamespaceStore<M>,
-) -> anyhow::Result<()>
+pub async fn run<M, A>(acceptor: A, namespaces: NamespaceStore<M>) -> anyhow::Result<()>
 where
     A: crate::net::Accept,
     M: MakeNamespace,
 {
     use axum::routing::{get, post};
     let router = axum::Router::new()
         .route("/", get(handle_get_index))
-        .route("/v1/config", get(handle_get_config))
-        .route("/v1/block", post(handle_post_block))
+        .route(
+            "/v1/namespaces/:namespace/config",
+            get(handle_get_config).post(handle_post_config),
+        )
         .route(
             "/v1/namespaces/:namespace/fork/:to",
             post(handle_fork_namespace),
@@ -50,10 +47,7 @@ where
         )
         .route("/v1/namespaces/:namespace", delete(handle_delete_namespace))
         .route("/v1/namespaces/:namespace/stats", get(stats::handle_stats))
-        .with_state(Arc::new(AppState {
-            db_config_store,
-            namespaces,
-        }));
+        .with_state(Arc::new(AppState { namespaces }));
 
     hyper::server::Server::builder(acceptor)
         .serve(router.into_make_service())
@@ -68,8 +62,10 @@ async fn handle_get_index() -> &'static str {
 
 async fn handle_get_config<M: MakeNamespace>(
     State(app_state): State<Arc<AppState<M>>>,
-) -> Json<Arc<DatabaseConfig>> {
-    Json(app_state.db_config_store.get())
+    Path(namespace): Path<String>,
+) -> crate::Result<Json<Arc<DatabaseConfig>>> {
+    let store = app_state.namespaces.config_store(namespace.into()).await?;
+    Ok(Json(store.get()))
 }
 
 #[derive(Debug, Deserialize)]
@@ -85,22 +81,20 @@ struct CreateNamespaceReq {
     dump_url: Option<Url>,
 }
 
-async fn handle_post_block<M: MakeNamespace>(
+async fn handle_post_config<M: MakeNamespace>(
     State(app_state): State<Arc<AppState<M>>>,
+    Path(namespace): Path<String>,
     Json(req): Json<BlockReq>,
-) -> (axum::http::StatusCode, &'static str) {
-    let mut config = (*app_state.db_config_store.get()).clone();
+) -> crate::Result<()> {
+    let store = app_state.namespaces.config_store(namespace.into()).await?;
+    let mut config = (*store.get()).clone();
     config.block_reads = req.block_reads;
     config.block_writes = req.block_writes;
     config.block_reason = req.block_reason;
 
-    match app_state.db_config_store.store(config) {
-        Ok(()) => (axum::http::StatusCode::OK, "OK"),
-        Err(err) => {
-            tracing::warn!("Could not store database config: {err}");
-            (axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Failed")
-        }
-    }
+    store.store(config)?;
+
+    Ok(())
 }
 
 async fn handle_create_namespace<M: MakeNamespace>(
diff --git a/sqld/src/lib.rs b/sqld/src/lib.rs
@@ -32,7 +32,6 @@ use url::Url;
 use utils::services::idle_shutdown::IdleShutdownKicker;
 
 use crate::auth::Auth;
-use crate::connection::config::DatabaseConfigStore;
 use crate::connection::{Connection, MakeConnection};
 use crate::error::Error;
 use crate::migration::maybe_migrate;
@@ -91,7 +90,6 @@ pub struct Server<C = HttpConnector, A = AddrIncoming> {
 struct Services<M: MakeNamespace, A, P, S> {
     namespaces: NamespaceStore<M>,
     idle_shutdown_kicker: Option<IdleShutdownKicker>,
-    db_config_store: Arc<DatabaseConfigStore>,
     proxy_service: P,
     replication_service: S,
     user_api_config: UserApiConfig<A>,
@@ -130,11 +128,7 @@ where
         user_http.configure(join_set);
 
         if let Some(AdminApiConfig { acceptor }) = self.admin_api_config {
-            join_set.spawn(http::admin::run(
-                acceptor,
-                self.db_config_store,
-                self.namespaces,
-            ));
+            join_set.spawn(http::admin::run(acceptor, self.namespaces));
         }
     }
 }
@@ -316,9 +310,6 @@ where
         let db_is_dirty = init_sentinel_file(&self.path)?;
         let idle_shutdown_kicker = self.setup_shutdown();
 
-        let db_config_store = Arc::new(
-            DatabaseConfigStore::load(&self.path).context("Could not load database config")?,
-        );
         let snapshot_callback = self.make_snapshot_callback();
         let auth = self.user_api_config.get_auth()?.into();
         let extensions = self.db_config.validate_extensions()?;
@@ -328,7 +319,6 @@ where
                 let replica = Replica {
                     rpc_config,
                     stats_sender,
-                    db_config_store: db_config_store.clone(),
                     extensions,
                     db_config: self.db_config.clone(),
                     base_path: self.path.clone(),
@@ -337,7 +327,6 @@ where
                 let services = Services {
                     namespaces,
                     idle_shutdown_kicker,
-                    db_config_store,
                     proxy_service,
                     replication_service,
                     user_api_config: self.user_api_config,
@@ -357,7 +346,6 @@ where
                     db_config: self.db_config.clone(),
                     idle_shutdown_kicker: idle_shutdown_kicker.clone(),
                     stats_sender,
-                    db_config_store: db_config_store.clone(),
                     db_is_dirty,
                     snapshot_callback,
                     extensions,
@@ -371,7 +359,6 @@ where
                 let services = Services {
                     namespaces,
                     idle_shutdown_kicker,
-                    db_config_store,
                     proxy_service,
                     replication_service,
                     user_api_config: self.user_api_config,
@@ -415,7 +402,6 @@ struct Primary<'a, A> {
     db_config: DbConfig,
     idle_shutdown_kicker: Option<IdleShutdownKicker>,
     stats_sender: StatsSender,
-    db_config_store: Arc<DatabaseConfigStore>,
     db_is_dirty: bool,
     snapshot_callback: NamespacedSnapshotCallback,
     extensions: Arc<[PathBuf]>,
@@ -445,7 +431,6 @@ where
             bottomless_replication: self.db_config.bottomless_replication,
             extensions: self.extensions,
             stats_sender: self.stats_sender.clone(),
-            config_store: self.db_config_store,
             max_response_size: self.db_config.max_response_size,
             max_total_response_size: self.db_config.max_total_response_size,
             checkpoint_interval: self.db_config.checkpoint_interval,
@@ -491,7 +476,6 @@ where
 struct Replica<C> {
     rpc_config: RpcClientConfig<C>,
     stats_sender: StatsSender,
-    db_config_store: Arc<DatabaseConfigStore>,
     extensions: Arc<[PathBuf]>,
     db_config: DbConfig,
     base_path: Arc<Path>,
@@ -512,7 +496,6 @@ impl<C: Connector> Replica<C> {
             uri: uri.clone(),
             extensions: self.extensions.clone(),
             stats_sender: self.stats_sender.clone(),
-            config_store: self.db_config_store.clone(),
             base_path: self.base_path,
             max_response_size: self.db_config.max_response_size,
             max_total_response_size: self.db_config.max_total_response_size,
diff --git a/sqld/src/namespace/mod.rs b/sqld/src/namespace/mod.rs
@@ -403,9 +403,16 @@ impl<M: MakeNamespace> NamespaceStore<M> {
         Ok(())
     }
 
-    pub async fn stats(&self, namespace: Bytes) -> crate::Result<Arc<Stats>> {
+    pub(crate) async fn stats(&self, namespace: Bytes) -> crate::Result<Arc<Stats>> {
         self.with(namespace, |ns| ns.stats.clone()).await
     }
+
+    pub(crate) async fn config_store(
+        &self,
+        namespace: Bytes,
+    ) -> crate::Result<Arc<DatabaseConfigStore>> {
+        self.with(namespace, |ns| ns.db_config_store.clone()).await
+    }
 }
 
 /// A namspace isolates the resources pertaining to a database of type T
@@ -415,6 +422,7 @@ pub struct Namespace<T: Database> {
     /// The set of tasks associated with this namespace
     tasks: JoinSet<anyhow::Result<()>>,
     stats: Arc<Stats>,
+    db_config_store: Arc<DatabaseConfigStore>,
 }
 
 impl<T: Database> Namespace<T> {
@@ -438,8 +446,6 @@ pub struct ReplicaNamespaceConfig {
     pub extensions: Arc<[PathBuf]>,
     /// Stats monitor
     pub stats_sender: StatsSender,
-    /// Reference to the config store
-    pub config_store: Arc<DatabaseConfigStore>,
 }
 
 impl Namespace<ReplicaDatabase> {
@@ -459,6 +465,10 @@ impl Namespace<ReplicaDatabase> {
             ));
         }
 
+        let db_config_store = Arc::new(
+            DatabaseConfigStore::load(&db_path).context("Could not load database config")?,
+        );
+
         let mut join_set = JoinSet::new();
         let replicator = Replicator::new(
             db_path.clone(),
@@ -489,7 +499,7 @@ impl Namespace<ReplicaDatabase> {
             config.channel.clone(),
             config.uri.clone(),
             stats.clone(),
-            config.config_store.clone(),
+            db_config_store.clone(),
             applied_frame_no_receiver,
             config.max_response_size,
             config.max_total_response_size,
@@ -507,6 +517,7 @@ impl Namespace<ReplicaDatabase> {
                 connection_maker: Arc::new(connection_maker),
             },
             stats,
+            db_config_store,
         })
     }
 }
@@ -520,7 +531,6 @@ pub struct PrimaryNamespaceConfig {
     pub bottomless_replication: Option<bottomless::replicator::Options>,
     pub extensions: Arc<[PathBuf]>,
     pub stats_sender: StatsSender,
-    pub config_store: Arc<DatabaseConfigStore>,
     pub max_response_size: u64,
     pub max_total_response_size: u64,
     pub checkpoint_interval: Option<Duration>,
@@ -625,12 +635,16 @@ impl Namespace<PrimaryDatabase> {
         )
         .await?;
 
+        let db_config_store = Arc::new(
+            DatabaseConfigStore::load(&db_path).context("Could not load database config")?,
+        );
+
         let connection_maker: Arc<_> = LibSqlDbFactory::new(
             db_path.clone(),
             &REPLICATION_METHODS,
             ctx_builder.clone(),
             stats.clone(),
-            config.config_store.clone(),
+            db_config_store.clone(),
             config.extensions.clone(),
             config.max_response_size,
             config.max_total_response_size,
@@ -673,6 +687,7 @@ impl Namespace<PrimaryDatabase> {
                 connection_maker,
             },
             stats,
+            db_config_store,
         })
     }
 }
