store: Count entities during copying

We used to count entities after all the data had been copied, but that is a
very slow operation for large subgraphs; that in turn can lead to bad
side-effects like connection timeouts and the copy ultimately failing.

We now keep track of the number of current entity versions that we copy
with each batch and calculate the entity count incrementally

Author: lutter

store/postgres/src/copy.rs

@@ -33,7 +33,7 @@ use graph::{
 use itertools::Itertools;
 
 use crate::{
-    advisory_lock, catalog,
+    advisory_lock, catalog, deployment,
     dynds::DataSourcesTable,
     primary::{DeploymentId, Site},
     relational::index::IndexList,
@@ -475,12 +475,18 @@ impl TableState {
     }
 
     fn copy_batch(&mut self, conn: &mut PgConnection) -> Result<Status, StoreError> {
-        let (duration, _) = self.batcher.step(|start, end| {
-            rq::CopyEntityBatchQuery::new(self.dst.as_ref(), &self.src, start, end)?
-                .execute(conn)?;
-            Ok(())
+        let (duration, count) = self.batcher.step(|start, end| {
+            let count = rq::CopyEntityBatchQuery::new(self.dst.as_ref(), &self.src, start, end)?
+                .count_current()
+                .get_result::<i64>(conn)
+                .optional()?;
+            Ok(count.unwrap_or(0) as i32)
         })?;
 
+        let count = count.unwrap_or(0);
+
+        deployment::update_entity_count(conn, &self.dst_site, count)?;
+
         self.record_progress(conn, duration)?;
 
         if self.finished() {

store/postgres/src/deployment_store.rs

@@ -1592,13 +1592,10 @@ impl DeploymentStore {
                     .number
                     .checked_add(1)
                     .expect("block numbers fit into an i32");
-                dst.revert_block(conn, block_to_revert)?;
-                info!(logger, "Rewound subgraph to block {}", block.number;
-                      "time_ms" => start.elapsed().as_millis());
+                let (_, count) = dst.revert_block(conn, block_to_revert)?;
+                deployment::update_entity_count(conn, &dst.site, count)?;
 
-                let start = Instant::now();
-                deployment::set_entity_count(conn, &dst.site, &dst.count_query)?;
-                info!(logger, "Counted the entities";
+                info!(logger, "Rewound subgraph to block {}", block.number;
                       "time_ms" => start.elapsed().as_millis());
 
                 deployment::set_history_blocks(

store/postgres/src/relational.rs

@@ -1037,6 +1037,9 @@ impl Layout {
     /// numbers. After this operation, only entity versions inserted or
     /// updated at blocks with numbers strictly lower than `block` will
     /// remain
+    ///
+    /// The `i32` that is returned is the amount by which the entity count
+    /// for the subgraph needs to be adjusted
     pub fn revert_block(
         &self,
         conn: &mut PgConnection,

store/postgres/src/relational_queries.rs

@@ -4799,6 +4799,10 @@ impl<'a> CopyEntityBatchQuery<'a> {
             last_vid,
         })
     }
+
+    pub fn count_current(self) -> CountCurrentVersionsQuery<'a> {
+        CountCurrentVersionsQuery::new(self)
+    }
 }
 
 impl<'a> QueryFragment<Pg> for CopyEntityBatchQuery<'a> {
@@ -4810,6 +4814,8 @@ impl<'a> QueryFragment<Pg> for CopyEntityBatchQuery<'a> {
         // Construct a query
         //   insert into {dst}({columns})
         //   select {columns} from {src}
+        //    where vid >= {first_vid} and vid <= {last_vid}
+        //   returning {upper_inf(block_range)|true}
         out.push_sql("insert into ");
         out.push_sql(self.dst.qualified_name.as_str());
         out.push_sql("(");
@@ -4905,6 +4911,12 @@ impl<'a> QueryFragment<Pg> for CopyEntityBatchQuery<'a> {
         out.push_bind_param::<BigInt, _>(&self.first_vid)?;
         out.push_sql(" and vid <= ");
         out.push_bind_param::<BigInt, _>(&self.last_vid)?;
+        out.push_sql("\n returning ");
+        if self.dst.immutable {
+            out.push_sql("true");
+        } else {
+            out.push_sql(BLOCK_RANGE_CURRENT);
+        }
         Ok(())
     }
 }
@@ -4917,6 +4929,40 @@ impl<'a> QueryId for CopyEntityBatchQuery<'a> {
 
 impl<'a, Conn> RunQueryDsl<Conn> for CopyEntityBatchQuery<'a> {}
 
+#[derive(Debug, Clone)]
+pub struct CountCurrentVersionsQuery<'a> {
+    copy: CopyEntityBatchQuery<'a>,
+}
+
+impl<'a> CountCurrentVersionsQuery<'a> {
+    pub fn new(copy: CopyEntityBatchQuery<'a>) -> Self {
+        Self { copy }
+    }
+}
+impl<'a> QueryFragment<Pg> for CountCurrentVersionsQuery<'a> {
+    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
+        // Generate a query
+        // with copy_cte as ( {copy} )
+        // select count(*) from copy_cte where {block_range_current}
+        out.push_sql("with copy_cte(current) as (");
+        self.copy.walk_ast(out.reborrow())?;
+        out.push_sql(")\nselect count(*) from copy_cte where current");
+        Ok(())
+    }
+}
+
+impl<'a> QueryId for CountCurrentVersionsQuery<'a> {
+    type QueryId = ();
+
+    const HAS_STATIC_QUERY_ID: bool = false;
+}
+
+impl<'a> Query for CountCurrentVersionsQuery<'a> {
+    type SqlType = BigInt;
+}
+
+impl<'a, Conn> RunQueryDsl<Conn> for CountCurrentVersionsQuery<'a> {}
+
 /// Helper struct for returning the id's touched by the RevertRemove and
 /// RevertExtend queries
 #[derive(QueryableByName, PartialEq, Eq, Hash)]