Add AbstractMetrics#removeTableMetrics to drop every metric series for a table

pinot-common/src/main/java/org/apache/pinot/common/metrics/AbstractMetrics.java

@@ -19,8 +19,11 @@
 package org.apache.pinot.common.metrics;
 
 import com.google.common.base.Preconditions;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
@@ -49,6 +52,10 @@ public abstract class AbstractMetrics<QP extends AbstractMetrics.QueryPhase, M e
 
   private static final Logger LOGGER = LoggerFactory.getLogger(AbstractMetrics.class);
 
+  /// Name every table's series collapses into when table-level metrics are off. Shared by every table, so a
+  /// table-scoped sweep must never target it.
+  private static final String ALL_TABLES = "allTables";
+
   protected final String _metricPrefix;
 
   protected final PinotMetricsRegistry _metricsRegistry;
@@ -711,6 +718,106 @@ public String composePluginGaugeName(String pluginName, Gauge gauge) {
     return gauge.getGaugeName() + "." + pluginName;
   }
 
+  /// Removes every series this instance registered for the given table.
+  ///
+  /// Unlike the targeted `removeTable*` methods, this does not rebuild names from the rules used to emit them -- it
+  /// scans what is actually registered. That is the whole point. A series emitted with an extra key, or with a
+  /// composite table name, embeds a segment no caller can rediscover from the table name alone, so a sweep built on
+  /// reconstruction strands exactly those series and keeps stranding each new one that gets added.
+  ///
+  /// Matching is deliberately narrow:
+  ///
+  ///   - Only names under this instance's metric prefix are considered, so a table named after a component
+  ///     (`broker`) cannot match the prefix itself.
+  ///   - The table name must occupy whole `.`-delimited segments, never part of one -- `foo` does not match
+  ///     `foobar`, and a database-qualified `db.tbl_OFFLINE` matches only as a unit.
+  ///   - A sibling [AbstractMetrics] sharing this registry and prefix keeps its **gauges**. The ownership check
+  ///     below -- re-deriving the key under this instance's class -- is exact only where the registry key carries
+  ///     the owning class; yammer's does, dropwizard's discards it. What protects the case that actually matters,
+  ///     on every implementation, is the vocabulary check above: a sibling's gauge name is absent from this
+  ///     instance's [#getGauges()], so `<siblingGauge>.<table>` reads as meter-shaped and the table is not at
+  ///     offset 0, so it cannot match. Gauges are the only kind with a re-registration gate ([#_gaugeValues]), so
+  ///     dropping one from under its owner would silence it for the life of the process. A sibling's meter or
+  ///     timer may be dropped early where the key cannot distinguish owners; that is harmless -- they carry no
+  ///     gate and re-register on the next emission. Every instance should still run its own sweep, since that is
+  ///     what clears its own [#_gaugeValues].
+  ///
+  /// A table folded into the shared `allTables` aggregate is safe without a special case: no registered name
+  /// contains its name, so nothing matches. Passing `allTables` itself is rejected for the same reason it would be
+  /// a disaster -- it would delete the aggregate for every table at once.
+  ///
+  /// Two residual false positives are accepted: a workload or remote-cluster name exactly equal to a table name
+  /// sits in the same slot and would be swept. Both re-register on next use, so the cost is one counter reset.
+  ///
+  /// Two things this deliberately does **not** reach, both of which need their owner to clean up:
+  ///
+  ///   - Series a component registers outside any [AbstractMetrics] -- [ValidationMetrics] composes its own
+  ///     `pinot.controller.<table>.<gauge>` names against its own class and keeps its own value map, so the
+  ///     ownership check above skips it. Dropping its registry entries from here would strand that map and retire
+  ///     those gauges for the life of the process.
+  ///   - Names where the table is not followed by a path separator, such as the consumer client id form
+  ///     `<gauge>.<table>-<topic>-<partition>`. Matching those would mean accepting any prefix match, which is
+  ///     what makes `tbl` match `tbl_OFFLINE` and `db.tbl`.
+  ///
+  /// @param tableName the table to sweep, in whichever name form its emitters used (raw or with type)
+  /// @return the number of series removed
+  public int removeTableMetrics(String tableName) {
+    return removeTableMetrics(List.of(tableName));
+  }
+
+  /// Like [#removeTableMetrics(String)], for several tables at once. Prefer this when sweeping a batch: the
+  /// registry is scanned once per call, and yammer and dropwizard both materialise a fresh map on every
+  /// `allMetrics()`.
+  public int removeTableMetrics(Collection<String> tableNames) {
+    Set<String> targets = tableNames.stream().filter(t -> !ALL_TABLES.equals(t)).collect(Collectors.toSet());
+    if (targets.isEmpty()) {
+      return 0;
+    }
+    Set<String> gaugeNames =
+        Arrays.stream(getGauges()).map(Gauge::getGaugeName).collect(Collectors.toCollection(HashSet::new));
+    int removed = 0;
+    // Snapshot the keys before mutating: the compound registry hands back its live map.
+    for (PinotMetricName registeredName : new ArrayList<>(_metricsRegistry.allMetrics().keySet())) {
+      String name = registeredName.getName();
+      if (!name.startsWith(_metricPrefix)
+          || !matchesAnyTable(name.substring(_metricPrefix.length()), targets, gaugeNames)) {
+        continue;
+      }
+      // Re-deriving the key under this class is the ownership test: an identically named series registered by a
+      // sibling AbstractMetrics is a different key, so it compares unequal and is left for that instance to sweep.
+      if (registeredName.equals(PinotMetricUtils.makePinotMetricName(_clazz, name))) {
+        PinotMetricUtils.removeMetric(_metricsRegistry, registeredName);
+        removed++;
+      }
+    }
+    // The deprecated gauge paths gate re-registration on _gaugeValues, so an entry left here would stop a removed
+    // gauge from ever coming back. Swept from this instance's own map rather than from what matched above, so it
+    // stays correct even where the registry cannot tell two instances' series apart.
+    synchronized (_gaugeValues) {
+      _gaugeValues.keySet().removeIf(gaugeName -> matchesAnyTable(gaugeName, targets, gaugeNames));
+    }
+    return removed;
+  }
+
+  /// Whether the prefix-stripped metric name names one of the given tables.
+  ///
+  /// The table sits at exactly one offset, decided by the shape: gauges compose `<gauge>.<table>[.<key>]`, while
+  /// meters, timers and query phases compose `<table>.<rest>`. Which one applies is settled by asking whether the
+  /// leading segment is a known gauge name -- and that question is what keeps a bare `tbl_OFFLINE` from matching
+  /// `db.tbl_OFFLINE`, a genuinely different table whose series must survive. A free search for the name anywhere
+  /// in the string cannot tell those two apart.
+  private static boolean matchesAnyTable(String name, Set<String> tableNames, Set<String> gaugeNames) {
+    int firstDot = name.indexOf('.');
+    int start = firstDot > 0 && gaugeNames.contains(name.substring(0, firstDot)) ? firstDot + 1 : 0;
+    for (String tableName : tableNames) {
+      int end = start + tableName.length();
+      if (name.startsWith(tableName, start) && (end == name.length() || name.charAt(end) == '.')) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   /// Remove gauge from Pinot metrics.
   /// @param gaugeName gauge name
   public void removeGauge(final String gaugeName) {
@@ -739,6 +846,6 @@ private void removeGaugeFromMetricRegistry(String metricName) {
   protected abstract G[] getGauges();
 
   protected String getTableName(String tableName) {
-    return _isTableLevelMetricsEnabled || _allowedTables.contains(tableName) ? tableName : "allTables";
+    return _isTableLevelMetricsEnabled || _allowedTables.contains(tableName) ? tableName : ALL_TABLES;
   }
 }

pinot-common/src/test/java/org/apache/pinot/common/metrics/AbstractMetricsTest.java

@@ -27,6 +27,7 @@
 import java.util.function.IntConsumer;
 import java.util.function.Supplier;
 import java.util.stream.IntStream;
+import org.apache.pinot.common.restlet.resources.RebalanceResult;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.metrics.PinotMeter;
 import org.apache.pinot.spi.metrics.PinotMetricName;
@@ -79,6 +80,107 @@ public void cleanUpMetricsFactory() {
     PinotMetricUtils.cleanUp();
   }
 
+  @Test
+  public void testRemoveTableMetricsReachesNamesNoCallerCanReconstruct() {
+    ControllerMetrics metrics = buildTestMetrics();
+    String table = "myTable_OFFLINE";
+
+    // The four shapes a table series can take. The keyed timer is the one that matters: it is how
+    // tableRebalanceExecutionTimeMs is emitted, and no sweep that rebuilds names from the enums can reach it.
+    metrics.addMeteredTableValue(table, ControllerMeter.LLC_STREAM_DATA_LOSS, 1);
+    metrics.setValueOfTableGauge(table, ControllerGauge.NUMBER_OF_REPLICAS, 3);
+    metrics.setOrUpdateTableGauge(table, "someTenant", ControllerGauge.TABLE_TENANT_INFO, 1);
+    metrics.addTimedTableValue(table, RebalanceResult.Status.DONE.toString(),
+        ControllerTimer.TABLE_REBALANCE_EXECUTION_TIME_MS, 100, TimeUnit.MILLISECONDS);
+    Assert.assertEquals(metrics.getMetricsRegistry().allMetrics().size(), 4);
+
+    Assert.assertEquals(metrics.removeTableMetrics(table), 4);
+    Assert.assertTrue(metrics.getMetricsRegistry().allMetrics().isEmpty());
+
+    // A gauge removed here must be able to come back: _gaugeValues gates re-registration, so a stale entry there
+    // would silently retire the series for the life of the process.
+    metrics.setValueOfTableGauge(table, ControllerGauge.NUMBER_OF_REPLICAS, 5);
+    Assert.assertEquals(getGaugeValue(metrics, ControllerGauge.NUMBER_OF_REPLICAS.getGaugeName() + "." + table), 5);
+  }
+
+  /// Two AbstractMetrics routinely share a registry and a prefix -- OSS and a vendor extension do exactly that.
+  /// A sweep by one must not take the other's gauges with it: gauges are the only kind carrying a
+  /// re-registration gate, so dropping one from under its owner silences it for the life of the process. The
+  /// guarantee comes from the gauge-vocabulary check rather than from key identity, so it holds even on a
+  /// registry whose key does not record the owning class.
+  @Test
+  public void testSweepLeavesASiblingInstancesGaugesAlone() {
+    PinotConfiguration config = new PinotConfiguration();
+    config.setProperty(CONFIG_OF_METRICS_FACTORY_CLASS_NAME, metricsFactoryClassName());
+    PinotMetricUtils.init(config);
+    PinotMetricsRegistry registry = buildRegistry();
+    String table = "shared_OFFLINE";
+
+    ControllerMetrics mine = new ControllerMetrics(registry);
+    ServerMetrics sibling = new ServerMetrics(mine.getMetricPrefix(), registry, true, java.util.Set.of());
+    mine.setValueOfTableGauge(table, ControllerGauge.NUMBER_OF_REPLICAS, 3);
+    sibling.setValueOfTableGauge(table, ServerGauge.LLC_PARTITION_CONSUMING, 1);
+    Assert.assertEquals(registry.allMetrics().size(), 2);
+
+    mine.removeTableMetrics(table);
+
+    // Asserted against the registry rather than through getGaugeValue: that helper derives the prefix from the
+    // instance type, so it would look the sibling up under "pinot.server." and miss the point of the test.
+    Assert.assertEquals(registry.allMetrics().size(), 1, "only the sweeping instance's gauge should be gone");
+    String survivor = mine.getMetricPrefix() + ServerGauge.LLC_PARTITION_CONSUMING.getGaugeName() + "." + table;
+    Assert.assertTrue(registry.allMetrics().keySet().stream().anyMatch(name -> survivor.equals(name.getName())),
+        "a sibling instance's gauge must survive a sweep by an instance sharing its prefix");
+  }
+
+  @Test
+  public void testRemoveTableMetricsMatchesWholeSegmentsOnly() {
+    ControllerMetrics metrics = buildTestMetrics();
+    metrics.addMeteredTableValue("foo_OFFLINE", ControllerMeter.LLC_STREAM_DATA_LOSS, 1);
+    metrics.addMeteredTableValue("foobar_OFFLINE", ControllerMeter.LLC_STREAM_DATA_LOSS, 1);
+    // Database-qualified names span two segments and must still match, but only as a unit.
+    metrics.addMeteredTableValue("db.foo_OFFLINE", ControllerMeter.LLC_STREAM_DATA_LOSS, 1);
+    // Gauges put the table in a different slot, so both shapes need covering.
+    metrics.setValueOfTableGauge("foo_OFFLINE", ControllerGauge.NUMBER_OF_REPLICAS, 1);
+    metrics.setValueOfTableGauge("db.foo_OFFLINE", ControllerGauge.NUMBER_OF_REPLICAS, 1);
+
+    // `db.foo_OFFLINE` is a different table, not a suffix of this one -- in either slot.
+    Assert.assertEquals(metrics.removeTableMetrics("foo_OFFLINE"), 2);
+    Assert.assertEquals(metrics.getMetricsRegistry().allMetrics().size(), 3);
+
+    Assert.assertEquals(metrics.removeTableMetrics("db.foo_OFFLINE"), 2);
+    // foobar survived every sweep: a prefix match is not a segment match.
+    Assert.assertEquals(metrics.getMetricsRegistry().allMetrics().size(), 1);
+  }
+
+  @Test
+  public void testRemoveTableMetricsLeavesGlobalAndOtherTablesAlone() {
+    ControllerMetrics metrics = buildTestMetrics();
+    metrics.addMeteredTableValue("doomed_OFFLINE", ControllerMeter.LLC_STREAM_DATA_LOSS, 1);
+    metrics.addMeteredTableValue("keep_OFFLINE", ControllerMeter.LLC_STREAM_DATA_LOSS, 1);
+    metrics.addMeteredGlobalValue(ControllerMeter.LLC_STREAM_DATA_LOSS, 1);
+    metrics.setValueOfGlobalGauge(ControllerGauge.VERSION, "1.0", 1);
+
+    Assert.assertEquals(metrics.removeTableMetrics("doomed_OFFLINE"), 1);
+    Assert.assertEquals(metrics.getMetricsRegistry().allMetrics().size(), 3);
+  }
+
+  /// With table-level metrics off every table folds into the shared `allTables` series. A sweep must not touch it:
+  /// deleting one table would otherwise zero the aggregate for all of them.
+  @Test
+  public void testRemoveTableMetricsCannotDeleteTheAllTablesAggregate() {
+    PinotConfiguration config = new PinotConfiguration();
+    config.setProperty(CONFIG_OF_METRICS_FACTORY_CLASS_NAME, metricsFactoryClassName());
+    PinotMetricUtils.init(config);
+    ServerMetrics metrics = new ServerMetrics(buildRegistry(), false, java.util.Set.of());
+
+    metrics.addMeteredTableValue("folded_OFFLINE", ServerMeter.QUERIES_ON_TABLE, 1);
+    Assert.assertEquals(metrics.getMetricsRegistry().allMetrics().size(), 1);
+
+    Assert.assertEquals(metrics.removeTableMetrics("folded_OFFLINE"), 0);
+    Assert.assertEquals(metrics.removeTableMetrics("allTables"), 0);
+    Assert.assertEquals(metrics.getMetricsRegistry().allMetrics().size(), 1);
+  }
+
   @Test
   public void testAddOrUpdateGauge() {
     ControllerMetrics controllerMetrics = buildTestMetrics();

pinot-common/src/test/java/org/apache/pinot/plugin/metrics/fake/FakePinotMetricName.java

@@ -23,14 +23,25 @@
 
 
 public class FakePinotMetricName implements PinotMetricName {
+  /// Class-qualified, so that two `AbstractMetrics` sharing a metric prefix stay distinct keys -- mirroring the
+  /// yammer registry, where the owning class is part of the metric identity.
+  private final String _qualifiedName;
   private final String _name;
 
   public FakePinotMetricName(Class<?> clazz, String name) {
-    _name = clazz.getName() + "." + name;
+    _qualifiedName = clazz.getName() + "." + name;
+    _name = name;
   }
 
   @Override
   public Object getMetricName() {
+    return _qualifiedName;
+  }
+
+  /// The bare composed name. Must not include the class qualifier: callers use this to match a registered series
+  /// against a metric prefix, which the qualifier would push out of the way.
+  @Override
+  public String getName() {
     return _name;
   }
 
@@ -42,16 +53,16 @@ public boolean equals(Object o) {
     if (!(o instanceof FakePinotMetricName)) {
       return false;
     }
-    return Objects.equals(_name, ((FakePinotMetricName) o)._name);
+    return Objects.equals(_qualifiedName, ((FakePinotMetricName) o)._qualifiedName);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(_name);
+    return Objects.hashCode(_qualifiedName);
   }
 
   @Override
   public String toString() {
-    return _name;
+    return _qualifiedName;
   }
 }

pinot-plugins/pinot-metrics/pinot-compound-metrics/src/main/java/org/apache/pinot/plugin/metrics/compound/CompoundPinotMetricName.java

@@ -42,6 +42,11 @@ public List<PinotMetricName> getMetricName() {
     return _names;
   }
 
+  @Override
+  public String getName() {
+    return _toString;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

pinot-plugins/pinot-metrics/pinot-dropwizard/src/main/java/org/apache/pinot/plugin/metrics/dropwizard/DropwizardMetricName.java

@@ -38,6 +38,11 @@ public String getMetricName() {
     return _metricName;
   }
 
+  @Override
+  public String getName() {
+    return _metricName;
+  }
+
   /// Overrides equals method by calling the equals from the actual metric name.
   @Override
   public boolean equals(Object obj) {

pinot-plugins/pinot-metrics/pinot-yammer/src/main/java/org/apache/pinot/plugin/metrics/yammer/YammerMetricName.java

@@ -38,6 +38,12 @@ public MetricName getMetricName() {
     return _metricName;
   }
 
+  /// Overridden because [#toString()] here renders the JMX object name, not the bare metric name.
+  @Override
+  public String getName() {
+    return _metricName.getName();
+  }
+
   /// Overrides equals method by calling the equals from the actual metric name.
   @Override
   public boolean equals(Object obj) {

pinot-spi/src/main/java/org/apache/pinot/spi/metrics/PinotMetricName.java

@@ -24,6 +24,19 @@ public interface PinotMetricName {
   /// Returns the actual metric name.
   Object getMetricName();
 
+  /// Returns the composed metric name, free of the registry-specific decoration that some implementations render
+  /// from [#toString()] (yammer renders a JMX object name there, for instance).
+  ///
+  /// This is what lets a caller reason about a series that is *already registered* without reconstructing its name
+  /// from the rules that produced it. Reconstruction cannot reach a name that carries a caller-supplied key, so a
+  /// bulk removal built on it silently strands exactly those series.
+  ///
+  /// The default returns [#toString()], which is correct for implementations whose string form is already the bare
+  /// name; implementations that decorate it must override.
+  default String getName() {
+    return toString();
+  }
+
   /// Overrides the equals method. This is needed as [PinotMetricName] is used as the key of the key-value pair
   /// inside the hashmap in MetricsRegistry. Without overriding equals() and hashCode() methods, all the existing k-v
   /// pairs