store: Move WritableStore into its own module

Author: lutter

store/postgres/src/lib.rs

@@ -48,6 +48,7 @@ mod store;
 mod store_events;
 mod subgraph_store;
 pub mod transaction_receipt;
+mod writable;
 
 #[cfg(debug_assertions)]
 pub mod layout_for_tests {

store/postgres/src/subgraph_store.rs

@@ -5,7 +5,6 @@ use diesel::{
     types::{FromSql, ToSql},
 };
 use std::{
-    collections::BTreeMap,
     collections::HashMap,
     sync::{Arc, Mutex},
 };
@@ -17,32 +16,29 @@ use graph::{
     components::{
         server::index_node::VersionInfo,
         store::{
-            self, DeploymentLocator, EnsLookup as EnsLookupTrait, EntityType,
+            self, DeploymentLocator, EnsLookup as EnsLookupTrait,
             WritableStore as WritableStoreTrait,
         },
     },
     constraint_violation,
     data::query::QueryTarget,
-    data::subgraph::schema::{self, SubgraphError},
     data::subgraph::status,
     prelude::StoreEvent,
     prelude::SubgraphDeploymentEntity,
     prelude::{
         anyhow, futures03::future::join_all, lazy_static, o, web3::types::Address, ApiSchema,
-        BlockPtr, DeploymentHash, Entity, EntityKey, EntityModification, Error, Logger, NodeId,
-        Schema, StopwatchMetrics, StoreError, SubgraphName, SubgraphStore as SubgraphStoreTrait,
-        SubgraphVersionSwitchingMode,
+        BlockPtr, DeploymentHash, Logger, NodeId, Schema, StoreError, SubgraphName,
+        SubgraphStore as SubgraphStoreTrait, SubgraphVersionSwitchingMode,
     },
-    slog::{error, warn},
-    util::{backoff::ExponentialBackoff, timed_cache::TimedCache},
+    util::timed_cache::TimedCache,
 };
-use store::StoredDynamicDataSource;
 
 use crate::{
     connection_pool::ConnectionPool,
     primary,
     primary::{DeploymentId, Mirror as PrimaryMirror, Site},
     relational::Layout,
+    writable::WritableStore,
     NotificationSender,
 };
 use crate::{
@@ -382,13 +378,13 @@ impl SubgraphStoreInner {
         Ok((store, site))
     }
 
-    fn for_site(&self, site: &Site) -> Result<&Arc<DeploymentStore>, StoreError> {
+    pub(crate) fn for_site(&self, site: &Site) -> Result<&Arc<DeploymentStore>, StoreError> {
         self.stores
             .get(&site.shard)
             .ok_or(StoreError::UnknownShard(site.shard.as_str().to_string()))
     }
 
-    fn layout(&self, id: &DeploymentHash) -> Result<Arc<Layout>, StoreError> {
+    pub(crate) fn layout(&self, id: &DeploymentHash) -> Result<Arc<Layout>, StoreError> {
         let (store, site) = self.store(id)?;
         store.find_layout(site)
     }
@@ -669,7 +665,7 @@ impl SubgraphStoreInner {
     /// connections can deadlock the entire process if the pool runs out
     /// of connections in between getting the first one and trying to get the
     /// second one.
-    fn primary_conn(&self) -> Result<primary::Connection, StoreError> {
+    pub(crate) fn primary_conn(&self) -> Result<primary::Connection, StoreError> {
         let conn = self.mirror.primary().get()?;
         Ok(primary::Connection::new(conn))
     }
@@ -926,7 +922,7 @@ impl SubgraphStoreInner {
     pub fn find(
         &self,
         query: graph::prelude::EntityQuery,
-    ) -> Result<Vec<Entity>, graph::prelude::QueryExecutionError> {
+    ) -> Result<Vec<graph::prelude::Entity>, graph::prelude::QueryExecutionError> {
         let (store, site) = self.store(&query.subgraph_id)?;
         store.find(site, query)
     }
@@ -1107,295 +1103,3 @@ impl SubgraphStoreTrait for SubgraphStore {
             .collect())
     }
 }
-
-/// A wrapper around `SubgraphStore` that only exposes functions that are
-/// safe to call from `WritableStore`, i.e., functions that either do not
-/// deal with anything that depends on a specific deployment
-/// location/instance, or where the result is independent of the deployment
-/// instance
-struct WritableSubgraphStore(SubgraphStore);
-
-impl WritableSubgraphStore {
-    fn primary_conn(&self) -> Result<primary::Connection, StoreError> {
-        self.0.primary_conn()
-    }
-
-    pub(crate) fn send_store_event(&self, event: &StoreEvent) -> Result<(), StoreError> {
-        self.0.send_store_event(event)
-    }
-
-    fn layout(&self, id: &DeploymentHash) -> Result<Arc<Layout>, StoreError> {
-        self.0.layout(id)
-    }
-}
-
-struct WritableStore {
-    logger: Logger,
-    store: WritableSubgraphStore,
-    writable: Arc<DeploymentStore>,
-    site: Arc<Site>,
-    input_schema: Arc<Schema>,
-}
-
-impl WritableStore {
-    const BACKOFF_BASE: Duration = Duration::from_millis(100);
-    const BACKOFF_CEIL: Duration = Duration::from_secs(10);
-
-    fn new(
-        subgraph_store: SubgraphStore,
-        logger: Logger,
-        site: Arc<Site>,
-    ) -> Result<Self, StoreError> {
-        let store = WritableSubgraphStore(subgraph_store.clone());
-        let writable = subgraph_store.for_site(site.as_ref())?.clone();
-        let input_schema = subgraph_store.input_schema(&site.deployment)?;
-        Ok(Self {
-            logger,
-            store,
-            writable,
-            site,
-            input_schema,
-        })
-    }
-
-    fn log_backoff_warning(&self, op: &str, backoff: &ExponentialBackoff) {
-        warn!(self.logger,
-            "database unavailable, will retry";
-            "operation" => op,
-            "attempt" => backoff.attempt,
-            "delay_ms" => backoff.delay().as_millis());
-    }
-
-    fn retry<T, F>(&self, op: &str, f: F) -> Result<T, StoreError>
-    where
-        F: Fn() -> Result<T, StoreError>,
-    {
-        let mut backoff = ExponentialBackoff::new(Self::BACKOFF_BASE, Self::BACKOFF_CEIL);
-        loop {
-            match f() {
-                Ok(v) => return Ok(v),
-                Err(StoreError::DatabaseUnavailable) => {
-                    self.log_backoff_warning(op, &backoff);
-                }
-                Err(e) => return Err(e),
-            }
-            backoff.sleep();
-        }
-    }
-
-    async fn retry_async<T, F, Fut>(&self, op: &str, f: F) -> Result<T, StoreError>
-    where
-        F: Fn() -> Fut,
-        Fut: std::future::Future<Output = Result<T, StoreError>>,
-    {
-        let mut backoff = ExponentialBackoff::new(Self::BACKOFF_BASE, Self::BACKOFF_CEIL);
-        loop {
-            match f().await {
-                Ok(v) => return Ok(v),
-                Err(StoreError::DatabaseUnavailable) => {
-                    self.log_backoff_warning(op, &backoff);
-                }
-                Err(e) => return Err(e),
-            }
-            backoff.sleep_async().await;
-        }
-    }
-
-    /// Try to send a `StoreEvent`; if sending fails, log the error but
-    /// return `Ok(())`
-    fn try_send_store_event(&self, event: StoreEvent) -> Result<(), StoreError> {
-        if *SEND_SUBSCRIPTION_NOTIFICATIONS {
-            let _ = self.store.send_store_event(&event).map_err(
-                |e| error!(self.logger, "Could not send store event"; "error" => e.to_string()),
-            );
-            Ok(())
-        } else {
-            Ok(())
-        }
-    }
-}
-
-#[async_trait::async_trait]
-impl WritableStoreTrait for WritableStore {
-    fn block_ptr(&self) -> Result<Option<BlockPtr>, StoreError> {
-        self.retry("block_ptr", || self.writable.block_ptr(self.site.as_ref()))
-    }
-
-    fn block_cursor(&self) -> Result<Option<String>, StoreError> {
-        self.writable.block_cursor(self.site.as_ref())
-    }
-
-    fn start_subgraph_deployment(&self, logger: &Logger) -> Result<(), StoreError> {
-        self.retry("start_subgraph_deployment", || {
-            let store = &self.writable;
-
-            let graft_base = match store.graft_pending(&self.site.deployment)? {
-                Some((base_id, base_ptr)) => {
-                    let src = self.store.layout(&base_id)?;
-                    Some((src, base_ptr))
-                }
-                None => None,
-            };
-            store.start_subgraph(logger, self.site.clone(), graft_base)?;
-            self.store.primary_conn()?.copy_finished(self.site.as_ref())
-        })
-    }
-
-    fn revert_block_operations(&self, block_ptr_to: BlockPtr) -> Result<(), StoreError> {
-        self.retry("revert_block_operations", || {
-            let event = self
-                .writable
-                .revert_block_operations(self.site.clone(), block_ptr_to.clone())?;
-            self.try_send_store_event(event)
-        })
-    }
-
-    fn unfail_deterministic_error(
-        &self,
-        current_ptr: &BlockPtr,
-        parent_ptr: &BlockPtr,
-    ) -> Result<(), StoreError> {
-        self.retry("unfail_deterministic_error", || {
-            self.writable
-                .unfail_deterministic_error(self.site.clone(), current_ptr, parent_ptr)
-        })
-    }
-
-    fn unfail_non_deterministic_error(&self, current_ptr: &BlockPtr) -> Result<(), StoreError> {
-        self.retry("unfail_non_deterministic_error", || {
-            self.writable
-                .unfail_non_deterministic_error(self.site.clone(), current_ptr)
-        })
-    }
-
-    async fn fail_subgraph(&self, error: SubgraphError) -> Result<(), StoreError> {
-        self.retry_async("fail_subgraph", || {
-            let error = error.clone();
-            async {
-                self.writable
-                    .clone()
-                    .fail_subgraph(self.site.deployment.clone(), error)
-                    .await
-            }
-        })
-        .await
-    }
-
-    async fn supports_proof_of_indexing(&self) -> Result<bool, StoreError> {
-        self.retry_async("supports_proof_of_indexing", || async {
-            self.writable
-                .supports_proof_of_indexing(self.site.clone())
-                .await
-        })
-        .await
-    }
-
-    fn get(&self, key: &EntityKey) -> Result<Option<Entity>, StoreError> {
-        self.retry("get", || self.writable.get(self.site.cheap_clone(), key))
-    }
-
-    fn transact_block_operations(
-        &self,
-        block_ptr_to: BlockPtr,
-        firehose_cursor: Option<String>,
-        mods: Vec<EntityModification>,
-        stopwatch: StopwatchMetrics,
-        data_sources: Vec<StoredDynamicDataSource>,
-        deterministic_errors: Vec<SubgraphError>,
-    ) -> Result<(), StoreError> {
-        assert!(
-            same_subgraph(&mods, &self.site.deployment),
-            "can only transact operations within one shard"
-        );
-        self.retry("transact_block_operations", move || {
-            let event = self.writable.transact_block_operations(
-                self.site.clone(),
-                &block_ptr_to,
-                firehose_cursor.as_deref(),
-                &mods,
-                stopwatch.cheap_clone(),
-                &data_sources,
-                &deterministic_errors,
-            )?;
-
-            let _section = stopwatch.start_section("send_store_event");
-            self.try_send_store_event(event)
-        })
-    }
-
-    fn get_many(
-        &self,
-        ids_for_type: BTreeMap<&EntityType, Vec<&str>>,
-    ) -> Result<BTreeMap<EntityType, Vec<Entity>>, StoreError> {
-        self.retry("get_many", || {
-            self.writable
-                .get_many(self.site.cheap_clone(), &ids_for_type)
-        })
-    }
-
-    async fn is_deployment_synced(&self) -> Result<bool, StoreError> {
-        self.retry_async("is_deployment_synced", || async {
-            self.writable
-                .exists_and_synced(self.site.deployment.cheap_clone())
-                .await
-        })
-        .await
-    }
-
-    fn unassign_subgraph(&self) -> Result<(), StoreError> {
-        self.retry("unassign_subgraph", || {
-            let pconn = self.store.primary_conn()?;
-            pconn.transaction(|| -> Result<_, StoreError> {
-                let changes = pconn.unassign_subgraph(self.site.as_ref())?;
-                pconn.send_store_event(&self.store.0.sender, &StoreEvent::new(changes))
-            })
-        })
-    }
-
-    async fn load_dynamic_data_sources(&self) -> Result<Vec<StoredDynamicDataSource>, StoreError> {
-        self.retry_async("load_dynamic_data_sources", || async {
-            self.writable
-                .load_dynamic_data_sources(self.site.deployment.clone())
-                .await
-        })
-        .await
-    }
-
-    fn deployment_synced(&self) -> Result<(), StoreError> {
-        self.retry("deployment_synced", || {
-            let event = {
-                // Make sure we drop `pconn` before we call into the deployment
-                // store so that we do not hold two database connections which
-                // might come from the same pool and could therefore deadlock
-                let pconn = self.store.primary_conn()?;
-                pconn.transaction(|| -> Result<_, Error> {
-                    let changes = pconn.promote_deployment(&self.site.deployment)?;
-                    Ok(StoreEvent::new(changes))
-                })?
-            };
-
-            self.writable.deployment_synced(&self.site.deployment)?;
-
-            self.store.send_store_event(&event)
-        })
-    }
-
-    fn shard(&self) -> &str {
-        self.site.shard.as_str()
-    }
-
-    async fn health(&self, id: &DeploymentHash) -> Result<schema::SubgraphHealth, StoreError> {
-        self.retry_async("health", || async {
-            self.writable.health(id).await.map(Into::into)
-        })
-        .await
-    }
-
-    fn input_schema(&self) -> Arc<Schema> {
-        self.input_schema.clone()
-    }
-}
-
-fn same_subgraph(mods: &Vec<EntityModification>, id: &DeploymentHash) -> bool {
-    mods.iter().all(|md| &md.entity_key().subgraph_id == id)
-}