/*
 * 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.samza.storage.kv;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import org.apache.samza.SamzaException;
import org.apache.samza.metrics.Counter;
import org.apache.samza.metrics.Gauge;
import org.apache.samza.metrics.MetricsRegistryMap;
import org.apache.samza.metrics.MetricsVisitor;
import org.apache.samza.metrics.Timer;
import org.apache.samza.serializers.Serde;
import org.apache.samza.serializers.StringSerde;
import org.apache.samza.storage.kv.inmemory.InMemoryKeyValueStore;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;


/**
 * This class is used to test whether the key value size metrics in {@link SerializedKeyValueStoreMetrics} works correctly.
 */
public class TestKeyValueSizeHistogramMetric {

  private static String storeName = "teststore";
  private static String keyPrefix = "key-size-bytes-histogram";
  private static String valuePrefix = "value-size-bytes-histogram";
  private static MetricsRegistryMap metricsRegistry = new MetricsRegistryMap();
  private static KeyValueStoreMetrics keyValueStoreMetrics = new KeyValueStoreMetrics(storeName, metricsRegistry);
  private static SerializedKeyValueStoreMetrics serializedKeyValueStoreMetrics =
      new SerializedKeyValueStoreMetrics(storeName, metricsRegistry);
  private static Serde<String> stringSerde = new StringSerde();
  private static Random random = new Random();

  private KeyValueStore<String, String> store = null;

  @Before
  public void setup() {
    KeyValueStore<byte[], byte[]> kvStore = new InMemoryKeyValueStore(keyValueStoreMetrics);
    KeyValueStore<String, String> serializedStore =
        new SerializedKeyValueStore<>(kvStore, stringSerde, stringSerde, serializedKeyValueStoreMetrics);
    store = new NullSafeKeyValueStore<>(serializedStore);
  }

  /**
   * Make sure that the histograms can record the key value size and we can use a
   * {@link MetricsVisitor} to get access to the value store in the histograms
   */
  @Test
  public void testHistogramMetric() {

    List<String> keys = new ArrayList<>();
    List<String> values = new ArrayList<>();

    for (int i = 0; i < 300; i++) {
      keys.add(getRandomString());
      values.add(getRandomString());
    }

    for (int i = 0; i < keys.size(); i++) {
      store.put(keys.get(i), values.get(i));
    }

    Set<String> names = new HashSet<>();
    for (Double p : serializedKeyValueStoreMetrics.record_key_size_percentiles()) {
      names.add(storeName + "-" + keyPrefix + "_" + p);
    }

    for (Double p : serializedKeyValueStoreMetrics.record_value_size_percentiles()) {
      names.add(storeName + "-" + valuePrefix + "_" + p);
    }

    metricsRegistry.getGroups().forEach(group -> metricsRegistry.getGroup(group.toString()).forEach((name, metric) -> {
        if (names.contains(name)) {
          //System.out.println(name);
          metric.visit(new MetricsVisitor() {
            @Override
            public void counter(Counter counter) {

            }

            @Override
            public <T> void gauge(Gauge<T> gauge) {
              Double num = (Double) gauge.getValue();
              //System.out.println(num);
              Assert.assertNotEquals(0D, (Double) gauge.getValue(), 0.0001);
            }

            @Override
            public void timer(Timer timer) {

            }
          });
        }
      }));
  }

  private String getRandomString() {
    int leftLimit = 97; // letter 'a'
    int rightLimit = 122; // letter 'z'
    int maxLength = 100;

    String generatedString = random.ints(leftLimit, rightLimit + 1)
        .limit(random.nextInt(maxLength))
        .collect(StringBuilder::new, StringBuilder::appendCodePoint, StringBuilder::append)
        .toString();

    return generatedString;
  }

  @After
  public void teardown() {
    try {
      store.close();
    } catch (SamzaException e) {
      e.printStackTrace();
    }
  }
}

@@ -840,16 +840,37 @@
 ge.kv;%0A%0A
+import java.io.File;%0A
 import j
@@ -885,24 +885,24 @@
 .ArrayList;%0A
-
 import java.
@@ -1025,16 +1025,97 @@
 eption;%0A
+import org.apache.samza.config.Config;%0Aimport org.apache.samza.config.MapConfig;%0A
 import o
@@ -1581,16 +1581,110 @@
 t.Test;%0A
+import org.rocksdb.FlushOptions;%0Aimport org.rocksdb.Options;%0Aimport org.rocksdb.WriteOptions;%0A
 %0A%0A/**%0A *
@@ -2589,41 +2589,253 @@
 ) %7B%0A
+%0A
     
-KeyValueStore%3Cbyte%5B%5D, byte%5B%5D%3E
+Config config = new MapConfig();%0A    Options options = new Options();%0A    options.setCreateIfMissing(true);%0A%0A    File dbDir = new File(System.getProperty(%22java.io.tmpdir%22) + %22/dbStore%22 + System.currentTimeMillis());%0A    RocksDbKeyValueStore
  kvS
@@ -2845,24 +2845,23 @@
 e = new 
-InMemory
+RocksDb
 KeyValue
@@ -2866,36 +2866,167 @@
 ueStore(
-keyValueStoreMetrics
+dbDir, options, config, false, %22dbStore%22,%0A        new WriteOptions(), new FlushOptions(), new KeyValueStoreMetrics(%22dbStore%22, new MetricsRegistryMap())
 );%0A    K

Use RocksDbKeyValueStore instead of InMemoryKeyValueStore