core, graph, store: Asyncify SubgraphStore.assignment_status

Author: lutter

core/src/subgraph/registrar.rs

@@ -121,7 +121,6 @@ where
         self.start_assigned_subgraphs().await?;
 
         // Spawn a task to handle assignment events.
-        // Blocking due to store interactions. Won't be blocking after #905.
         let assignment_event_stream_cancel_handle =
             self.assignment_event_stream_cancel_guard.handle();
 
@@ -147,14 +146,17 @@ where
                     }
                 });
 
-        graph::spawn_blocking(fut);
+        graph::spawn(fut);
         Ok(())
     }
 
     /// Maps an assignment change to an assignment event by checking the
     /// current state in the database, ignoring changes that do not affect
     /// this node or do not require anything to change.
-    fn map_assignment(&self, change: AssignmentChange) -> Result<Option<AssignmentEvent>, Error> {
+    async fn map_assignment(
+        &self,
+        change: AssignmentChange,
+    ) -> Result<Option<AssignmentEvent>, Error> {
         let (deployment, operation) = change.into_parts();
 
         trace!(self.logger, "Received assignment change";
@@ -167,6 +169,7 @@ where
                 let assigned = self
                     .store
                     .assignment_status(&deployment)
+                    .await
                     .map_err(|e| anyhow!("Failed to get subgraph assignment entity: {}", e))?;
 
                 let logger = self.logger.new(o!("subgraph_id" => deployment.hash.to_string(), "node_id" => self.node_id.to_string()));
@@ -223,7 +226,7 @@ where
                     let this = this.cheap_clone();
 
                     async move {
-                        match this.map_assignment(change) {
+                        match this.map_assignment(change).await {
                             Ok(Some(event)) => stream::once(futures03::future::ok(event)).boxed(),
                             Ok(None) => stream::empty().boxed(),
                             Err(e) => stream::once(futures03::future::err(e)).boxed(),

graph/src/components/store/traits.rs

@@ -121,7 +121,7 @@ pub trait SubgraphStore: Send + Sync + 'static {
     /// the subgraph is assigned to, and `is_paused` is true if the
     /// subgraph is paused.
     /// Returns None if the deployment does not exist.
-    fn assignment_status(
+    async fn assignment_status(
         &self,
         deployment: &DeploymentLocator,
     ) -> Result<Option<(NodeId, bool)>, StoreError>;

store/postgres/src/primary.rs

@@ -2088,8 +2088,12 @@ impl Mirror {
     /// the subgraph is assigned to, and `is_paused` is true if the
     /// subgraph is paused.
     /// Returns None if the deployment does not exist.
-    pub fn assignment_status(&self, site: &Site) -> Result<Option<(NodeId, bool)>, StoreError> {
-        self.read(|conn| queries::assignment_status(conn, site))
+    pub async fn assignment_status(
+        &self,
+        site: Arc<Site>,
+    ) -> Result<Option<(NodeId, bool)>, StoreError> {
+        self.read_async(move |conn| queries::assignment_status(conn, &site))
+            .await
     }
 
     pub fn find_active_site(&self, subgraph: &DeploymentHash) -> Result<Option<Site>, StoreError> {

store/postgres/src/subgraph_store.rs

@@ -1455,12 +1455,12 @@ impl SubgraphStoreTrait for SubgraphStore {
     /// the subgraph is assigned to, and `is_paused` is true if the
     /// subgraph is paused.
     /// Returns None if the deployment does not exist.
-    fn assignment_status(
+    async fn assignment_status(
         &self,
         deployment: &DeploymentLocator,
     ) -> Result<Option<(NodeId, bool)>, StoreError> {
         let site = self.find_site(deployment.id.into())?;
-        self.mirror.assignment_status(site.as_ref())
+        self.mirror.assignment_status(site).await
     }
 
     fn assignments(&self, node: &NodeId) -> Result<Vec<DeploymentLocator>, StoreError> {