store: Asyncify primary::Connection.least_used_shard

Author: lutter

store/postgres/src/primary.rs

@@ -1585,7 +1585,10 @@ impl Connection {
     /// that are stored in it. Unassigned deployments are ignored; in
     /// particular, that ignores deployments that are going to be removed
     /// soon.
-    pub fn least_used_shard(&mut self, shards: &[Shard]) -> Result<Option<Shard>, StoreError> {
+    pub async fn least_used_shard(
+        &mut self,
+        shards: &[Shard],
+    ) -> Result<Option<Shard>, StoreError> {
         use deployment_schemas as ds;
         use subgraph_deployment_assignment as a;
 

store/postgres/src/subgraph_store.rs

@@ -666,15 +666,15 @@ impl Inner {
         }
     }
 
-    fn place_in_shard(&self, mut shards: Vec<Shard>) -> Result<Shard, StoreError> {
+    async fn place_in_shard(&self, mut shards: Vec<Shard>) -> Result<Shard, StoreError> {
         match shards.len() {
             0 => Ok(PRIMARY_SHARD.clone()),
             1 => Ok(shards.pop().unwrap()),
             _ => {
                 let mut conn = self.primary_conn()?;
 
                 // unwrap is fine since shards is not empty
-                let shard = conn.least_used_shard(&shards)?.unwrap();
+                let shard = conn.least_used_shard(&shards).await?.unwrap();
                 Ok(shard)
             }
         }
@@ -700,7 +700,7 @@ impl Inner {
             None => Ok((PRIMARY_SHARD.clone(), default_node)),
             Some((shards, nodes)) => {
                 let node = self.place_on_node(nodes, default_node).await?;
-                let shard = self.place_in_shard(shards)?;
+                let shard = self.place_in_shard(shards).await?;
 
                 Ok((shard, node))
             }