content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.streams.state.internals.metrics;
  
  import org.apache.kafka.common.Metric;
  import org.apache.kafka.common.MetricName;
  import org.apache.kafka.common.metrics.KafkaMetric;
  import org.apache.kafka.common.metrics.Metrics;
  import org.apache.kafka.common.utils.MockTime;
  import org.apache.kafka.streams.processor.TaskId;
  import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
  
  import org.junit.jupiter.api.Test;
  import org.rocksdb.Cache;
  import org.rocksdb.RocksDB;
  import org.rocksdb.Statistics;
  
  import java.math.BigInteger;
  import java.util.Map;
  
  import static org.apache.kafka.common.utils.Utils.mkEntry;
  import static org.apache.kafka.common.utils.Utils.mkMap;
  import static org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl.STATE_STORE_LEVEL_GROUP;
  import static org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl.STORE_ID_TAG;
  import static org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl.TASK_ID_TAG;
  import static org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl.THREAD_ID_TAG;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.CAPACITY_OF_BLOCK_CACHE;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.COMPACTION_PENDING;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.CURRENT_SIZE_OF_ACTIVE_MEMTABLE;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.CURRENT_SIZE_OF_ALL_MEMTABLES;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.ESTIMATED_BYTES_OF_PENDING_COMPACTION;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.ESTIMATED_MEMORY_OF_TABLE_READERS;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.ESTIMATED_NUMBER_OF_KEYS;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.LIVE_SST_FILES_SIZE;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.MEMTABLE_FLUSH_PENDING;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_BACKGROUND_ERRORS;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_DELETES_ACTIVE_MEMTABLE;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_DELETES_IMMUTABLE_MEMTABLES;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_ENTRIES_ACTIVE_MEMTABLE;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_ENTRIES_IMMUTABLE_MEMTABLES;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_IMMUTABLE_MEMTABLES;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_RUNNING_COMPACTIONS;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.NUMBER_OF_RUNNING_FLUSHES;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.PINNED_USAGE_OF_BLOCK_CACHE;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.SIZE_OF_ALL_MEMTABLES;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.TOTAL_SST_FILES_SIZE;
  import static org.apache.kafka.streams.state.internals.metrics.RocksDBMetrics.USAGE_OF_BLOCK_CACHE;
  import static org.hamcrest.CoreMatchers.is;
  import static org.hamcrest.CoreMatchers.notNullValue;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  public class RocksDBMetricsRecorderGaugesTest {
      private static final String METRICS_SCOPE = "metrics-scope";
      private static final TaskId TASK_ID = new TaskId(0, 0);
      private static final String STORE_NAME = "store-name";
      private static final String SEGMENT_STORE_NAME_1 = "segment-store-name-1";
      private static final String SEGMENT_STORE_NAME_2 = "segment-store-name-2";
      private static final String ROCKSDB_PROPERTIES_PREFIX = "rocksdb.";
  
      private final RocksDB dbToAdd1 = mock(RocksDB.class);
      private final RocksDB dbToAdd2 = mock(RocksDB.class);
      private final Cache cacheToAdd1 = mock(Cache.class);
      private final Cache cacheToAdd2 = mock(Cache.class);
      private final Statistics statisticsToAdd1 = mock(Statistics.class);
      private final Statistics statisticsToAdd2 = mock(Statistics.class);
  
      @Test
      public void shouldGetNumberOfImmutableMemTables() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_IMMUTABLE_MEMTABLES);
      }
  
      @Test
      public void shouldGetCurrentSizeofActiveMemTable() throws Exception {
          runAndVerifySumOfProperties(CURRENT_SIZE_OF_ACTIVE_MEMTABLE);
      }
  
      @Test
      public void shouldGetCurrentSizeofAllMemTables() throws Exception {
          runAndVerifySumOfProperties(CURRENT_SIZE_OF_ALL_MEMTABLES);
      }
  
      @Test
      public void shouldGetSizeofAllMemTables() throws Exception {
          runAndVerifySumOfProperties(SIZE_OF_ALL_MEMTABLES);
      }
  
      @Test
      public void shouldGetNumberOfEntriesActiveMemTable() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_ENTRIES_ACTIVE_MEMTABLE);
      }
  
      @Test
      public void shouldGetNumberOfDeletesActiveMemTable() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_DELETES_ACTIVE_MEMTABLE);
      }
  
      @Test
      public void shouldGetNumberOfEntriesImmutableMemTables() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_ENTRIES_IMMUTABLE_MEMTABLES);
      }
  
      @Test
      public void shouldGetNumberOfDeletesImmutableMemTables() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_DELETES_IMMUTABLE_MEMTABLES);
      }
  
      @Test
      public void shouldGetMemTableFlushPending() throws Exception {
          runAndVerifySumOfProperties(MEMTABLE_FLUSH_PENDING);
      }
  
      @Test
      public void shouldGetNumberOfRunningFlushes() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_RUNNING_FLUSHES);
      }
  
      @Test
      public void shouldGetCompactionPending() throws Exception {
          runAndVerifySumOfProperties(COMPACTION_PENDING);
      }
  
      @Test
      public void shouldGetNumberOfRunningCompactions() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_RUNNING_COMPACTIONS);
      }
  
      @Test
      public void shouldGetEstimatedBytesOfPendingCompactions() throws Exception {
          runAndVerifySumOfProperties(ESTIMATED_BYTES_OF_PENDING_COMPACTION);
      }
  
      @Test
      public void shouldGetTotalSstFilesSize() throws Exception {
          runAndVerifySumOfProperties(TOTAL_SST_FILES_SIZE);
      }
  
      @Test
      public void shouldGetLiveSstFilesSize() throws Exception {
          runAndVerifySumOfProperties(LIVE_SST_FILES_SIZE);
      }
  
      @Test
      public void shouldGetEstimatedNumberOfKeys() throws Exception {
          runAndVerifySumOfProperties(ESTIMATED_NUMBER_OF_KEYS);
      }
  
      @Test
      public void shouldGetEstimatedMemoryOfTableReaders() throws Exception {
          runAndVerifySumOfProperties(ESTIMATED_MEMORY_OF_TABLE_READERS);
      }
  
      @Test
      public void shouldGetNumberOfBackgroundErrors() throws Exception {
          runAndVerifySumOfProperties(NUMBER_OF_BACKGROUND_ERRORS);
      }
  
      @Test
      public void shouldGetCapacityOfBlockCacheWithMultipleCaches() throws Exception {
          runAndVerifyBlockCacheMetricsWithMultipleCaches(CAPACITY_OF_BLOCK_CACHE);
      }
  
      @Test
      public void shouldGetCapacityOfBlockCacheWithSingleCache() throws Exception {
          runAndVerifyBlockCacheMetricsWithSingleCache(CAPACITY_OF_BLOCK_CACHE);
      }
  
      @Test
      public void shouldGetUsageOfBlockCacheWithMultipleCaches() throws Exception {
          runAndVerifyBlockCacheMetricsWithMultipleCaches(USAGE_OF_BLOCK_CACHE);
      }
  
      @Test
      public void shouldGetUsageOfBlockCacheWithSingleCache() throws Exception {
          runAndVerifyBlockCacheMetricsWithSingleCache(USAGE_OF_BLOCK_CACHE);
      }
  
      @Test
      public void shouldGetPinnedUsageOfBlockCacheWithMultipleCaches() throws Exception {
          runAndVerifyBlockCacheMetricsWithMultipleCaches(PINNED_USAGE_OF_BLOCK_CACHE);
      }
  
      @Test
      public void shouldGetPinnedUsageOfBlockCacheWithSingleCache() throws Exception {
          runAndVerifyBlockCacheMetricsWithSingleCache(PINNED_USAGE_OF_BLOCK_CACHE);
      }
  
      private void runAndVerifySumOfProperties(final String propertyName) throws Exception {
          final StreamsMetricsImpl streamsMetrics =
              new StreamsMetricsImpl(new Metrics(), "test-client", "processId", new MockTime());
          final RocksDBMetricsRecorder recorder = new RocksDBMetricsRecorder(METRICS_SCOPE, STORE_NAME);
  
          recorder.init(streamsMetrics, TASK_ID);
          recorder.addValueProviders(SEGMENT_STORE_NAME_1, dbToAdd1, cacheToAdd1, statisticsToAdd1);
          recorder.addValueProviders(SEGMENT_STORE_NAME_2, dbToAdd2, cacheToAdd2, statisticsToAdd2);
  
          final long recordedValue1 = 5L;
          final long recordedValue2 = 3L;
          when(dbToAdd1.getAggregatedLongProperty(ROCKSDB_PROPERTIES_PREFIX + propertyName)).thenReturn(recordedValue1);
          when(dbToAdd2.getAggregatedLongProperty(ROCKSDB_PROPERTIES_PREFIX + propertyName)).thenReturn(recordedValue2);
  
          verifyMetrics(streamsMetrics, propertyName, recordedValue1 + recordedValue2);
      }
  
      private void runAndVerifyBlockCacheMetricsWithMultipleCaches(final String propertyName) throws Exception {
          final StreamsMetricsImpl streamsMetrics =
                  new StreamsMetricsImpl(new Metrics(), "test-client", "processId", new MockTime());
          final RocksDBMetricsRecorder recorder = new RocksDBMetricsRecorder(METRICS_SCOPE, STORE_NAME);
  
          recorder.init(streamsMetrics, TASK_ID);
          recorder.addValueProviders(SEGMENT_STORE_NAME_1, dbToAdd1, cacheToAdd1, statisticsToAdd1);
          recorder.addValueProviders(SEGMENT_STORE_NAME_2, dbToAdd2, cacheToAdd2, statisticsToAdd2);
  
          final long recordedValue1 = 5L;
          final long recordedValue2 = 3L;
          when(dbToAdd1.getLongProperty(ROCKSDB_PROPERTIES_PREFIX + propertyName)).thenReturn(recordedValue1);
          when(dbToAdd2.getLongProperty(ROCKSDB_PROPERTIES_PREFIX + propertyName)).thenReturn(recordedValue2);
  
          verifyMetrics(streamsMetrics, propertyName, recordedValue1 + recordedValue2);
      }
  
      private void runAndVerifyBlockCacheMetricsWithSingleCache(final String propertyName) throws Exception {
          final StreamsMetricsImpl streamsMetrics =
              new StreamsMetricsImpl(new Metrics(), "test-client", "processId", new MockTime());
          final RocksDBMetricsRecorder recorder = new RocksDBMetricsRecorder(METRICS_SCOPE, STORE_NAME);
  
          recorder.init(streamsMetrics, TASK_ID);
          recorder.addValueProviders(SEGMENT_STORE_NAME_1, dbToAdd1, cacheToAdd1, statisticsToAdd1);
          recorder.addValueProviders(SEGMENT_STORE_NAME_2, dbToAdd2, cacheToAdd1, statisticsToAdd2);
  
          final long recordedValue = 5L;
          when(dbToAdd1.getLongProperty(ROCKSDB_PROPERTIES_PREFIX + propertyName)).thenReturn(recordedValue);
          when(dbToAdd2.getLongProperty(ROCKSDB_PROPERTIES_PREFIX + propertyName)).thenReturn(recordedValue);
  
          verifyMetrics(streamsMetrics, propertyName, recordedValue);
      }
  
      private void verifyMetrics(final StreamsMetricsImpl streamsMetrics,
                                 final String propertyName,
                                 final long expectedValue) {
  
          final Map<MetricName, ? extends Metric> metrics = streamsMetrics.metrics();
          final Map<String, String> tagMap = mkMap(
              mkEntry(THREAD_ID_TAG, Thread.currentThread().getName()),
              mkEntry(TASK_ID_TAG, TASK_ID.toString()),
              mkEntry(METRICS_SCOPE + "-" + STORE_ID_TAG, STORE_NAME)
          );
          final KafkaMetric metric = (KafkaMetric) metrics.get(new MetricName(
              propertyName,
              STATE_STORE_LEVEL_GROUP,
              "description is ignored",
              tagMap
          ));
  
          assertThat(metric, notNullValue());
          assertThat(metric.metricValue(), is(BigInteger.valueOf(expectedValue)));
      }
  }
  