[controller][server][da-vinci][protocol][doc] Fix version lifecycle leaks

clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/AggVersionedStorageEngineStats.java

@@ -93,6 +93,24 @@ public void setStorageEngine(String topicName, StorageEngine storageEngine) {
     }
   }
 
+  public void unsetStorageEngine(String topicName) {
+    if (!Version.isVersionTopicOrStreamReprocessingTopic(topicName)) {
+      LOGGER.warn("Invalid topic name: {}", topicName);
+      return;
+    }
+    String storeName = Version.parseStoreFromKafkaTopicName(topicName);
+    int version = Version.parseVersionFromKafkaTopicName(topicName);
+    try {
+      getStats(storeName, version).setStorageEngine(null);
+      otelStatsMap.computeIfPresent(storeName, (k, stats) -> {
+        stats.onVersionRemoved(version);
+        return stats;
+      });
+    } catch (Exception e) {
+      LOGGER.warn("Failed to unset StorageEngine for store: {}, version: {}", storeName, version, e);
+    }
+  }
+
   public void recordRocksDBOpenFailure(String topicName) {
     if (!Version.isVersionTopicOrStreamReprocessingTopic(topicName)) {
       LOGGER.warn("Invalid topic name: {}", topicName);

clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/StorageEngineOtelMetricEntity.java

@@ -17,21 +17,28 @@
 /**
  * OTel metric entities for storage engine statistics.
  *
- * <p>Consolidates 4 Tehuti AsyncGauge sensors into 3 OTel metrics:
+ * <p>Consolidates storage-engine Tehuti sensors and adds bounded role-level visibility:
  * <ul>
  *   <li>{@code disk_usage_in_bytes} + {@code rmd_disk_usage_in_bytes} consolidated into
  *       {@link #DISK_USAGE} with {@code RECORD_TYPE} dimension (DATA vs REPLICATION_METADATA)</li>
+ *   <li>{@link #VERSION_COUNT} reports how many local storage engines have each version role</li>
  *   <li>{@code rocksdb_open_failure_count} maps to {@link #ROCKSDB_OPEN_FAILURE_COUNT} (COUNTER)</li>
  *   <li>{@code rocksdb_key_count_estimate} maps to {@link #KEY_COUNT_ESTIMATE} (ASYNC_GAUGE)</li>
  * </ul>
  */
 public enum StorageEngineOtelMetricEntity implements ModuleMetricEntityInterface {
   DISK_USAGE(
       "ingestion.disk.used", MetricType.ASYNC_GAUGE, MetricUnit.BYTES,
-      "Disk usage in bytes by record type (data or replication metadata)",
+      "Total disk usage in bytes across all versions with each role, by record type",
       setOf(VENICE_CLUSTER_NAME, VENICE_STORE_NAME, VENICE_VERSION_ROLE, VENICE_RECORD_TYPE)
   ),
 
+  VERSION_COUNT(
+      "ingestion.disk.version_count", MetricType.ASYNC_GAUGE, MetricUnit.NUMBER,
+      "Number of locally loaded storage-engine versions by role",
+      setOf(VENICE_CLUSTER_NAME, VENICE_STORE_NAME, VENICE_VERSION_ROLE)
+  ),
+
   ROCKSDB_OPEN_FAILURE_COUNT(
       "rocksdb.open.failure_count", MetricType.COUNTER, MetricUnit.NUMBER,
       "Count of RocksDB open failures; VERSION_ROLE reflects the version's role at failure time, not its current role",

clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/StorageEngineOtelStats.java

@@ -5,6 +5,7 @@
 import static com.linkedin.davinci.stats.StorageEngineOtelMetricEntity.DISK_USAGE;
 import static com.linkedin.davinci.stats.StorageEngineOtelMetricEntity.KEY_COUNT_ESTIMATE;
 import static com.linkedin.davinci.stats.StorageEngineOtelMetricEntity.ROCKSDB_OPEN_FAILURE_COUNT;
+import static com.linkedin.davinci.stats.StorageEngineOtelMetricEntity.VERSION_COUNT;
 import static com.linkedin.venice.meta.Store.NON_EXISTING_VERSION;
 
 import com.linkedin.davinci.stats.AggVersionedStorageEngineStats.StorageEngineStatsWrapper;
@@ -26,10 +27,11 @@
 /**
  * Per-store OTel stats for storage engine metrics.
  *
- * <p>Holds 3 OTel metrics:
+ * <p>Holds 4 OTel metrics:
  * <ul>
  *   <li>{@code ingestion.disk.used} — ASYNC_GAUGE with VERSION_ROLE + RECORD_TYPE dimensions,
  *       implemented via {@link AsyncMetricEntityStateTwoEnums}{@code <VeniceRecordType, VersionRole>}.</li>
+ *   <li>{@code ingestion.disk.version_count} — ASYNC_GAUGE with VERSION_ROLE dimension</li>
  *   <li>{@code rocksdb.key.estimated_count} — ASYNC_GAUGE with VERSION_ROLE dimension</li>
  *   <li>{@code rocksdb.open.failure_count} — COUNTER with VERSION_ROLE dimension</li>
  * </ul>
@@ -61,6 +63,9 @@ public class StorageEngineOtelStats implements Closeable {
   /** Key count ASYNC_GAUGE with VersionRole dimension */
   private final AsyncMetricEntityStateOneEnum<VersionRole> keyCountMetric;
 
+  /** Number of local storage engines with each VersionRole */
+  private final AsyncMetricEntityStateOneEnum<VersionRole> versionCountMetric;
+
   /** RocksDB open failure COUNTER with VersionRole dimension */
   private final MetricEntityStateOneEnum<VersionRole> openFailureMetric;
 
@@ -78,18 +83,17 @@ public StorageEngineOtelStats(MetricsRepository metricsRepository, String storeN
       Map<VeniceMetricsDimensions, String> baseDimensionsMap = otelSetup.getBaseDimensionsMap();
 
       /*
-       * Two-callback contract: the liveStateResolver returns the wrapper or null (null -> dormant,
-       * no emission); the valueResolver reads the metric value from that wrapper. The null return
-       * is the liveness signal, enforced by the API.
+       * The live-state resolver returns the bounded per-version map when a role is present, or null
+       * when dormant. The value resolver aggregates that role without creating per-version attributes.
        */
       this.diskUsageMetrics = AsyncMetricEntityStateTwoEnums.create(
           DISK_USAGE.getMetricEntity(),
           otelRepository,
           baseDimensionsMap,
           VeniceRecordType.class,
           VersionRole.class,
-          (recordType, role) -> getWrapperForRole(role),
-          (wrapper, recordType, role) -> diskUsage(wrapper, recordType));
+          (recordType, role) -> hasVersionForRole(role) ? wrappersByVersion : null,
+          (wrappers, recordType, role) -> diskUsageForRole(wrappers, role, recordType));
 
       this.keyCountMetric = AsyncMetricEntityStateOneEnum.create(
           KEY_COUNT_ESTIMATE.getMetricEntity(),
@@ -99,12 +103,21 @@ public StorageEngineOtelStats(MetricsRepository metricsRepository, String storeN
           role -> getWrapperForRole(role),
           (wrapper, role) -> wrapper.getKeyCountEstimate());
 
+      this.versionCountMetric = AsyncMetricEntityStateOneEnum.create(
+          VERSION_COUNT.getMetricEntity(),
+          otelRepository,
+          baseDimensionsMap,
+          VersionRole.class,
+          role -> wrappersByVersion.isEmpty() ? null : wrappersByVersion,
+          (wrappers, role) -> countVersionsForRole(wrappers, role));
+
       // RocksDB open failure count: COUNTER with VersionRole dimension
       this.openFailureMetric = MetricEntityStateOneEnum
           .create(ROCKSDB_OPEN_FAILURE_COUNT.getMetricEntity(), otelRepository, baseDimensionsMap, VersionRole.class);
     } else {
       this.diskUsageMetrics = null;
       this.keyCountMetric = null;
+      this.versionCountMetric = null;
       this.openFailureMetric = null;
     }
   }
@@ -175,6 +188,33 @@ private StorageEngineStatsWrapper getWrapperForRole(VersionRole role) {
     return wrappersByVersion.get(version);
   }
 
+  private boolean hasVersionForRole(VersionRole role) {
+    return countVersionsForRole(wrappersByVersion, role) > 0;
+  }
+
+  private long countVersionsForRole(Map<Integer, StorageEngineStatsWrapper> wrappers, VersionRole role) {
+    VersionInfo snapshot = versionInfo;
+    return wrappers.keySet().stream().filter(version -> classifyVersion(version, snapshot) == role).count();
+  }
+
+  private long diskUsageForRole(
+      Map<Integer, StorageEngineStatsWrapper> wrappers,
+      VersionRole role,
+      VeniceRecordType recordType) {
+    VersionInfo snapshot = versionInfo;
+    if (role == VersionRole.BACKUP) {
+      return wrappers.entrySet()
+          .stream()
+          .filter(entry -> classifyVersion(entry.getKey(), snapshot) == VersionRole.BACKUP)
+          .mapToLong(entry -> diskUsage(entry.getValue(), recordType))
+          .sum();
+    }
+
+    int version = getVersionForRole(role, snapshot, wrappers.keySet());
+    StorageEngineStatsWrapper wrapper = wrappers.get(version);
+    return wrapper == null ? 0 : diskUsage(wrapper, recordType);
+  }
+
   /** Reads disk usage (data or RMD) from a resolved wrapper. */
   private static long diskUsage(StorageEngineStatsWrapper wrapper, VeniceRecordType recordType) {
     switch (recordType) {
@@ -187,14 +227,15 @@ private static long diskUsage(StorageEngineStatsWrapper wrapper, VeniceRecordTyp
     }
   }
 
-  /**
-   * Clears internal wrapper references. On subsequent collections each async-gauge's
-   * {@code liveStateResolver} will return {@code null} for every role and no data points will be
-   * emitted. The SDK instruments themselves are NOT deregistered — they remain registered and
-   * are polled until the SDK is shut down.
-   */
+  /** Stops observable callbacks and clears all storage-engine references. */
   @Override
   public void close() {
     wrappersByVersion.clear();
+    versionInfo = VersionInfo.NON_EXISTING;
+    if (emitOtelMetrics) {
+      diskUsageMetrics.close();
+      keyCountMetric.close();
+      versionCountMetric.close();
+    }
   }
 }

clients/da-vinci-client/src/main/java/com/linkedin/davinci/storage/StorageService.java

@@ -526,6 +526,7 @@ public synchronized void removeStorageEngine(String kafkaTopic) {
       LOGGER.warn("Storage engine {} does not exist, ignoring remove request.", kafkaTopic);
       return;
     }
+    aggVersionedStorageEngineStats.unsetStorageEngine(kafkaTopic);
     storageEngine.drop();
 
     VeniceStoreVersionConfig storeConfig = configLoader.getStoreConfig(kafkaTopic);
@@ -552,6 +553,7 @@ public synchronized void closeStorageEngine(String kafkaTopic) {
       LOGGER.warn("Storage engine {} does not exist, ignoring close request.", kafkaTopic);
       return;
     }
+    aggVersionedStorageEngineStats.unsetStorageEngine(kafkaTopic);
     storageEngine.close();
 
     VeniceStoreVersionConfig storeConfig = configLoader.getStoreConfig(kafkaTopic);

clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/AggVersionedStorageEngineStatsTest.java

@@ -84,6 +84,31 @@ public void testSetStorageEngineWiresOtelStats() {
     }
   }
 
+  @Test
+  public void testUnsetStorageEngineRemovesAndReopenRestoresOtelStats() {
+    InMemoryMetricReader reader = InMemoryMetricReader.create();
+    try (VeniceMetricsRepository veniceRepo = createOtelMetricsRepository(reader)) {
+      OtelTestContext ctx = createOtelTestContext(veniceRepo);
+      StorageEngineStats mockEngineStats = mock(StorageEngineStats.class);
+      doReturn(5000L).when(mockEngineStats).getStoreSizeInBytes();
+      StorageEngine mockEngine = mock(StorageEngine.class);
+      doReturn(mockEngineStats).when(mockEngine).getStats();
+      Attributes attrs = buildDiskUsageDataAttrs(ctx.clusterName, ctx.storeName, VersionRole.BACKUP);
+      String diskMetric = StorageEngineOtelMetricEntity.DISK_USAGE.getMetricEntity().getMetricName();
+
+      ctx.stats.setStorageEngine(ctx.topicName, mockEngine);
+      OpenTelemetryDataTestUtils.validateLongPointDataFromGauge(reader, 5000, attrs, diskMetric, OTEL_PREFIX);
+
+      ctx.stats.unsetStorageEngine(ctx.topicName);
+      LongPointData pointAfterClose = OpenTelemetryDataTestUtils
+          .getLongPointDataFromGaugeIfPresent(reader.collectAllMetrics(), diskMetric, OTEL_PREFIX, attrs);
+      Assert.assertNull(pointAfterClose, "Closed local engine must not emit a version metric");
+
+      ctx.stats.setStorageEngine(ctx.topicName, mockEngine);
+      OpenTelemetryDataTestUtils.validateLongPointDataFromGauge(reader, 5000, attrs, diskMetric, OTEL_PREFIX);
+    }
+  }
+
   @Test
   public void testRecordRocksDBOpenFailureWiresBothTehutiAndOtel() {
     InMemoryMetricReader reader = InMemoryMetricReader.create();

clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/StorageEngineOtelMetricEntityTest.java

@@ -29,8 +29,16 @@ private static Map<StorageEngineOtelMetricEntity, MetricEntityExpectation> expec
             "ingestion.disk.used",
             MetricType.ASYNC_GAUGE,
             MetricUnit.BYTES,
-            "Disk usage in bytes by record type (data or replication metadata)",
+            "Total disk usage in bytes across all versions with each role, by record type",
             setOf(VENICE_CLUSTER_NAME, VENICE_STORE_NAME, VENICE_VERSION_ROLE, VENICE_RECORD_TYPE)));
+    map.put(
+        StorageEngineOtelMetricEntity.VERSION_COUNT,
+        new MetricEntityExpectation(
+            "ingestion.disk.version_count",
+            MetricType.ASYNC_GAUGE,
+            MetricUnit.NUMBER,
+            "Number of locally loaded storage-engine versions by role",
+            setOf(VENICE_CLUSTER_NAME, VENICE_STORE_NAME, VENICE_VERSION_ROLE)));
     map.put(
         StorageEngineOtelMetricEntity.ROCKSDB_OPEN_FAILURE_COUNT,
         new MetricEntityExpectation(

clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/StorageEngineOtelStatsTest.java

@@ -32,6 +32,8 @@ public class StorageEngineOtelStatsTest {
 
   private static final String DISK_USAGE_METRIC =
       StorageEngineOtelMetricEntity.DISK_USAGE.getMetricEntity().getMetricName();
+  private static final String VERSION_COUNT_METRIC =
+      StorageEngineOtelMetricEntity.VERSION_COUNT.getMetricEntity().getMetricName();
   private static final String KEY_COUNT_METRIC =
       StorageEngineOtelMetricEntity.KEY_COUNT_ESTIMATE.getMetricEntity().getMetricName();
   private static final String OPEN_FAILURE_METRIC =
@@ -144,17 +146,23 @@ public void testDiskUsageRmdBackupVersion() {
   }
 
   @Test
-  public void testDiskUsageBackupSelectsSmallestVersion() {
-    // current=1, future=2. Versions 3 and 5 are both backups — should select version 3 (smallest)
+  public void testDiskUsageAggregatesAllBackupVersions() {
+    // current=1, future=2. Versions 3 and 5 are both backups.
     stats.setStatsWrapper(3, new MockWrapper(3000, 300, 30));
     stats.setStatsWrapper(5, new MockWrapper(5000, 500, 50));
 
     OpenTelemetryDataTestUtils.validateLongPointDataFromGauge(
         inMemoryMetricReader,
-        3000,
+        8000,
         buildDiskUsageAttributes(VersionRole.BACKUP, VeniceRecordType.DATA),
         DISK_USAGE_METRIC,
         METRIC_PREFIX);
+    OpenTelemetryDataTestUtils.validateLongPointDataFromGauge(
+        inMemoryMetricReader,
+        800,
+        buildDiskUsageAttributes(VersionRole.BACKUP, VeniceRecordType.REPLICATION_METADATA),
+        DISK_USAGE_METRIC,
+        METRIC_PREFIX);
   }
 
   @Test
@@ -326,6 +334,29 @@ public void testLiveValueUpdatesAfterVersionInfoChange() {
         METRIC_PREFIX);
   }
 
+  @Test
+  public void testVersionCountTracksAddRemoveAndRoleChanges() {
+    stats.setStatsWrapper(1, new MockWrapper(1000, 100, 10));
+    stats.setStatsWrapper(2, new MockWrapper(2000, 200, 20));
+    stats.setStatsWrapper(3, new MockWrapper(3000, 300, 30));
+    stats.setStatsWrapper(5, new MockWrapper(5000, 500, 50));
+
+    assertVersionCount(VersionRole.CURRENT, 1);
+    assertVersionCount(VersionRole.FUTURE, 1);
+    assertVersionCount(VersionRole.BACKUP, 2);
+
+    stats.onVersionRemoved(3);
+    assertVersionCount(VersionRole.BACKUP, 1);
+
+    stats.updateVersionInfo(5, 2);
+    assertVersionCount(VersionRole.CURRENT, 1);
+    assertVersionCount(VersionRole.FUTURE, 1);
+    assertVersionCount(VersionRole.BACKUP, 1);
+
+    stats.onVersionRemoved(1);
+    assertVersionCount(VersionRole.BACKUP, 0);
+  }
+
   @Test
   public void testRemoveVersionClearsWrapper() {
     AggVersionedStorageEngineStats.StorageEngineStatsWrapper wrapper = new MockWrapper(1000, 100, 10);
@@ -394,6 +425,15 @@ private static Attributes buildVersionRoleAttributes(VersionRole role) {
         .build();
   }
 
+  private void assertVersionCount(VersionRole role, long expectedCount) {
+    OpenTelemetryDataTestUtils.validateLongPointDataFromGauge(
+        inMemoryMetricReader,
+        expectedCount,
+        buildVersionRoleAttributes(role),
+        VERSION_COUNT_METRIC,
+        METRIC_PREFIX);
+  }
+
   @Test
   public void testVersionRoleEnumCount() {
     // getVersionForRole returns NON_EXISTING_VERSION for unknown roles — but a new VersionRole

clients/da-vinci-client/src/test/java/com/linkedin/davinci/storage/StorageServiceTest.java

@@ -16,6 +16,7 @@
 import com.linkedin.davinci.config.VeniceStoreVersionConfig;
 import com.linkedin.davinci.stats.AggVersionedStorageEngineStats;
 import com.linkedin.davinci.stats.RocksDBMemoryStats;
+import com.linkedin.davinci.store.DelegatingStorageEngine;
 import com.linkedin.davinci.store.StorageEngine;
 import com.linkedin.davinci.store.StorageEngineFactory;
 import com.linkedin.venice.exceptions.VeniceNoStoreException;
@@ -76,6 +77,47 @@ public void testDeleteStorageEngineOnRocksDBError() {
     verify(factory, times(2)).removeStorageEngine(storageEngineName);
   }
 
+  @Test
+  public void testRemoveAndCloseStorageEngineDetachStats() {
+    VeniceConfigLoader configLoader = mock(VeniceConfigLoader.class);
+    VeniceServerConfig serverConfig = mock(VeniceServerConfig.class);
+    when(serverConfig.getDataBasePath()).thenReturn("/tmp");
+    when(configLoader.getVeniceServerConfig()).thenReturn(serverConfig);
+    VeniceStoreVersionConfig storeConfig = mock(VeniceStoreVersionConfig.class);
+    when(storeConfig.getStorePersistenceType()).thenReturn(PersistenceType.BLACK_HOLE);
+    when(configLoader.getStoreConfig(storageEngineName)).thenReturn(storeConfig);
+
+    AggVersionedStorageEngineStats storageEngineStats = mock(AggVersionedStorageEngineStats.class);
+    StorageEngineFactory factory = mock(StorageEngineFactory.class);
+    Map<PersistenceType, StorageEngineFactory> factories = new HashMap<>();
+    factories.put(PersistenceType.BLACK_HOLE, factory);
+    StorageService storageService = new StorageService(
+        configLoader,
+        storageEngineStats,
+        mock(RocksDBMemoryStats.class),
+        mock(InternalAvroSpecificSerializer.class),
+        mock(InternalAvroSpecificSerializer.class),
+        storeRepository,
+        false,
+        false,
+        ignored -> true,
+        Optional.of(factories));
+
+    DelegatingStorageEngine storageEngine = mock(DelegatingStorageEngine.class);
+    when(storageEngine.getStoreVersionName()).thenReturn(storageEngineName);
+    when(storageEngine.getType()).thenReturn(PersistenceType.BLACK_HOLE);
+    storageService.getStorageEngineRepository().addLocalStorageEngine(storageEngine);
+
+    storageService.removeStorageEngine(storageEngineName);
+    verify(storageEngineStats).unsetStorageEngine(storageEngineName);
+    verify(storageEngine).drop();
+
+    storageService.getStorageEngineRepository().addLocalStorageEngine(storageEngine);
+    storageService.closeStorageEngine(storageEngineName);
+    verify(storageEngineStats, times(2)).unsetStorageEngine(storageEngineName);
+    verify(storageEngine).close();
+  }
+
   @Test
   public void testGetStoreAndUserPartitionsMapping() {
     VeniceConfigLoader configLoader = mock(VeniceConfigLoader.class);