store: Move AdaptiveBatchSize to its own module

Author: lutter

store/postgres/src/copy.rs

@@ -19,21 +19,17 @@ use std::{
 };
 
 use diesel::{
-    deserialize::FromSql,
     dsl::sql,
     insert_into,
-    pg::Pg,
     r2d2::{ConnectionManager, PooledConnection},
-    select,
-    serialize::{Output, ToSql},
-    sql_query,
+    select, sql_query,
     sql_types::{BigInt, Integer},
     update, Connection as _, ExpressionMethods, OptionalExtension, PgConnection, QueryDsl,
     RunQueryDsl,
 };
 use graph::{
     constraint_violation,
-    prelude::{info, o, warn, BlockNumber, BlockPtr, Logger, StoreError, ENV_VARS},
+    prelude::{info, o, warn, BlockNumber, BlockPtr, Logger, StoreError},
     schema::EntityType,
 };
 use itertools::Itertools;
@@ -43,17 +39,11 @@ use crate::{
     dynds::DataSourcesTable,
     primary::{DeploymentId, Site},
     relational::index::IndexList,
+    vid_batcher::AdaptiveBatchSize,
 };
 use crate::{connection_pool::ConnectionPool, relational::Layout};
 use crate::{relational::Table, relational_queries as rq};
 
-/// The initial batch size for tables that do not have an array column
-const INITIAL_BATCH_SIZE: i64 = 10_000;
-/// The initial batch size for tables that do have an array column; those
-/// arrays can be large and large arrays will slow down copying a lot. We
-/// therefore tread lightly in that case
-const INITIAL_BATCH_SIZE_LIST: i64 = 100;
-
 const LOG_INTERVAL: Duration = Duration::from_secs(3 * 60);
 
 /// If replicas are lagging by more than this, the copying code will pause
@@ -299,51 +289,6 @@ pub(crate) fn source(
         .map_err(StoreError::from)
 }
 
-/// Track the desired size of a batch in such a way that doing the next
-/// batch gets close to TARGET_DURATION for the time it takes to copy one
-/// batch, but don't step up the size by more than 2x at once
-#[derive(Debug, Queryable)]
-pub(crate) struct AdaptiveBatchSize {
-    pub size: i64,
-}
-
-impl AdaptiveBatchSize {
-    pub fn new(table: &Table) -> Self {
-        let size = if table.columns.iter().any(|col| col.is_list()) {
-            INITIAL_BATCH_SIZE_LIST
-        } else {
-            INITIAL_BATCH_SIZE
-        };
-
-        Self { size }
-    }
-
-    // adjust batch size by trying to extrapolate in such a way that we
-    // get close to TARGET_DURATION for the time it takes to copy one
-    // batch, but don't step up batch_size by more than 2x at once
-    pub fn adapt(&mut self, duration: Duration) {
-        // Avoid division by zero
-        let duration = duration.as_millis().max(1);
-        let new_batch_size = self.size as f64
-            * ENV_VARS.store.batch_target_duration.as_millis() as f64
-            / duration as f64;
-        self.size = (2 * self.size).min(new_batch_size.round() as i64);
-    }
-}
-
-impl ToSql<BigInt, Pg> for AdaptiveBatchSize {
-    fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> diesel::serialize::Result {
-        <i64 as ToSql<BigInt, Pg>>::to_sql(&self.size, out)
-    }
-}
-
-impl FromSql<BigInt, Pg> for AdaptiveBatchSize {
-    fn from_sql(bytes: diesel::pg::PgValue) -> diesel::deserialize::Result<Self> {
-        let size = <i64 as FromSql<BigInt, Pg>>::from_sql(bytes)?;
-        Ok(AdaptiveBatchSize { size })
-    }
-}
-
 /// A helper to copy entities from one table to another in batches that are
 /// small enough to not interfere with the rest of the operations happening
 /// in the database. The `src` and `dst` table must have the same structure

store/postgres/src/lib.rs

@@ -36,6 +36,7 @@ mod store;
 mod store_events;
 mod subgraph_store;
 pub mod transaction_receipt;
+mod vid_batcher;
 mod writable;
 
 pub mod graphman;

store/postgres/src/relational/prune.rs

@@ -18,10 +18,9 @@ use graph::{
 use itertools::Itertools;
 
 use crate::{
-    catalog,
-    copy::AdaptiveBatchSize,
-    deployment,
+    catalog, deployment,
     relational::{Table, VID_COLUMN},
+    vid_batcher::AdaptiveBatchSize,
 };
 
 use super::{Catalog, Layout, Namespace};

store/postgres/src/vid_batcher.rs

@@ -0,0 +1,63 @@
+use std::time::Duration;
+
+use diesel::{
+    deserialize::FromSql,
+    pg::Pg,
+    serialize::{Output, ToSql},
+    sql_types::BigInt,
+};
+use graph::env::ENV_VARS;
+
+use crate::relational::Table;
+
+/// The initial batch size for tables that do not have an array column
+const INITIAL_BATCH_SIZE: i64 = 10_000;
+/// The initial batch size for tables that do have an array column; those
+/// arrays can be large and large arrays will slow down copying a lot. We
+/// therefore tread lightly in that case
+const INITIAL_BATCH_SIZE_LIST: i64 = 100;
+
+/// Track the desired size of a batch in such a way that doing the next
+/// batch gets close to TARGET_DURATION for the time it takes to copy one
+/// batch, but don't step up the size by more than 2x at once
+#[derive(Debug, Queryable)]
+pub(crate) struct AdaptiveBatchSize {
+    pub size: i64,
+}
+
+impl AdaptiveBatchSize {
+    pub fn new(table: &Table) -> Self {
+        let size = if table.columns.iter().any(|col| col.is_list()) {
+            INITIAL_BATCH_SIZE_LIST
+        } else {
+            INITIAL_BATCH_SIZE
+        };
+
+        Self { size }
+    }
+
+    // adjust batch size by trying to extrapolate in such a way that we
+    // get close to TARGET_DURATION for the time it takes to copy one
+    // batch, but don't step up batch_size by more than 2x at once
+    pub fn adapt(&mut self, duration: Duration) {
+        // Avoid division by zero
+        let duration = duration.as_millis().max(1);
+        let new_batch_size = self.size as f64
+            * ENV_VARS.store.batch_target_duration.as_millis() as f64
+            / duration as f64;
+        self.size = (2 * self.size).min(new_batch_size.round() as i64);
+    }
+}
+
+impl ToSql<BigInt, Pg> for AdaptiveBatchSize {
+    fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> diesel::serialize::Result {
+        <i64 as ToSql<BigInt, Pg>>::to_sql(&self.size, out)
+    }
+}
+
+impl FromSql<BigInt, Pg> for AdaptiveBatchSize {
+    fn from_sql(bytes: diesel::pg::PgValue) -> diesel::deserialize::Result<Self> {
+        let size = <i64 as FromSql<BigInt, Pg>>::from_sql(bytes)?;
+        Ok(AdaptiveBatchSize { size })
+    }
+}