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.connect.runtime;
  
  import org.apache.kafka.clients.CommonClientConfigs;
  import org.apache.kafka.common.MetricName;
  import org.apache.kafka.common.MetricNameTemplate;
  import org.apache.kafka.common.internals.Plugin;
  import org.apache.kafka.common.metrics.Gauge;
  import org.apache.kafka.common.metrics.KafkaMetricsContext;
  import org.apache.kafka.common.metrics.MetricConfig;
  import org.apache.kafka.common.metrics.Metrics;
  import org.apache.kafka.common.metrics.MetricsContext;
  import org.apache.kafka.common.metrics.MetricsReporter;
  import org.apache.kafka.common.metrics.Sensor;
  import org.apache.kafka.common.metrics.internals.MetricsUtils;
  import org.apache.kafka.common.metrics.internals.PluginMetricsImpl;
  import org.apache.kafka.common.utils.AppInfoParser;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.connect.connector.ConnectRecord;
  import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
  import org.apache.kafka.connect.storage.Converter;
  import org.apache.kafka.connect.storage.HeaderConverter;
  import org.apache.kafka.connect.transforms.Transformation;
  import org.apache.kafka.connect.transforms.predicates.Predicate;
  import org.apache.kafka.connect.util.ConnectorTaskId;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.TimeUnit;
  import java.util.function.Supplier;
  
  /**
   * The Connect metrics with configurable {@link MetricsReporter}s.
   */
  public class ConnectMetrics {
  
      public static final String JMX_PREFIX = "kafka.connect";
  
      private static final Logger LOG = LoggerFactory.getLogger(ConnectMetrics.class);
  
      private final Metrics metrics;
      private final Time time;
      private final String workerId;
      private final ConcurrentMap<MetricGroupId, MetricGroup> groupsByName = new ConcurrentHashMap<>();
      private final ConnectMetricsRegistry registry = new ConnectMetricsRegistry();
  
      /**
       * Create an instance.
       *
       * @param workerId the worker identifier; may not be null
       * @param config   the worker configuration; may not be null
       * @param time     the time; may not be null
       * @param clusterId the Kafka cluster ID
       */
      public ConnectMetrics(String workerId, WorkerConfig config, Time time, String clusterId) {
          this.workerId = workerId;
          this.time = time;
  
          int numSamples = config.getInt(CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG);
          long sampleWindowMs = config.getLong(CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_CONFIG);
          String metricsRecordingLevel = config.getString(CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG);
          List<MetricsReporter> reporters = CommonClientConfigs.metricsReporters(workerId, config);
          MetricConfig metricConfig = new MetricConfig().samples(numSamples)
                  .timeWindow(sampleWindowMs, TimeUnit.MILLISECONDS).recordLevel(
                          Sensor.RecordingLevel.forName(metricsRecordingLevel));
  
          Map<String, Object> contextLabels = new HashMap<>(config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX));
          contextLabels.put(WorkerConfig.CONNECT_KAFKA_CLUSTER_ID, clusterId);
          Object groupId = config.originals().get(DistributedConfig.GROUP_ID_CONFIG);
          if (groupId != null) {
              contextLabels.put(WorkerConfig.CONNECT_GROUP_ID, groupId);
          }
          MetricsContext metricsContext = new KafkaMetricsContext(JMX_PREFIX, contextLabels);
          this.metrics = new Metrics(metricConfig, reporters, time, metricsContext);
  
          LOG.debug("Registering Connect metrics with JMX for worker '{}'", workerId);
          AppInfoParser.registerAppInfo(JMX_PREFIX, workerId, metrics, time.milliseconds());
      }
  
      /**
       * Get the worker identifier.
       *
       * @return the worker ID; never null
       */
      public String workerId() {
          return workerId;
      }
  
      /**
       * Get the {@link Metrics Kafka Metrics} that are managed by this object and that should be used to
       * add sensors and individual metrics.
       *
       * @return the Kafka Metrics instance; never null
       */
      public Metrics metrics() {
          return metrics;
      }
  
      /**
       * Get the registry of metric names.
       *
       * @return the registry for the Connect metrics; never null
       */
      public ConnectMetricsRegistry registry() {
          return registry;
      }
  
      /**
       * Get or create a {@link MetricGroup} with the specified group name and the given tags.
       * Each group is uniquely identified by the name and tags.
       *
       * @param groupName    the name of the metric group; may not be null
       * @param tagKeyValues pairs of tag name and values
       * @return the {@link MetricGroup} that can be used to create metrics; never null
       * @throws IllegalArgumentException if the group name is not valid
       */
      public MetricGroup group(String groupName, String... tagKeyValues) {
          MetricGroupId groupId = groupId(groupName, tagKeyValues);
          MetricGroup group = groupsByName.get(groupId);
          if (group == null) {
              group = new MetricGroup(groupId);
              MetricGroup previous = groupsByName.putIfAbsent(groupId, group);
              if (previous != null)
                  group = previous;
          }
          return group;
      }
  
      protected MetricGroupId groupId(String groupName, String... tagKeyValues) {
          Map<String, String> tags = MetricsUtils.getTags(tagKeyValues);
          return new MetricGroupId(groupName, tags);
      }
  
      /**
       * Get the time.
       *
       * @return the time; never null
       */
      public Time time() {
          return time;
      }
  
      /**
       * Stop and unregister the metrics from any reporters.
       */
      public void stop() {
          metrics.close();
          LOG.debug("Unregistering Connect metrics with JMX for worker '{}'", workerId);
          AppInfoParser.unregisterAppInfo(JMX_PREFIX, workerId, metrics);
      }
  
      public PluginMetricsImpl connectorPluginMetrics(String connectorId) {
          return new PluginMetricsImpl(metrics, connectorPluginTags(connectorId));
      }
  
      private static Map<String, String> connectorPluginTags(String connectorId) {
          Map<String, String> tags = new LinkedHashMap<>();
          tags.put("connector", connectorId);
          return tags;
      }
  
      PluginMetricsImpl taskPluginMetrics(ConnectorTaskId connectorTaskId) {
          return new PluginMetricsImpl(metrics, taskPluginTags(connectorTaskId));
      }
  
      private static Map<String, String> taskPluginTags(ConnectorTaskId connectorTaskId) {
          Map<String, String> tags = connectorPluginTags(connectorTaskId.connector());
          tags.put("task", String.valueOf(connectorTaskId.task()));
          return tags;
      }
  
      private static Supplier<Map<String, String>> converterPluginTags(ConnectorTaskId connectorTaskId, boolean isKey) {
          return () -> {
              Map<String, String> tags = taskPluginTags(connectorTaskId);
              tags.put("converter", isKey ? "key" : "value");
              return tags;
          };
      }
  
      private static Supplier<Map<String, String>> headerConverterPluginTags(ConnectorTaskId connectorTaskId) {
          return () -> {
              Map<String, String> tags = taskPluginTags(connectorTaskId);
              tags.put("converter", "header");
              return tags;
          };
      }
  
      private static Supplier<Map<String, String>> transformationPluginTags(ConnectorTaskId connectorTaskId, String transformationAlias) {
          return () -> {
              Map<String, String> tags = taskPluginTags(connectorTaskId);
              tags.put("transformation", transformationAlias);
              return tags;
          };
      }
  
      private static Supplier<Map<String, String>> predicatePluginTags(ConnectorTaskId connectorTaskId, String predicateAlias) {
          return () -> {
              Map<String, String> tags = taskPluginTags(connectorTaskId);
              tags.put("predicate", predicateAlias);
              return tags;
          };
      }
  
      public Plugin<HeaderConverter> wrap(HeaderConverter headerConverter, ConnectorTaskId connectorTaskId) {
          return Plugin.wrapInstance(headerConverter, metrics, headerConverterPluginTags(connectorTaskId));
      }
  
      public Plugin<Converter> wrap(Converter converter, ConnectorTaskId connectorTaskId, boolean isKey) {
          return Plugin.wrapInstance(converter, metrics, converterPluginTags(connectorTaskId, isKey));
      }
  
      public <R extends ConnectRecord<R>> Plugin<Transformation<R>> wrap(Transformation<R> transformation, ConnectorTaskId connectorTaskId, String alias) {
          return Plugin.wrapInstance(transformation, metrics, transformationPluginTags(connectorTaskId, alias));
      }
  
      public <R extends ConnectRecord<R>> Plugin<Predicate<R>> wrap(Predicate<R> predicate, ConnectorTaskId connectorTaskId, String alias) {
          return Plugin.wrapInstance(predicate, metrics, predicatePluginTags(connectorTaskId, alias));
      }
  
      public static class MetricGroupId {
          private final String groupName;
          private final Map<String, String> tags;
          private final int hc;
          private final String str;
  
          public MetricGroupId(String groupName, Map<String, String> tags) {
              Objects.requireNonNull(groupName);
              Objects.requireNonNull(tags);
              this.groupName = groupName;
              this.tags = Collections.unmodifiableMap(new LinkedHashMap<>(tags));
              this.hc = Objects.hash(this.groupName, this.tags);
              StringBuilder sb = new StringBuilder(this.groupName);
              for (Map.Entry<String, String> entry : this.tags.entrySet()) {
                  sb.append(";").append(entry.getKey()).append('=').append(entry.getValue());
              }
              this.str = sb.toString();
          }
  
          /**
           * Get the group name.
           *
           * @return the group name; never null
           */
          public String groupName() {
              return groupName;
          }
  
          /**
           * Get the immutable map of tag names and values.
           *
           * @return the tags; never null
           */
          public Map<String, String> tags() {
              return tags;
          }
  
          /**
           * Determine if the supplied metric name is part of this group identifier.
           *
           * @param metricName the metric name
           * @return true if the metric name's group and tags match this group identifier, or false otherwise
           */
          public boolean includes(MetricName metricName) {
              return metricName != null && groupName.equals(metricName.group()) && tags.equals(metricName.tags());
          }
  
          @Override
          public int hashCode() {
              return hc;
          }
  
          @Override
          public boolean equals(Object obj) {
              if (obj == this)
                  return true;
              if (obj instanceof MetricGroupId that) {
                  return this.groupName.equals(that.groupName) && this.tags.equals(that.tags);
              }
              return false;
          }
  
          @Override
          public String toString() {
              return str;
          }
      }
  
      /**
       * A group of metrics. Each group maps to a JMX MBean and each metric maps to an MBean attribute.
       * <p>
       * Sensors should be added via the {@code sensor} methods on this class, rather than directly through
       * the {@link Metrics} class, so that the sensor names are made to be unique (based on the group name)
       * and so the sensors are removed when this group is {@link #close() closed}.
       */
      public class MetricGroup implements AutoCloseable {
          private final MetricGroupId groupId;
          private final Set<String> sensorNames = new HashSet<>();
          private final String sensorPrefix;
  
          /**
           * Create a group of Connect metrics.
           *
           * @param groupId the identifier of the group; may not be null and must be valid
           */
          protected MetricGroup(MetricGroupId groupId) {
              Objects.requireNonNull(groupId);
              this.groupId = groupId;
              sensorPrefix = "connect-sensor-group: " + groupId + ";";
          }
  
          /**
           * Get the group identifier.
           *
           * @return the group identifier; never null
           */
          public MetricGroupId groupId() {
              return groupId;
          }
  
          /**
           * Create the name of a metric that belongs to this group and has the group's tags.
           *
           * @param template the name template for the metric; may not be null
           * @return the metric name; never null
           * @throws IllegalArgumentException if the name is not valid
           */
          public MetricName metricName(MetricNameTemplate template) {
              return metrics.metricInstance(template, groupId.tags());
          }
  
          // for testing only
          MetricName metricName(String name) {
              return metrics.metricName(name, groupId.groupName(), "", groupId.tags());
          }
  
          /**
           * The {@link Metrics} that this group belongs to.
           * <p>
           * Do not use this to add {@link Sensor Sensors}, since they will not be removed when this group is
           * {@link #close() closed}. Metrics can be added directly, as long as the metric names are obtained from
           * this group via the {@link #metricName(MetricNameTemplate)} method.
           *
           * @return the metrics; never null
           */
          public Metrics metrics() {
              return metrics;
          }
  
          /**
           * The tags of this group.
           *
           * @return the unmodifiable tags; never null but may be empty
           */
          Map<String, String> tags() {
              return groupId.tags();
          }
  
          /**
           * Add to this group an indicator metric with a function that returns the current value.
           *
           * @param nameTemplate the name template for the metric; may not be null
           * @param supplier     the function used to determine the literal value of the metric; may not be null
           * @throws IllegalArgumentException if the name is not valid
           */
          public <T> void addValueMetric(MetricNameTemplate nameTemplate, final LiteralSupplier<T> supplier) {
              MetricName metricName = metricName(nameTemplate);
              metrics().addMetricIfAbsent(metricName, null, (Gauge<T>) (config, now) -> supplier.metricValue(now));
          }
  
          /**
           * Add to this group an indicator metric that always returns the specified value.
           *
           * @param nameTemplate the name template for the metric; may not be null
           * @param value        the value; may not be null
           * @throws IllegalArgumentException if the name is not valid
           */
          public <T> void addImmutableValueMetric(MetricNameTemplate nameTemplate, final T value) {
              MetricName metricName = metricName(nameTemplate);
              metrics().addMetricIfAbsent(metricName, null, (Gauge<T>) (config, now) -> value);
          }
  
          /**
           * Get or create a sensor with the given unique name and no parent sensors. This uses
           * a default recording level of INFO.
           *
           * @param name The sensor name
           * @return The sensor
           */
          public Sensor sensor(String name) {
              return sensor(name, null, Sensor.RecordingLevel.INFO);
          }
  
          /**
           * Get or create a sensor with the given unique name and no parent sensors. This uses
           * a default recording level of INFO.
           *
           * @param name The sensor name
           * @return The sensor
           */
          public Sensor sensor(String name, Sensor... parents) {
              return sensor(name, null, Sensor.RecordingLevel.INFO, parents);
          }
  
          /**
           * Get or create a sensor with the given unique name and zero or more parent sensors. All parent sensors will
           * receive every value recorded with this sensor.
           *
           * @param name           The name of the sensor
           * @param recordingLevel The recording level.
           * @param parents        The parent sensors
           * @return The sensor that is created
           */
          public Sensor sensor(String name, Sensor.RecordingLevel recordingLevel, Sensor... parents) {
              return sensor(name, null, recordingLevel, parents);
          }
  
          /**
           * Get or create a sensor with the given unique name and zero or more parent sensors. All parent sensors will
           * receive every value recorded with this sensor.
           *
           * @param name    The name of the sensor
           * @param config  A default configuration to use for this sensor for metrics that don't have their own config
           * @param parents The parent sensors
           * @return The sensor that is created
           */
          public Sensor sensor(String name, MetricConfig config, Sensor... parents) {
              return sensor(name, config, Sensor.RecordingLevel.INFO, parents);
          }
  
          /**
           * Get or create a sensor with the given unique name and zero or more parent sensors. All parent sensors will
           * receive every value recorded with this sensor.
           *
           * @param name           The name of the sensor
           * @param config         A default configuration to use for this sensor for metrics that don't have their own config
           * @param recordingLevel The recording level.
           * @param parents        The parent sensors
           * @return The sensor that is created
           */
          public synchronized Sensor sensor(String name, MetricConfig config, Sensor.RecordingLevel recordingLevel, Sensor... parents) {
              // We need to make sure that all sensor names are unique across all groups, so use the sensor prefix
              Sensor result = metrics.sensor(sensorPrefix + name, config, Long.MAX_VALUE, recordingLevel, parents);
              sensorNames.add(result.name());
              return result;
          }
  
          /**
           * Remove all sensors and metrics associated with this group.
           */
          public synchronized void close() {
              for (String sensorName : sensorNames) {
                  metrics.removeSensor(sensorName);
              }
              sensorNames.clear();
              for (MetricName metricName : new HashSet<>(metrics.metrics().keySet())) {
                  if (groupId.includes(metricName)) {
                      metrics.removeMetric(metricName);
                  }
              }
          }
      }
  
      /**
       * A simple functional interface that returns a literal value.
       */
      public interface LiteralSupplier<T> {
  
          /**
           * Return the literal value for the metric.
           *
           * @param now the current time in milliseconds
           * @return the literal metric value; may not be null
           */
          T metricValue(long now);
      }
  
      /**
       * Utility to generate the documentation for the Connect metrics.
       *
       * @param args the arguments
       */
      public static void main(String[] args) {
          ConnectMetricsRegistry metrics = new ConnectMetricsRegistry();
          System.out.println(Metrics.toHtmlTable(JMX_PREFIX, metrics.getAllTemplates()));
      }
  }
  