alluxio-master-2.x/core/server/common/src/main/java/alluxio/metrics/sink/RatisDropwizardExports.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.metrics.sink;

import com.codahale.metrics.MetricRegistry;
import io.prometheus.client.Collector;
import io.prometheus.client.CollectorRegistry;
import io.prometheus.client.dropwizard.DropwizardExports;
import io.prometheus.client.dropwizard.samplebuilder.DefaultSampleBuilder;
import org.apache.ratis.metrics.MetricRegistries;
import org.apache.ratis.metrics.RatisMetricRegistry;

import java.util.Map;

/**
 * Collect Dropwizard metrics, but rename ratis specific metrics.
 */
public class RatisDropwizardExports extends DropwizardExports {

  /**
   * Creates a new DropwizardExports with a {@link DefaultSampleBuilder}.
   *
   * @param registry a metric registry to export in prometheus
   */
  public RatisDropwizardExports(MetricRegistry registry) {
    super(registry, new RatisNameRewriteSampleBuilder());
  }

  /**
   * Register ratis metric to metricRegistry.
   * @param ratisMetricsMap a map to store the registered metrics
   */
  public static void registerRatisMetricReporters(
      Map<String, RatisDropwizardExports> ratisMetricsMap) {
    MetricRegistries.global().addReporterRegistration(
        r1 -> registerDropwizard(r1, ratisMetricsMap),
        r2 -> deregisterDropwizard(r2, ratisMetricsMap));
  }

  private static void registerDropwizard(RatisMetricRegistry registry,
      Map<String, RatisDropwizardExports> ratisMetricsMap) {
    RatisDropwizardExports rde = new RatisDropwizardExports(
        registry.getDropWizardMetricRegistry());
    CollectorRegistry.defaultRegistry.register(rde);
    String name = registry.getMetricRegistryInfo().getName();
    ratisMetricsMap.putIfAbsent(name, rde);
  }

  private static void deregisterDropwizard(RatisMetricRegistry registry,
      Map<String, RatisDropwizardExports> ratisMetricsMap) {
    String name = registry.getMetricRegistryInfo().getName();
    Collector c = ratisMetricsMap.remove(name);
    if (c != null) {
      CollectorRegistry.defaultRegistry.unregister(c);
    }
  }
}
