client-api: Make ControlStateReadAccess an async trait (#3357)

Historically, controldb reads could be treated as just a datastructure,
but that became a lie when reconnections were introduced.

Some tricks were employed to enter the async context when needed, but
those always bear the risk of introducing a deadlock somewhere.

So, just make it async.


# Expected complexity level and risk

2

It may or may not be safe to use sled in an async context. We did
already for
the write path, but if it's a problem it'll show now.

# Testing

Not a functional change.

Author: kim

crates/client-api/src/lib.rs

@@ -201,28 +201,29 @@ pub trait ControlStateDelegate: ControlStateReadAccess + ControlStateWriteAccess
 impl<T: ControlStateReadAccess + ControlStateWriteAccess + Send + Sync> ControlStateDelegate for T {}
 
 /// Query API of the SpacetimeDB control plane.
+#[async_trait]
 pub trait ControlStateReadAccess {
     // Nodes
-    fn get_node_id(&self) -> Option<u64>;
-    fn get_node_by_id(&self, node_id: u64) -> anyhow::Result<Option<Node>>;
-    fn get_nodes(&self) -> anyhow::Result<Vec<Node>>;
+    async fn get_node_id(&self) -> Option<u64>;
+    async fn get_node_by_id(&self, node_id: u64) -> anyhow::Result<Option<Node>>;
+    async fn get_nodes(&self) -> anyhow::Result<Vec<Node>>;
 
     // Databases
-    fn get_database_by_id(&self, id: u64) -> anyhow::Result<Option<Database>>;
-    fn get_database_by_identity(&self, database_identity: &Identity) -> anyhow::Result<Option<Database>>;
-    fn get_databases(&self) -> anyhow::Result<Vec<Database>>;
+    async fn get_database_by_id(&self, id: u64) -> anyhow::Result<Option<Database>>;
+    async fn get_database_by_identity(&self, database_identity: &Identity) -> anyhow::Result<Option<Database>>;
+    async fn get_databases(&self) -> anyhow::Result<Vec<Database>>;
 
     // Replicas
-    fn get_replica_by_id(&self, id: u64) -> anyhow::Result<Option<Replica>>;
-    fn get_replicas(&self) -> anyhow::Result<Vec<Replica>>;
-    fn get_leader_replica_by_database(&self, database_id: u64) -> Option<Replica>;
+    async fn get_replica_by_id(&self, id: u64) -> anyhow::Result<Option<Replica>>;
+    async fn get_replicas(&self) -> anyhow::Result<Vec<Replica>>;
+    async fn get_leader_replica_by_database(&self, database_id: u64) -> Option<Replica>;
 
     // Energy
-    fn get_energy_balance(&self, identity: &Identity) -> anyhow::Result<Option<EnergyBalance>>;
+    async fn get_energy_balance(&self, identity: &Identity) -> anyhow::Result<Option<EnergyBalance>>;
 
     // DNS
-    fn lookup_identity(&self, domain: &str) -> anyhow::Result<Option<Identity>>;
-    fn reverse_lookup(&self, database_identity: &Identity) -> anyhow::Result<Vec<DomainName>>;
+    async fn lookup_identity(&self, domain: &str) -> anyhow::Result<Option<Identity>>;
+    async fn reverse_lookup(&self, database_identity: &Identity) -> anyhow::Result<Vec<DomainName>>;
 }
 
 /// Write operations on the SpacetimeDB control plane.
@@ -281,53 +282,54 @@ pub trait ControlStateWriteAccess: Send + Sync {
     ) -> anyhow::Result<SetDomainsResult>;
 }
 
