From d7530053d3f42065b1b635594a734df42af3867b Mon Sep 17 00:00:00 2001
From: Joseph Koshakow <koshy44@gmail.com>
Date: Tue, 12 Sep 2023 12:00:09 -0400
Subject: [PATCH] test: Add coordinator consistency checks

This commit extends the catalog consistency checks by adding a
consistency check for the Coordinator in-memory state. This consistency
check is run as a soft-assert after catalog transactions. However, it's
not run after testdrive tests, or other catalog check locations, since
we don't have easy access to the Coordinator.
---
 src/adapter/src/catalog.rs           |  2 +-
 src/adapter/src/coord.rs             |  1 +
 src/adapter/src/coord/consistency.rs | 82 ++++++++++++++++++++++++++++
 src/adapter/src/coord/ddl.rs         |  2 +-
 4 files changed, 85 insertions(+), 2 deletions(-)
 create mode 100644 src/adapter/src/coord/consistency.rs

diff --git a/src/adapter/src/catalog.rs b/src/adapter/src/catalog.rs
index 2688e37e3def7..edbc0f1ce60e7 100644
--- a/src/adapter/src/catalog.rs
+++ b/src/adapter/src/catalog.rs
@@ -121,11 +121,11 @@ use crate::{AdapterError, AdapterNotice, ExecuteResponse};
 
 mod builtin_table_updates;
 mod config;
-mod consistency;
 mod error;
 mod migrate;
 
 pub mod builtin;
+pub(crate) mod consistency;
 mod inner;
 pub mod storage;
 
diff --git a/src/adapter/src/coord.rs b/src/adapter/src/coord.rs
index c6915ab7753f1..a84e621cd6657 100644
--- a/src/adapter/src/coord.rs
+++ b/src/adapter/src/coord.rs
@@ -159,6 +159,7 @@ pub(crate) mod timestamp_selection;
 
 mod appends;
 mod command_handler;
+mod consistency;
 mod ddl;
 mod indexes;
 mod introspection;
diff --git a/src/adapter/src/coord/consistency.rs b/src/adapter/src/coord/consistency.rs
new file mode 100644
index 0000000000000..faf7fa8e0611c
--- /dev/null
+++ b/src/adapter/src/coord/consistency.rs
@@ -0,0 +1,82 @@
+// Copyright Materialize, Inc. and contributors. All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+//! Internal consistency checks that validate invariants of [`Coordinator`].
+
+use super::Coordinator;
+use crate::catalog::consistency::CatalogInconsistencies;
+use mz_repr::GlobalId;
+use serde::Serialize;
+
+#[derive(Debug, Default, Serialize)]
+pub struct CoordinatorInconsistencies {
+    /// Inconsistencies found in the catalog.
+    catalog_inconsistencies: CatalogInconsistencies,
+    /// Inconsistencies found in read capabilities.
+    read_capabilities: Vec<ReadCapabilitiesInconsistency>,
+}
+
+impl CoordinatorInconsistencies {
+    pub fn is_empty(&self) -> bool {
+        self.catalog_inconsistencies.is_empty() && self.read_capabilities.is_empty()
+    }
+}
+
+impl Coordinator {
+    /// Checks the [`Coordinator`] to make sure we're internally consistent.
+    pub fn check_consistency(&self) -> Result<(), CoordinatorInconsistencies> {
+        let mut inconsistencies = CoordinatorInconsistencies::default();
+
+        if let Err(catalog_inconsistencies) = self.catalog().state().check_consistency() {
+            inconsistencies.catalog_inconsistencies = catalog_inconsistencies;
+        }
+
+        if let Err(read_capabilities) = self.check_read_capabilities() {
+            inconsistencies.read_capabilities = read_capabilities;
+        }
+
+        if inconsistencies.is_empty() {
+            Ok(())
+        } else {
+            Err(inconsistencies)
+        }
+    }
+
+    /// # Invariants:
+    ///
+    /// * Read capabilities should reference known objects.
+    ///
+    fn check_read_capabilities(&self) -> Result<(), Vec<ReadCapabilitiesInconsistency>> {
+        let mut read_capabilities_inconsistencies = Vec::new();
+        for (gid, _) in &self.storage_read_capabilities {
+            if self.catalog().try_get_entry(gid).is_none() {
+                read_capabilities_inconsistencies
+                    .push(ReadCapabilitiesInconsistency::Storage(gid.clone()));
+            }
+        }
+        for (gid, _) in &self.compute_read_capabilities {
+            if self.catalog().try_get_entry(gid).is_none() {
+                read_capabilities_inconsistencies
+                    .push(ReadCapabilitiesInconsistency::Compute(gid.clone()));
+            }
+        }
+
+        if read_capabilities_inconsistencies.is_empty() {
+            Ok(())
+        } else {
+            Err(read_capabilities_inconsistencies)
+        }
+    }
+}
+
+#[derive(Debug, Serialize)]
+enum ReadCapabilitiesInconsistency {
+    Storage(GlobalId),
+    Compute(GlobalId),
+}
diff --git a/src/adapter/src/coord/ddl.rs b/src/adapter/src/coord/ddl.rs
index 465ce8706f6fb..3f571e557f840 100644
--- a/src/adapter/src/coord/ddl.rs
+++ b/src/adapter/src/coord/ddl.rs
@@ -516,7 +516,7 @@ impl Coordinator {
 
         // Note: It's important that we keep the function call inside macro, this way we only run
         // the consistency checks if sort assertions are enabled.
-        mz_ore::soft_assert_eq!(self.catalog().check_consistency(), Ok(()));
+        mz_ore::soft_assert_eq!(self.check_consistency(), Ok(()));
 
         Ok(result)
     }