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.celeborn.common.network.util;
  
  import java.lang.reflect.Method;
  import java.lang.reflect.Modifier;
  import java.util.*;
  
  import com.codahale.metrics.MetricRegistry;
  import com.google.common.annotations.VisibleForTesting;
  import io.netty.buffer.*;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.metrics.source.AbstractSource;
  
  /** A Netty memory metrics class to collect metrics from Netty ByteBufAllocator. */
  public class NettyMemoryMetrics {
    private final Logger logger = LoggerFactory.getLogger(this.getClass());
  
    private final ByteBufAllocator allocator;
  
    private final boolean verboseMetricsEnabled;
  
    private final String metricPrefix;
  
    private final AbstractSource source;
  
    private final Map<String, String> labels;
  
    @VisibleForTesting static final Set<String> VERBOSE_METRICS = new HashSet<>();
  
    static {
      VERBOSE_METRICS.addAll(
          Arrays.asList(
              "numAllocations",
              "numTinyAllocations",
              "numSmallAllocations",
              "numNormalAllocations",
              "numHugeAllocations",
              "numDeallocations",
              "numTinyDeallocations",
              "numSmallDeallocations",
              "numNormalDeallocations",
              "numHugeDeallocations",
              "numActiveAllocations",
              "numActiveTinyAllocations",
              "numActiveSmallAllocations",
              "numActiveNormalAllocations",
              "numActiveHugeAllocations",
              "numActiveBytes"));
    }
  
    public NettyMemoryMetrics(
        ByteBufAllocator allocator,
        String metricPrefix,
        boolean verboseMetricsEnabled,
        AbstractSource source,
        Map<String, String> labels) {
      this.allocator = allocator;
      this.metricPrefix = metricPrefix;
      this.verboseMetricsEnabled = verboseMetricsEnabled;
      this.source = source;
      this.labels = labels;
  
      registerMetrics();
    }
  
    private void registerMetrics() {
      // Register general metrics.
      if (source != null) {
        if (allocator instanceof UnpooledByteBufAllocator) {
          logger.debug("Setup netty metrics for UnpooledByteBufAllocator");
          ByteBufAllocatorMetric unpooledMetric = ((UnpooledByteBufAllocator) allocator).metric();
          source.addGauge(
              MetricRegistry.name(metricPrefix, "usedHeapMemory"),
              labels,
              unpooledMetric::usedHeapMemory);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "usedDirectMemory"),
              labels,
              unpooledMetric::usedDirectMemory);
        } else if (allocator instanceof PooledByteBufAllocator) {
          logger.debug("Setup netty metrics for PooledByteBufAllocator");
          PooledByteBufAllocatorMetric pooledAllocatorMetric =
              ((PooledByteBufAllocator) allocator).metric();
          source.addGauge(
              MetricRegistry.name(metricPrefix, "usedHeapMemory"),
              labels,
              pooledAllocatorMetric::usedHeapMemory);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "usedDirectMemory"),
              labels,
              pooledAllocatorMetric::usedDirectMemory);
  
          source.addGauge(
              MetricRegistry.name(metricPrefix, "numHeapArenas"),
              labels,
              pooledAllocatorMetric::numHeapArenas);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "numDirectArenas"),
              labels,
              pooledAllocatorMetric::numDirectArenas);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "tinyCacheSize"),
              labels,
              pooledAllocatorMetric::tinyCacheSize);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "smallCacheSize"),
              labels,
              pooledAllocatorMetric::smallCacheSize);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "normalCacheSize"),
              labels,
              pooledAllocatorMetric::normalCacheSize);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "numThreadLocalCaches"),
              labels,
              pooledAllocatorMetric::numThreadLocalCaches);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "chunkSize"),
              labels,
              pooledAllocatorMetric::chunkSize);
          source.addGauge(
              MetricRegistry.name(metricPrefix, "pinnedHeapMemory"),
              labels,
              () -> ((PooledByteBufAllocator) allocator).pinnedHeapMemory());
          source.addGauge(
              MetricRegistry.name(metricPrefix, "pinnedDirectMemory"),
              labels,
              () -> ((PooledByteBufAllocator) allocator).pinnedDirectMemory());
          if (verboseMetricsEnabled) {
            int directArenaIndex = 0;
            for (PoolArenaMetric metric : pooledAllocatorMetric.directArenas()) {
              registerArenaMetric(metric, "directArena" + directArenaIndex);
              directArenaIndex++;
            }
  
            int heapArenaIndex = 0;
            for (PoolArenaMetric metric : pooledAllocatorMetric.heapArenas()) {
              registerArenaMetric(metric, "heapArena" + heapArenaIndex);
              heapArenaIndex++;
            }
          }
        }
      }
    }
  
    private void registerArenaMetric(PoolArenaMetric arenaMetric, String arenaName) {
      for (String methodName : VERBOSE_METRICS) {
        Method m;
        try {
          m = PoolArenaMetric.class.getMethod(methodName);
        } catch (Exception e) {
          // Failed to find metric related method, ignore this metric.
          continue;
        }
  
        if (!Modifier.isPublic(m.getModifiers())) {
          // Ignore non-public methods.
          continue;
        }
  
        Class<?> returnType = m.getReturnType();
        String metricName = MetricRegistry.name(metricPrefix, arenaName, m.getName());
        if (returnType.equals(int.class)) {
          source.addGauge(
              metricName,
              labels,
              () -> {
                try {
                  return (Integer) m.invoke(arenaMetric);
                } catch (Exception e) {
                  return -1; // Swallow the exceptions.
                }
              });
  
        } else if (returnType.equals(long.class)) {
          source.addGauge(
              metricName,
              labels,
              () -> {
                try {
                  return (Long) m.invoke(arenaMetric);
                } catch (Exception e) {
                  return -1L; // Swallow the exceptions.
                }
              });
        }
      }
    }
  }
  