catalog: reclaim storage metadata for ephemeral items on crash restart (SQL-150)

When a process crash-restarted, shards from temp tables were never cleaned up. However, on graceful close, we do clean up these shards. This was an existing bug from before the durable temporary objects change.

Author: SangJunBak

src/catalog/src/durable/transaction.rs

@@ -802,19 +802,90 @@ impl<'a> Transaction<'a> {
         }
     }
 
-    /// Removes every item owned by an ephemeral session from the transaction.
+    /// Removes every item owned by an ephemeral session from the transaction,
+    /// along with the durable state a graceful drop would have removed with
+    /// it: storage collection metadata (moving the backing shards to the
+    /// finalization WAL), comments, and source references.
     ///
     /// Used to reclaim temporary items when the catalog is opened with write
     /// intent, at which point every session that could own one is dead.
+    ///
+    /// This must mirror everything the graceful `Op::DropObjects` path
+    /// persists for a temporary item, because nothing revisits the leftovers:
+    /// bootstrap only ever inserts collection metadata for items present in
+    /// the catalog, and shard finalization is driven solely by the
+    /// `unfinalized_shards` collection, so a metadata row that outlives its
+    /// item leaks the persist shard permanently.
     pub fn remove_ephemeral_items(&mut self) {
-        let keys: Vec<_> = self
-            .items
-            .items()
+        let mut keys = Vec::new();
+        let mut item_ids = BTreeSet::new();
+        let mut global_ids = BTreeSet::new();
+        for (key, value) in self.items.items() {
+            if value.ephemeral_owner_session.is_none() {
+                continue;
+            }
+            item_ids.insert(key.id);
+            global_ids.insert(value.global_id);
+            global_ids.extend(value.extra_versions.values().copied());
+            keys.push(key.clone());
+        }
+        self.items.delete_by_keys(keys, self.op_id);
+
+        // Move the items' storage mappings to the finalization WAL, like
+        // `StorageCollections::prepare_state` does for a graceful drop. Every
+        // version of a table maps to the same shard, and a shard that a
+        // remaining mapping still references must not be finalized. No
+        // remaining mapping can reference one today (only replacement
+        // materialized views share shards, and those cannot be temporary),
+        // so this mirrors `prepare_state`'s guard defensively.
+        let dropped_mappings = self.delete_collection_metadata(global_ids);
+        let mut dropped_shards: BTreeSet<_> = dropped_mappings
             .into_iter()
-            .filter(|(_, value)| value.ephemeral_owner_session.is_some())
-            .map(|(key, _)| key.clone())
+            .map(|(_, shard)| shard)
             .collect();
-        self.items.delete_by_keys(keys, self.op_id);
+        let live_shards: BTreeSet<_> = self.get_collection_metadata().into_values().collect();
+        dropped_shards.retain(|shard| {
+            let live = live_shards.contains(shard);
+            if live {
+                soft_panic_or_log!(
+                    "shard {shard} of a reclaimed ephemeral item is still referenced by a \
+                     live collection, not finalizing it"
+                );
+            }
+            !live
+        });
+        self.insert_unfinalized_shards(dropped_shards).expect(
+            "inserting unfinalized shards only fails on duplicate values, which it ignores",
+        );
+
+        // Comments on ephemeral items would otherwise dangle and, because
+        // item ids are reused, could later re-attach to an unrelated object.
+        self.comments.delete(
+            |key, _value| match key.object_id {
+                CommentObjectId::Table(item_id)
+                | CommentObjectId::View(item_id)
+                | CommentObjectId::MaterializedView(item_id)
+                | CommentObjectId::Source(item_id)
+                | CommentObjectId::Sink(item_id)
+                | CommentObjectId::Index(item_id)
+                | CommentObjectId::Func(item_id)
+                | CommentObjectId::Connection(item_id)
+                | CommentObjectId::Type(item_id)
+                | CommentObjectId::Secret(item_id) => item_ids.contains(&item_id),
+                CommentObjectId::Role(_)
+                | CommentObjectId::Database(_)
+                | CommentObjectId::Schema(_)
+                | CommentObjectId::Cluster(_)
+                | CommentObjectId::ClusterReplica(_)
+                | CommentObjectId::NetworkPolicy(_) => false,
+            },
+            self.op_id,
+        );
+
+        // Only sources hold source references and sources cannot be temporary
+        // today, so this is defensive.
+        self.source_references
+            .delete(|key, _value| item_ids.contains(&key.source_id), self.op_id);
     }
 
     pub fn get_and_increment_id(&mut self, key: String) -> Result<u64, CatalogError> {

src/catalog/tests/read-write.rs

@@ -9,14 +9,14 @@
 
 #![recursion_limit = "256"]
 
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::sync::Arc;
 
 use insta::assert_debug_snapshot;
 use itertools::Itertools;
 use mz_audit_log::{EventDetails, EventType, EventV1, IdNameV1, VersionedEvent};
 use mz_catalog::durable::objects::serialization::proto;
-use mz_catalog::durable::objects::{DurableType, IdAlloc};
+use mz_catalog::durable::objects::{Comment, DurableType, IdAlloc};
 use mz_catalog::durable::{
     CatalogError, Database, DurableCatalogError, FenceError, Item, Metrics,
     TestCatalogStateBuilder, USER_ITEM_ALLOC_KEY, test_bootstrap_args,
@@ -25,12 +25,13 @@ use mz_ore::assert_ok;
 use mz_ore::collections::HashSet;
 use mz_ore::metrics::MetricsRegistry;
 use mz_ore::now::SYSTEM_TIME;
-use mz_persist_client::PersistClient;
+use mz_persist_client::{PersistClient, ShardId};
 use mz_proto::RustType;
 use mz_repr::role_id::RoleId;
-use mz_repr::{CatalogItemId, GlobalId};
+use mz_repr::{CatalogItemId, GlobalId, RelationVersion};
 use mz_sql::catalog::{RoleAttributesRaw, RoleMembership, RoleVars};
-use mz_sql::names::{DatabaseId, ResolvedDatabaseSpecifier, SchemaId};
+use mz_sql::names::{CommentObjectId, DatabaseId, ResolvedDatabaseSpecifier, SchemaId};
+use mz_storage_client::controller::StorageTxn;
 use uuid::Uuid;
 
 #[mz_ore::test(tokio::test)]
@@ -580,6 +581,49 @@ async fn test_ephemeral_items(state_builder: TestCatalogStateBuilder) {
     insert(&mut txn, 200, temp_schema, "tt", Some(session_a)).unwrap();
     insert(&mut txn, 300, temp_schema, "tt", Some(session_b)).unwrap();
 
+    // A temporary item with an ALTER history: two global ids, one shard.
+    txn.insert_item(
+        CatalogItemId::User(500),
+        20_500,
+        GlobalId::User(500),
+        temp_schema,
+        "versioned",
+        "CREATE TABLE versioned (a int)".to_string(),
+        RoleId::User(1),
+        vec![],
+        BTreeMap::from([(RelationVersion::root().bump(), GlobalId::User(501))]),
+        Some(session_a),
+    )
+    .unwrap();
+
+    // Storage mappings like the ones `prepare_state` writes at CREATE, for
+    // the normal item, one plain temporary item, and both versions of the
+    // versioned one.
+    let keep_shard = ShardId::new();
+    let temp_shard = ShardId::new();
+    let versioned_shard = ShardId::new();
+    txn.insert_collection_metadata(BTreeMap::from([
+        (GlobalId::User(100), keep_shard),
+        (GlobalId::User(200), temp_shard),
+        (GlobalId::User(500), versioned_shard),
+        (GlobalId::User(501), versioned_shard),
+    ]))
+    .unwrap();
+
+    // Comments on a temporary and a non-temporary item.
+    txn.update_comment(
+        CommentObjectId::View(CatalogItemId::User(100)),
+        None,
+        Some("keep comment".into()),
+    )
+    .unwrap();
+    txn.update_comment(
+        CommentObjectId::View(CatalogItemId::User(200)),
+        None,
+        Some("temp comment".into()),
+    )
+    .unwrap();
+
     // One session may not hold the same name twice, though.
     let err = insert(&mut txn, 400, temp_schema, "tt", Some(session_a)).unwrap_err();
     assert!(
@@ -627,6 +671,41 @@ async fn test_ephemeral_items(state_builder: TestCatalogStateBuilder) {
         "non-ephemeral item was removed: {snapshot_items:?}"
     );
 
+    // Only the non-ephemeral item's comment survives.
+    let snapshot_comments: Vec<Comment> = state
+        .snapshot()
+        .await
+        .unwrap()
+        .comments
+        .into_iter()
+        .map(RustType::from_proto)
+        .map_ok(|(k, v)| Comment::from_key_value(k, v))
+        .collect::<Result<_, _>>()
+        .unwrap();
+    assert_eq!(
+        snapshot_comments
+            .iter()
+            .map(|c| c.object_id.clone())
+            .collect::<Vec<_>>(),
+        vec![CommentObjectId::View(CatalogItemId::User(100))],
+        "comments on ephemeral items survived: {snapshot_comments:?}"
+    );
+
+    // The ephemeral items' storage mappings moved to the finalization WAL,
+    // deduped to one shard per item. The non-ephemeral mapping is untouched.
+    let txn = state.transaction().await.unwrap();
+    assert_eq!(
+        txn.get_collection_metadata(),
+        BTreeMap::from([(GlobalId::User(100), keep_shard)]),
+        "ephemeral collection metadata survived"
+    );
+    assert_eq!(
+        txn.get_unfinalized_shards(),
+        BTreeSet::from([temp_shard, versioned_shard]),
+        "ephemeral shards were not enqueued for finalization"
+    );
+    drop(txn);
+
     Box::new(state).expire().await;
 }
 

test/restart/mzcompose.py

@@ -1334,6 +1334,36 @@ def wait_for(sql: str, expected: list[tuple], what: str) -> None:
     cur_b.execute("SELECT count(*) FROM tv")
     assert cur_b.fetchall() == [(1,)], "session b lost its own temporary items"
 
+    # A comment on a temporary item is a durable catalog row too, and item ids
+    # are reused, so reclamation must drop it or it can re-attach to an
+    # unrelated later object.
+    cur_b.execute("COMMENT ON TABLE tt IS 'crash victim'")
+    temp_comment_count = """
+        SELECT count(*) FROM mz_internal.mz_catalog_raw
+        WHERE data->>'kind' = 'Comment'
+          AND data->'value'->>'comment' = 'crash victim'
+    """
+    comments = c.sql_query(temp_comment_count, port=6877, user="mz_system")
+    assert comments == [(1,)], f"the temp table's comment was not written: {comments}"
+
+    # Capture the shard backing session b's temp table: the metadata row of
+    # the one remaining ephemeral item that has storage (the temp view has
+    # none). It is what boot-time reclamation must clean up after the kill.
+    shards = c.sql_query(
+        """SELECT m.data->'value'->>'shard'
+           FROM mz_internal.mz_catalog_raw m
+           WHERE m.data->>'kind' = 'StorageCollectionMetadata'
+             AND m.data->'key'->'id' IN (
+               SELECT i.data->'value'->'global_id'
+               FROM mz_internal.mz_catalog_raw i
+               WHERE i.data->>'kind' = 'Item'
+                 AND i.data->'value'->>'ephemeral_owner_session' IS NOT NULL)""",
+        port=6877,
+        user="mz_system",
+    )
+    assert len(shards) == 1, f"expected one ephemeral storage mapping: {shards}"
+    temp_shard = shards[0][0]
+
     # --- kill -9, with session b's items still live ---------------------------
 
     c.kill("materialized")
@@ -1366,6 +1396,39 @@ def wait_for(sql: str, expected: list[tuple], what: str) -> None:
         (0,)
     ], f"ephemeral catalog items survived the restart: {ephemeral}"
 
+    # The temp table's storage mapping must have moved to the finalization
+    # WAL in the same reclamation, else the metadata row and its persist
+    # shard would leak forever. Both rows are stable to assert on here: the
+    # metadata deletion is permanent, and the WAL row survives until the
+    # next committed catalog transaction, which cannot have happened because
+    # nothing has run DDL since the restart.
+    metadata = c.sql_query(
+        f"""SELECT count(*) FROM mz_internal.mz_catalog_raw
+            WHERE data->>'kind' = 'StorageCollectionMetadata'
+              AND data->'value'->>'shard' = '{temp_shard}'""",
+        port=6877,
+        user="mz_system",
+    )
+    assert metadata == [
+        (0,)
+    ], f"temp table's storage metadata survived the restart: {temp_shard}"
+    unfinalized = c.sql_query(
+        f"""SELECT count(*) FROM mz_internal.mz_catalog_raw
+            WHERE data->>'kind' = 'UnfinalizedShard'
+              AND data->'key'->>'shard' = '{temp_shard}'""",
+        port=6877,
+        user="mz_system",
+    )
+    assert unfinalized == [
+        (1,)
+    ], f"temp table's shard was not enqueued for finalization: {temp_shard}"
+
+    # The comment row dies with its item.
+    comments = c.sql_query(temp_comment_count, port=6877, user="mz_system")
+    assert comments == [
+        (0,)
+    ], f"the temp table's comment survived the restart: {comments}"
+
     # conn_b's socket died with the process; closing is bookkeeping only.
     try:
         conn_b.close()