core, graph, store: Restart writable when restarting subgraph

When the subgraph runner encounters an error, it needs to restart the store
to clear any errors that might have happened.

Author: lutter

core/src/subgraph/inputs.rs

@@ -33,3 +33,42 @@ pub struct IndexingInputs<C: Blockchain> {
     /// possibly expensive and noisy, information
     pub instrument: bool,
 }
+
+impl<C: Blockchain> IndexingInputs<C> {
+    pub fn with_store(&self, store: Arc<dyn WritableStore>) -> Self {
+        let IndexingInputs {
+            deployment,
+            features,
+            start_blocks,
+            stop_block,
+            store: _,
+            debug_fork,
+            triggers_adapter,
+            chain,
+            templates,
+            unified_api_version,
+            static_filters,
+            poi_version,
+            network,
+            manifest_idx_and_name,
+            instrument,
+        } = self;
+        IndexingInputs {
+            deployment: deployment.clone(),
+            features: features.clone(),
+            start_blocks: start_blocks.clone(),
+            stop_block: stop_block.clone(),
+            store,
+            debug_fork: debug_fork.clone(),
+            triggers_adapter: triggers_adapter.clone(),
+            chain: chain.clone(),
+            templates: templates.clone(),
+            unified_api_version: unified_api_version.clone(),
+            static_filters: *static_filters,
+            poi_version: *poi_version,
+            network: network.clone(),
+            manifest_idx_and_name: manifest_idx_and_name.clone(),
+            instrument: *instrument,
+        }
+    }
+}

core/src/subgraph/runner.rs

@@ -179,7 +179,19 @@ where
                         self.inputs.store.flush().await?;
                         return Ok(self);
                     }
-                    Action::Restart => break,
+                    Action::Restart => {
+                        // Restart the store to clear any errors that it
+                        // might have encountered and use that from now on
+                        let store = self.inputs.store.cheap_clone();
+                        let store = store.restart().await?;
+                        self.inputs = Arc::new(self.inputs.with_store(store));
+                        // Also clear the entity cache since we might have
+                        // entries in there that never made it to the
+                        // database
+                        self.state.entity_lfu_cache = LfuCache::new();
+                        self.state.synced = self.inputs.store.is_deployment_synced().await?;
+                        break;
+                    }
                 };
             }
         }

graph/src/components/store/traits.rs

@@ -324,6 +324,17 @@ pub trait WritableStore: ReadStore + DeploymentCursorTracker {
 
     /// Wait for the background writer to finish processing its queue
     async fn flush(&self) -> Result<(), StoreError>;
+
+    /// Restart the `WritableStore`. This will clear any errors that have
+    /// been encountered. Code that calls this must not make any assumptions
+    /// about what has been written already, as the write queue might
+    /// contain unprocessed write requests that will be discarded by this
+    /// call.
+    ///
+    /// After this call, `self` should not be used anymore, as it will
+    /// continue to produce errors for any write requests, and instead, the
+    /// returned `WritableStore` should be used.
+    async fn restart(self: Arc<Self>) -> Result<Arc<dyn WritableStore>, StoreError>;
 }
 
 #[async_trait]

store/postgres/src/writable.rs

@@ -13,7 +13,7 @@ use graph::prelude::{
     BLOCK_NUMBER_MAX,
 };
 use graph::schema::InputSchema;
-use graph::slog::info;
+use graph::slog::{info, warn};
 use graph::tokio::task::JoinHandle;
 use graph::util::bounded_queue::BoundedQueue;
 use graph::{
@@ -1313,4 +1313,17 @@ impl WritableStoreTrait for WritableStore {
     async fn flush(&self) -> Result<(), StoreError> {
         self.writer.flush().await
     }
+
+    async fn restart(self: Arc<Self>) -> Result<Arc<dyn WritableStoreTrait>, StoreError> {
+        if self.poisoned() {
+            let logger = self.store.logger.clone();
+            if let Err(e) = self.stop().await {
+                warn!(logger, "Writable had error when stopping, it is safe to ignore this error"; "error" => e.to_string());
+            }
+            let store = Arc::new(self.store.store.0.clone());
+            store.writable(logger, self.store.site.id.into()).await
+        } else {
+            Ok(self)
+        }
+    }
 }

store/test-store/tests/graph/entity_cache.rs

@@ -169,6 +169,10 @@ impl WritableStore for MockStore {
     async fn causality_region_curr_val(&self) -> Result<Option<CausalityRegion>, StoreError> {
         unimplemented!()
     }
+
+    async fn restart(self: Arc<Self>) -> Result<Arc<dyn WritableStore>, StoreError> {
+        unimplemented!()
+    }
 }
 
 fn make_band_key(id: &'static str) -> EntityKey {

store/test-store/tests/postgres/writable.rs

@@ -16,7 +16,7 @@ use web3::types::H256;
 const SCHEMA_GQL: &str = "
     type Counter @entity {
         id: ID!,
-        count: Int,
+        count: Int!,
     }
 ";
 
@@ -93,8 +93,7 @@ where
 
         // Run test and wait for the background writer to finish its work so
         // it won't conflict with the next test
-        test(store, writable.clone(), deployment).await;
-        writable.flush().await.unwrap();
+        test(store, writable, deployment).await;
     });
 }
 
@@ -161,5 +160,58 @@ fn tracker() {
 
         resume_writer(&deployment, 1).await;
         assert_eq!(2, read_count());
+
+        // There shouldn't be anything left to do, but make sure of that
+        writable.flush().await.unwrap();
+    })
+}
+
+#[test]
+fn restart() {
+    run_test(|store, writable, deployment| async move {
+        let subgraph_store = store.subgraph_store();
+
+        // Cause an error by leaving out the non-nullable `count` attribute
+        let entity_ops = vec![EntityOperation::Set {
+            key: count_key("1"),
+            data: entity! { id: "1" },
+        }];
+        transact_entity_operations(
+            &subgraph_store,
+            &deployment,
+            block_pointer(1),
+            entity_ops.clone(),
+        )
+        .await
+        .unwrap();
+        // flush checks for errors and therefore fails
+        writable
+            .flush()
+            .await
+            .expect_err("writing with missing non-nullable field should fail");
+
+        // We now have a poisoned store. Restarting it gives us a new store
+        // that works again
+        let writable = writable.restart().await.unwrap();
+        writable.flush().await.unwrap();
+
+        // Retry our write with correct data
+        let entity_ops = vec![EntityOperation::Set {
+            key: count_key("1"),
+            data: entity! { id: "1", count: 1 },
+        }];
+        // `SubgraphStore` caches the correct writable so that this call
+        // uses the restarted writable, and is equivalent to using
+        // `writable` directly
+        transact_entity_operations(
+            &subgraph_store,
+            &deployment,
+            block_pointer(1),
+            entity_ops.clone(),
+        )
+        .await
+        .unwrap();
+        // Look, no errors
+        writable.flush().await.unwrap();
     })
 }