sql-150: reclaim storage metadata for ephemeral items on crash

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

@@ -1391,6 +1391,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")
@@ -1423,6 +1453,36 @@ 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")
+
     # conn_b's socket died with the process; closing is bookkeeping only.
     try:
         conn_b.close()