Merge pull request #29530 from michael-redpanda/sr/contexts/core-15182

[CORE-15182] - Add context labels to schema registry metrics

Author: pgellert

src/v/pandaproxy/schema_registry/store.h

@@ -474,6 +474,9 @@ class store {
             return soft_deleted(sub);
         }
 
+        // Track the subject's deleted status before modification
+        is_deleted was_deleted = sub_it->second.deleted;
+
         sub_it->second.written_at.push_back(marker);
         sub_it->second.deleted = is_deleted::yes;
 
@@ -487,8 +490,21 @@ class store {
         }
 
         if (permanent) {
+            // Permanent delete - decrement appropriate counter and erase
+            auto ctx_it = _context_stores.find(sub.ctx);
+            if (ctx_it != _context_stores.end()) {
+                // Decrement based on current status
+                ctx_it->second.decrement_subject_count(is_deleted::yes);
+            }
             _subjects.erase(sub_it);
         } else {
+            // Soft delete - move from not-deleted to deleted if it wasn't
+            // already deleted
+            if (!was_deleted) {
+                get_or_create_context_store(sub.ctx)
+                  .update_subject_deleted_status(
+                    is_deleted::no, is_deleted::yes);
+            }
             // Mark all versions within the store deleted too: this matters
             // if someone revives the subject with new versions later, as
             // these older versions should remain deleted.
@@ -563,8 +579,9 @@ class store {
     result<bool>
     set_mode(seq_marker marker, const context& ctx, mode m, force f) {
         BOOST_OUTCOME_TRYX(check_mode_mutability(f));
-        _context_stores[ctx]._mode_written_at.emplace_back(marker);
-        return std::exchange(_context_stores[ctx]._mode, m) != m;
+        auto& context = get_or_create_context_store(ctx);
+        context._mode_written_at.emplace_back(marker);
+        return std::exchange(context._mode, m) != m;
     }
 
     ///\brief Set the mode for a subject.
@@ -590,9 +607,9 @@ class store {
     ///\brief Clear the mode for a context.
     result<bool> clear_mode(const context& ctx, force f) {
         BOOST_OUTCOME_TRYX(check_mode_mutability(f));
-        _context_stores[ctx]._mode_written_at.clear();
-        return std::exchange(_context_stores[ctx]._mode, std::nullopt)
-               != std::nullopt;
+        auto& context = get_or_create_context_store(ctx);
+        context._mode_written_at.clear();
+        return std::exchange(context._mode, std::nullopt) != std::nullopt;
     }
 
     /// \brief Return the seq_marker write history of a context, but only
@@ -639,8 +656,9 @@ class store {
       seq_marker marker,
       const context& ctx,
       compatibility_level compatibility) {
-        _context_stores[ctx]._config_written_at.push_back(marker);
-        return std::exchange(_context_stores[ctx]._compatibility, compatibility)
+        auto& context = get_or_create_context_store(ctx);
+        context._config_written_at.push_back(marker);
+        return std::exchange(context._compatibility, compatibility)
                != compatibility;
     }
 
@@ -657,8 +675,9 @@ class store {
 
     ///\brief Clear the compatibility level of a context.
     result<bool> clear_compatibility(const context& ctx) {
-        _context_stores[ctx]._config_written_at.clear();
-        return std::exchange(_context_stores[ctx]._compatibility, std::nullopt)
+        auto& context = get_or_create_context_store(ctx);
+        context._config_written_at.clear();
+        return std::exchange(context._compatibility, std::nullopt)
                != std::nullopt;
     }
 
@@ -699,6 +718,11 @@ class store {
                           : std::prev(_schemas.end())->first.id + 1;
         auto [_, inserted] = _schemas.try_emplace(
           context_schema_id{ctx, id}, std::move(def));
+
+        if (inserted) {
+            get_or_create_context_store(ctx).increment_schema_count();
+        }
+
         return {id, inserted};
     }
 
@@ -707,12 +731,26 @@ class store {
         if (mark_schema) {
             _marked_schemas.push_back(id);
         }
-        return _schemas
-          .insert_or_assign(std::move(id), schema_entry(std::move(def)))
-          .second;
+        auto [it, inserted] = _schemas.insert_or_assign(
+          std::move(id), schema_entry(std::move(def)));
+
+        if (inserted) {
+            get_or_create_context_store(it->first.ctx).increment_schema_count();
+        }
+
+        return inserted;
     }
 
-    void delete_schema(const context_schema_id& id) { _schemas.erase(id); }
+    void delete_schema(const context_schema_id& id) {
+        auto it = _schemas.find(id);
+        if (it != _schemas.end()) {
+            auto ctx_it = _context_stores.find(id.ctx);
+            if (ctx_it != _context_stores.end()) {
+                ctx_it->second.decrement_schema_count();
+            }
+            _schemas.erase(it);
+        }
+    }
 
     // This function returns and unmarkes all marked schemas.
     chunked_vector<context_schema_id> extract_marked_schemas() {
@@ -724,16 +762,32 @@ class store {
         bool inserted;
     };
     insert_subject_result insert_subject(context_subject sub, schema_id id) {
-        auto& subject_entry = get_or_create_subject_entry(std::move(sub));
+        auto [it, inserted] = _subjects.try_emplace(sub, sub);
+        auto& subject_entry = it->second;
+
+        // Track the previous deleted status
+        is_deleted was_deleted = subject_entry.deleted;
         subject_entry.deleted = is_deleted::no;
+
+        // Update counters based on whether this is a new subject or revival
+        if (inserted) {
+            // New subject - increment not-deleted counter
+            get_or_create_context_store(sub.ctx).increment_subject_count(
+              is_deleted::no);
+        } else if (was_deleted) {
+            // Reviving a deleted subject - move from deleted to not-deleted
+            get_or_create_context_store(sub.ctx).update_subject_deleted_status(
+              is_deleted::yes, is_deleted::no);
+        }
+
         auto& versions = subject_entry.versions;
         const auto v_it = std::find_if(
           versions.begin(), versions.end(), [id](auto v) {
               return v.id == id;
           });
         if (v_it != versions.cend()) {
-            auto inserted = std::exchange(v_it->deleted, is_deleted::no);
-            return {v_it->version, bool(inserted)};
+            auto was_deleted = std::exchange(v_it->deleted, is_deleted::no);
+            return {v_it->version, bool(was_deleted)};
         }
 
         const auto version = versions.empty() ? schema_version{1}
@@ -748,7 +802,12 @@ class store {
       schema_version version,
       schema_id id,
       is_deleted deleted) {
-        auto& subject_entry = get_or_create_subject_entry(std::move(sub));
+        auto [it, inserted] = _subjects.try_emplace(sub, sub);
+        auto& subject_entry = it->second;
+
+        // Track if subject deletion status changed
+        is_deleted old_deleted = subject_entry.deleted;
+
         auto& versions = subject_entry.versions;
         subject_entry.written_at.push_back(marker);
 
@@ -784,6 +843,17 @@ class store {
             subject_entry.deleted = deleted;
         }
 
+        // Update counters based on whether this is new or status changed
+        if (inserted) {
+            // New subject - increment appropriate counter
+            get_or_create_context_store(sub.ctx).increment_subject_count(
+              deleted);
+        } else if (old_deleted != subject_entry.deleted) {
+            // Deletion status changed - move between counters
+            get_or_create_context_store(sub.ctx).update_subject_deleted_status(
+              old_deleted, subject_entry.deleted);
+        }
+
         return !found;
     }
 
@@ -799,24 +869,6 @@ class store {
 
     void setup_metrics() {
         namespace sm = ss::metrics;
-        const auto make_schema_count = [this]() {
-            return sm::make_gauge(
-              "schema_count",
-              [this] { return _schemas.size(); },
-              sm::description("The number of schemas in the store"));
-        };
-        const auto make_subject_count = [this](is_deleted deleted) {
-            return sm::make_gauge(
-              "subject_count",
-              [this, deleted] {
-                  return std::ranges::count_if(
-                    _subjects, [deleted](const auto& entry) {
-                        return entry.second.deleted == deleted;
-                    });
-              },
-              sm::description("The number of subjects in the store"),
-              {sm::label{"deleted"}(deleted)});
-        };
         const auto make_schema_bytes = [this]() {
             return sm::make_gauge(
               "schema_memory_bytes",
@@ -837,10 +889,7 @@ class store {
             _metrics.add_group(
               group_name,
               {
-                make_schema_count(),
                 make_schema_bytes(),
-                make_subject_count(is_deleted::no),
-                make_subject_count(is_deleted::yes),
               },
               {},
               agg);
@@ -850,16 +899,13 @@ class store {
             _public_metrics.add_group(
               group_name,
               {
-                make_schema_count().aggregate(agg),
                 make_schema_bytes().aggregate(agg),
-                make_subject_count(is_deleted::no).aggregate(agg),
-                make_subject_count(is_deleted::yes).aggregate(agg),
               });
         }
     };
 
     void maybe_update_max_schema_id(const context_schema_id& id) {
-        auto& nsi = _context_stores[id.ctx]._next_schema_id;
+        auto& nsi = get_or_create_context_store(id.ctx)._next_schema_id;
         auto old = nsi;
         nsi = std::max(nsi, id.id + schema_id{1});
         vlog(
@@ -873,7 +919,7 @@ class store {
         // operations if needed.  _next_schema_id gets updated
         // if the operation was successful, as a side effect
         // of applying the write to the store.
-        return _context_stores[ctx]._next_schema_id;
+        return get_or_create_context_store(ctx)._next_schema_id;
     }
 
     chunked_vector<context> get_materialized_contexts() const {
@@ -896,7 +942,7 @@ class store {
     }
 
     void set_context_materialized(const context& ctx, bool materialized) {
-        _context_stores[ctx]._materialized = materialized;
+        get_or_create_context_store(ctx)._materialized = materialized;
     }
 
 private:
@@ -1030,11 +1076,117 @@ class store {
         schema_id _next_schema_id{1};
         bool _materialized{false};
 
+        void increment_schema_count() { _schema_count++; }
+
+        void decrement_schema_count() {
+            if (_schema_count > 0) {
+                _schema_count--;
+            }
+        }
+
+        void clear_schema_count() { _schema_count = 0; }
+
+        void increment_subject_count(is_deleted deleted) {
+            if (deleted == is_deleted::yes) {
+                _subject_count_deleted++;
+            } else {
+                _subject_count_not_deleted++;
+            }
+        }
+
+        void decrement_subject_count(is_deleted deleted) {
+            if (deleted == is_deleted::yes) {
+                if (_subject_count_deleted > 0) {
+                    _subject_count_deleted--;
+                }
+            } else {
+                if (_subject_count_not_deleted > 0) {
+                    _subject_count_not_deleted--;
+                }
+            }
+        }
+
+        void update_subject_deleted_status(
+          is_deleted old_status, is_deleted new_status) {
+            if (old_status != new_status) {
+                decrement_subject_count(old_status);
+                increment_subject_count(new_status);
+            }
+        }
+
+        void clear_subject_counts() {
+            _subject_count_deleted = 0;
+            _subject_count_not_deleted = 0;
+        }
+
         chunked_vector<seq_marker> _config_written_at;
         chunked_vector<seq_marker> _mode_written_at;
+
+    private:
+        metrics::internal_metric_groups _metrics;
+        metrics::public_metric_groups _public_metrics;
+        size_t _schema_count{0};
+        size_t _subject_count_not_deleted{0};
+        size_t _subject_count_deleted{0};
+
+        void setup_metrics(const context& ctx) {
+            namespace sm = ss::metrics;
+            auto group_name = prometheus_sanitize::metrics_name(
+              "schema_registry_cache");
+
+            const auto make_schema_count = [this, &ctx]() {
+                return sm::make_gauge(
+                  "schema_count",
+                  [this] { return _schema_count; },
+                  sm::description("The number of schemas in the store"),
+                  {sm::label{"context"}(ctx)});
+            };
+
+            const auto make_subject_count = [this, &ctx](is_deleted deleted) {
+                return sm::make_gauge(
+                  "subject_count",
+                  [this, deleted] {
+                      return deleted == is_deleted::yes
+                               ? _subject_count_deleted
+                               : _subject_count_not_deleted;
+                  },
+                  sm::description("The number of subjects in the store"),
+                  {sm::label{"context"}(ctx), sm::label{"deleted"}(deleted)});
+            };
+
+            if (!config::shard_local_cfg().disable_metrics()) {
+                _metrics.add_group(
+                  group_name,
+                  {make_schema_count(),
+                   make_subject_count(is_deleted::no),
+                   make_subject_count(is_deleted::yes)},
+                  {},
+                  {sm::shard_label});
+            }
+
+            if (!config::shard_local_cfg().disable_public_metrics()) {
+                _public_metrics.add_group(
+                  group_name,
+                  {make_schema_count().aggregate({sm::shard_label}),
+                   make_subject_count(is_deleted::no)
+                     .aggregate({sm::shard_label}),
+                   make_subject_count(is_deleted::yes)
+                     .aggregate({sm::shard_label})});
+            }
+        }
+
+        friend class store;
     };
     using context_store_map = absl::node_hash_map<context, context_store>;
 
+    context_store& get_or_create_context_store(const context& ctx) {
+        auto [it, inserted] = _context_stores.try_emplace(ctx);
+        if (inserted) {
+            it->second.setup_metrics(ctx);
+        }
+        return it->second;
+    }
+
     // NOTE: sharded_store shards data into multiple store instances, so some
     // fields are only present on certain shards.
     // _schemas: sharded by (context, schema_id)