-impl<T: ControlStateReadAccess + ?Sized> ControlStateReadAccess for Arc<T> {
+#[async_trait]
+impl<T: ControlStateReadAccess + Send + Sync + Sync + ?Sized> ControlStateReadAccess for Arc<T> {
     // Nodes
-    fn get_node_id(&self) -> Option<u64> {
-        (**self).get_node_id()
+    async fn get_node_id(&self) -> Option<u64> {
+        (**self).get_node_id().await
     }
-    fn get_node_by_id(&self, node_id: u64) -> anyhow::Result<Option<Node>> {
-        (**self).get_node_by_id(node_id)
+    async fn get_node_by_id(&self, node_id: u64) -> anyhow::Result<Option<Node>> {
+        (**self).get_node_by_id(node_id).await
     }
-    fn get_nodes(&self) -> anyhow::Result<Vec<Node>> {
-        (**self).get_nodes()
+    async fn get_nodes(&self) -> anyhow::Result<Vec<Node>> {
+        (**self).get_nodes().await
     }
 
     // Databases
-    fn get_database_by_id(&self, id: u64) -> anyhow::Result<Option<Database>> {
-        (**self).get_database_by_id(id)
+    async fn get_database_by_id(&self, id: u64) -> anyhow::Result<Option<Database>> {
+        (**self).get_database_by_id(id).await
     }
-    fn get_database_by_identity(&self, identity: &Identity) -> anyhow::Result<Option<Database>> {
-        (**self).get_database_by_identity(identity)
+    async fn get_database_by_identity(&self, identity: &Identity) -> anyhow::Result<Option<Database>> {
+        (**self).get_database_by_identity(identity).await
     }
-    fn get_databases(&self) -> anyhow::Result<Vec<Database>> {
-        (**self).get_databases()
+    async fn get_databases(&self) -> anyhow::Result<Vec<Database>> {
+        (**self).get_databases().await
     }
 
     // Replicas
-    fn get_replica_by_id(&self, id: u64) -> anyhow::Result<Option<Replica>> {
-        (**self).get_replica_by_id(id)
+    async fn get_replica_by_id(&self, id: u64) -> anyhow::Result<Option<Replica>> {
+        (**self).get_replica_by_id(id).await
     }
-    fn get_replicas(&self) -> anyhow::Result<Vec<Replica>> {
-        (**self).get_replicas()
+    async fn get_replicas(&self) -> anyhow::Result<Vec<Replica>> {
+        (**self).get_replicas().await
     }
 
     // Energy
-    fn get_energy_balance(&self, identity: &Identity) -> anyhow::Result<Option<EnergyBalance>> {
-        (**self).get_energy_balance(identity)
+    async fn get_energy_balance(&self, identity: &Identity) -> anyhow::Result<Option<EnergyBalance>> {
+        (**self).get_energy_balance(identity).await
     }
 
     // DNS
-    fn lookup_identity(&self, domain: &str) -> anyhow::Result<Option<Identity>> {
-        (**self).lookup_identity(domain)
+    async fn lookup_identity(&self, domain: &str) -> anyhow::Result<Option<Identity>> {
+        (**self).lookup_identity(domain).await
     }
 
-    fn reverse_lookup(&self, database_identity: &Identity) -> anyhow::Result<Vec<DomainName>> {
-        (**self).reverse_lookup(database_identity)
+    async fn reverse_lookup(&self, database_identity: &Identity) -> anyhow::Result<Vec<DomainName>> {
+        (**self).reverse_lookup(database_identity).await
     }
 
-    fn get_leader_replica_by_database(&self, database_id: u64) -> Option<Replica> {
-        (**self).get_leader_replica_by_database(database_id)
+    async fn get_leader_replica_by_database(&self, database_id: u64) -> Option<Replica> {
+        (**self).get_leader_replica_by_database(database_id).await
     }
 }
 

crates/client-api/src/routes/database.rs

@@ -447,6 +447,7 @@ where
 
     let replica = worker_ctx
         .get_leader_replica_by_database(database.id)
+        .await
         .ok_or((StatusCode::NOT_FOUND, "Replica not scheduled to this node yet."))?;
     let replica_id = replica.id;
 
@@ -506,6 +507,7 @@ pub(crate) async fn worker_ctx_find_database(
 ) -> axum::response::Result<Option<Database>> {
     worker_ctx
         .get_database_by_identity(database_identity)
+        .await
         .map_err(log_and_500)
 }
 
@@ -594,6 +596,7 @@ pub async fn get_names<S: ControlStateDelegate>(
 
     let names = ctx
         .reverse_lookup(&database_identity)
+        .await
         .map_err(log_and_500)?
         .into_iter()
         .filter_map(|x| String::from(x).try_into().ok())
@@ -697,7 +700,12 @@ pub async fn publish<S: NodeDelegate + ControlStateDelegate + Authorization>(
             .as_ref()
             .ok_or_else(|| bad_request("Clear database requires database name or identity".into()))?;
         if let Ok(identity) = name_or_identity.try_resolve(&ctx).await.map_err(log_and_500)? {
-            if ctx.get_database_by_identity(&identity).map_err(log_and_500)?.is_some() {
+            if ctx
+                .get_database_by_identity(&identity)
+                .await
+                .map_err(log_and_500)?
+                .is_some()
+            {
                 return reset(
                     State(ctx),
                     Path(ResetDatabaseParams {
@@ -727,7 +735,10 @@ pub async fn publish<S: NodeDelegate + ControlStateDelegate + Authorization>(
     log::trace!("Publishing to the identity: {}", database_identity.to_hex());
 
     // Check if the database already exists.
-    let existing = ctx.get_database_by_identity(&database_identity).map_err(log_and_500)?;
+    let existing = ctx
+        .get_database_by_identity(&database_identity)
+        .await
+        .map_err(log_and_500)?;
     match existing.as_ref() {
         // If not, check that the we caller is sufficiently authenticated.
         None => {
@@ -1055,7 +1066,10 @@ pub async fn set_names<S: ControlStateDelegate + Authorization>(
 
     let database_identity = name_or_identity.resolve(&ctx).await?;
 
-    let database = ctx.get_database_by_identity(&database_identity).map_err(log_and_500)?;
+    let database = ctx
+        .get_database_by_identity(&database_identity)
+        .await
+        .map_err(log_and_500)?;
     let Some(database) = database else {
         return Ok((
             StatusCode::NOT_FOUND,
@@ -1077,7 +1091,7 @@ pub async fn set_names<S: ControlStateDelegate + Authorization>(
         })?;
 
     for name in &validated_names {
-        if ctx.lookup_identity(name.as_str()).unwrap().is_some() {
+        if ctx.lookup_identity(name.as_str()).await.unwrap().is_some() {
             return Ok((
                 StatusCode::BAD_REQUEST,
                 axum::Json(name::SetDomainsResult::OtherError(format!(

crates/client-api/src/routes/energy.rs

@@ -22,7 +22,7 @@ pub async fn get_energy_balance<S: ControlStateDelegate>(
     Path(IdentityParams { identity }): Path<IdentityParams>,
 ) -> axum::response::Result<impl IntoResponse> {
     let identity = Identity::from(identity);
-    get_budget_inner(ctx, &identity)
+    get_budget_inner(ctx, &identity).await
 }
 
 #[serde_with::serde_as]
@@ -57,15 +57,20 @@ pub async fn add_energy<S: ControlStateDelegate>(
     // TODO: is this guaranteed to pull the updated balance?
     let balance = ctx
         .get_energy_balance(&auth.claims.identity)
+        .await
         .map_err(log_and_500)?
         .map_or(0, |quanta| quanta.get());
 
     Ok(axum::Json(BalanceResponse { balance }))
 }
 
-fn get_budget_inner(ctx: impl ControlStateDelegate, identity: &Identity) -> axum::response::Result<impl IntoResponse> {
+async fn get_budget_inner(
+    ctx: impl ControlStateDelegate,
+    identity: &Identity,
+) -> axum::response::Result<impl IntoResponse> {
     let balance = ctx
         .get_energy_balance(identity)
+        .await
         .map_err(log_and_500)?
         .map_or(0, |quanta| quanta.get());
 
@@ -103,6 +108,7 @@ pub async fn set_energy_balance<S: ControlStateDelegate>(
         .unwrap_or(0);
     let current_balance = ctx
         .get_energy_balance(&identity)
+        .await
         .map_err(log_and_500)?
         .map_or(0, |quanta| quanta.get());
 

crates/client-api/src/routes/health.rs

@@ -11,6 +11,7 @@ pub async fn health<S: ControlStateDelegate + NodeDelegate>(
 ) -> axum::response::Result<impl IntoResponse> {
     let nodes: Vec<u64> = ctx
         .get_nodes()
+        .await
         .map_err(|_| {
             (
                 StatusCode::INTERNAL_SERVER_ERROR,
@@ -23,8 +24,10 @@ pub async fn health<S: ControlStateDelegate + NodeDelegate>(
     let schedulable = !ctx
         .get_node_by_id(
             ctx.get_node_id()
+                .await
                 .ok_or((StatusCode::INTERNAL_SERVER_ERROR, "Can't get node id"))?,
         )
+        .await
         .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Couldn't get node info"))?
         .map(|n| n.unschedulable)
         .unwrap_or(false);

crates/client-api/src/routes/identity.rs

@@ -79,7 +79,7 @@ pub async fn get_databases<S: ControlStateDelegate>(
 ) -> axum::response::Result<impl IntoResponse> {
     let identity = identity.into();
     // Linear scan for all databases that have this owner, and return their identities
-    let all_dbs = ctx.get_databases().map_err(|e| {
+    let all_dbs = ctx.get_databases().await.map_err(|e| {
         log::error!("Failure when retrieving databases for search: {e}");
         StatusCode::INTERNAL_SERVER_ERROR
     })?;

crates/client-api/src/routes/prometheus.rs

@@ -14,7 +14,7 @@ pub async fn get_sd_config<S: ControlStateReadAccess>(
     State(ctx): State<S>,
 ) -> axum::response::Result<impl IntoResponse> {
     // TODO(cloutiertyler): security
-    let nodes = ctx.get_nodes().map_err(log_and_500)?;
+    let nodes = ctx.get_nodes().await.map_err(log_and_500)?;
 
     let mut targets = Vec::new();
     let labels = HashMap::default();

crates/client-api/src/routes/subscribe.rs

@@ -143,6 +143,7 @@ where
 
     let database = ctx
         .get_database_by_identity(&db_identity)
+        .await
         .unwrap()
         .ok_or(StatusCode::NOT_FOUND)?;
 

crates/client-api/src/util.rs

@@ -100,7 +100,7 @@ impl NameOrIdentity {
     ) -> anyhow::Result<Result<Identity, &DatabaseName>> {
         Ok(match self {
             Self::Identity(identity) => Ok(Identity::from(*identity)),
-            Self::Name(name) => ctx.lookup_identity(name.as_ref())?.ok_or(name),
+            Self::Name(name) => ctx.lookup_identity(name.as_ref()).await?.ok_or(name),
         })
     }
 

crates/standalone/src/lib.rs

@@ -148,13 +148,14 @@ impl NodeDelegate for StandaloneEnv {
     }
 }
 
+#[async_trait]
 impl spacetimedb_client_api::ControlStateReadAccess for StandaloneEnv {
     // Nodes
-    fn get_node_id(&self) -> Option<u64> {
+    async fn get_node_id(&self) -> Option<u64> {
         Some(0)
     }
 
-    fn get_node_by_id(&self, node_id: u64) -> anyhow::Result<Option<Node>> {
+    async fn get_node_by_id(&self, node_id: u64) -> anyhow::Result<Option<Node>> {
         if node_id == 0 {
             return Ok(Some(Node {
                 id: 0,
@@ -166,46 +167,46 @@ impl spacetimedb_client_api::ControlStateReadAccess for StandaloneEnv {
         Ok(None)
     }
 
-    fn get_nodes(&self) -> anyhow::Result<Vec<Node>> {
-        Ok(vec![self.get_node_by_id(0)?.unwrap()])
+    async fn get_nodes(&self) -> anyhow::Result<Vec<Node>> {
+        Ok(vec![self.get_node_by_id(0).await?.unwrap()])
     }
 
     // Databases
-    fn get_database_by_id(&self, id: u64) -> anyhow::Result<Option<Database>> {
+    async fn get_database_by_id(&self, id: u64) -> anyhow::Result<Option<Database>> {
         Ok(self.control_db.get_database_by_id(id)?)
     }
 
-    fn get_database_by_identity(&self, database_identity: &Identity) -> anyhow::Result<Option<Database>> {
+    async fn get_database_by_identity(&self, database_identity: &Identity) -> anyhow::Result<Option<Database>> {
         Ok(self.control_db.get_database_by_identity(database_identity)?)
     }
 
-    fn get_databases(&self) -> anyhow::Result<Vec<Database>> {
+    async fn get_databases(&self) -> anyhow::Result<Vec<Database>> {
         Ok(self.control_db.get_databases()?)
     }
 
     // Replicas
-    fn get_replica_by_id(&self, id: u64) -> anyhow::Result<Option<Replica>> {
+    async fn get_replica_by_id(&self, id: u64) -> anyhow::Result<Option<Replica>> {
         Ok(self.control_db.get_replica_by_id(id)?)
     }
 
-    fn get_replicas(&self) -> anyhow::Result<Vec<Replica>> {
+    async fn get_replicas(&self) -> anyhow::Result<Vec<Replica>> {
         Ok(self.control_db.get_replicas()?)
     }
 
-    fn get_leader_replica_by_database(&self, database_id: u64) -> Option<Replica> {
+    async fn get_leader_replica_by_database(&self, database_id: u64) -> Option<Replica> {
         self.control_db.get_leader_replica_by_database(database_id)
     }
     // Energy
-    fn get_energy_balance(&self, identity: &Identity) -> anyhow::Result<Option<EnergyBalance>> {
+    async fn get_energy_balance(&self, identity: &Identity) -> anyhow::Result<Option<EnergyBalance>> {
         Ok(self.control_db.get_energy_balance(identity)?)
     }
 
     // DNS
-    fn lookup_identity(&self, domain: &str) -> anyhow::Result<Option<Identity>> {
+    async fn lookup_identity(&self, domain: &str) -> anyhow::Result<Option<Identity>> {
         Ok(self.control_db.spacetime_dns(domain)?)
     }
 
-    fn reverse_lookup(&self, database_identity: &Identity) -> anyhow::Result<Vec<DomainName>> {
+    async fn reverse_lookup(&self, database_identity: &Identity) -> anyhow::Result<Vec<DomainName>> {
         Ok(self.control_db.spacetime_reverse_dns(database_identity)?)
     }
 }
@@ -446,7 +447,8 @@ impl spacetimedb_client_api::Authorization for StandaloneEnv {
         action: spacetimedb_client_api::Action,
     ) -> Result<(), spacetimedb_client_api::Unauthorized> {
         let database = self
-            .get_database_by_identity(&database)?
+            .get_database_by_identity(&database)
+            .await?
             .with_context(|| format!("database {database} not found"))
             .with_context(|| format!("Unable to authorize {subject} to perform {action:?})"))?;
         if subject == database.owner_identity {
@@ -467,7 +469,8 @@ impl spacetimedb_client_api::Authorization for StandaloneEnv {
         database: Identity,
     ) -> Result<AuthCtx, spacetimedb_client_api::Unauthorized> {
         let database = self
-            .get_database_by_identity(&database)?
+            .get_database_by_identity(&database)
+            .await?
             .with_context(|| format!("database {database} not found"))
             .with_context(|| format!("Unable to authorize {subject} for SQL"))?;
 

crates/testing/src/modules.rs

@@ -217,8 +217,8 @@ impl CompiledModule {
         .await
         .unwrap();
 
-        let database = env.get_database_by_identity(&db_identity).unwrap().unwrap();
-        let instance = env.get_leader_replica_by_database(database.id).unwrap();
+        let database = env.get_database_by_identity(&db_identity).await.unwrap().unwrap();
+        let instance = env.get_leader_replica_by_database(database.id).await.unwrap();
 
         let client_id = ClientActorId {
             identity,