<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package com.uber.m3.tally.m3;

import com.uber.m3.tally.Buckets;
import com.uber.m3.tally.Capabilities;
import com.uber.m3.tally.CapableOf;
import com.uber.m3.tally.DurationBuckets;
import com.uber.m3.tally.StatsReporter;
import com.uber.m3.tally.ValueBuckets;
import com.uber.m3.tally.m3.thrift.TCalcTransport;
import com.uber.m3.tally.m3.thrift.TMultiUdpClient;
import com.uber.m3.tally.m3.thrift.TUdpClient;
import com.uber.m3.tally.m3.thrift.TUdpTransport;
import com.uber.m3.thrift.gen.CountValue;
import com.uber.m3.thrift.gen.GaugeValue;
import com.uber.m3.thrift.gen.M3;
import com.uber.m3.thrift.gen.Metric;
import com.uber.m3.thrift.gen.MetricBatch;
import com.uber.m3.thrift.gen.MetricTag;
import com.uber.m3.thrift.gen.MetricValue;
import com.uber.m3.thrift.gen.TimerValue;
import com.uber.m3.util.Duration;
import com.uber.m3.util.ImmutableMap;
import org.apache.http.annotation.NotThreadSafe;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TProtocolFactory;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetAddress;
import java.net.SocketAddress;
import java.net.SocketException;
import java.net.UnknownHostException;
import java.time.Clock;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 * An M3 implementation of a {@link StatsReporter}.
 */
public class M3Reporter implements StatsReporter, AutoCloseable {
    public static final String SERVICE_TAG = "service";
    public static final String ENV_TAG = "env";
    public static final String HOST_TAG = "host";
    public static final String DEFAULT_TAG_VALUE = "default";

    public static final String DEFAULT_HISTOGRAM_BUCKET_ID_NAME = "bucketid";
    public static final String DEFAULT_HISTOGRAM_BUCKET_NAME = "bucket";
    public static final int DEFAULT_HISTOGRAM_BUCKET_TAG_PRECISION = 6;

    /**
     * NOTE: DO NOT CHANGE THIS NUMBER!
     *       Reporter architecture is not suited for multi-processor setup and might cause some disruption
     *       to how metrics are processed and eventually submitted to M3 collectors;
     */
    static final int NUM_PROCESSORS = 1;

    private static final Logger LOG = LoggerFactory.getLogger(M3Reporter.class);

    private static final int MAX_DELAY_BEFORE_FLUSHING_MILLIS = 1_000;

    private static final int MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS = 5_000;

    private static final int DEFAULT_METRIC_SIZE = 100;
    private static final int DEFAULT_MAX_QUEUE_SIZE = 4096;
    private static final int DEFAULT_MAX_PACKET_SIZE = TUdpTransport.PACKET_DATA_PAYLOAD_MAX_SIZE;

    // NOTE: 256 bytes of overhead is reserved for Thrift metadata within UDP datagram payload
    private static final int THRIFT_METADATA_PADDING = 256;

    private static final int MIN_METRIC_BUCKET_ID_TAG_LENGTH = 4;

    private static final ThreadLocal<SerializedPayloadSizeEstimator> PAYLOAD_SIZE_ESTIMATOR =
            ThreadLocal.withInitial(SerializedPayloadSizeEstimator::new);

    private final Duration maxBufferingDelay;

    private final int payloadCapacity;

    private final String bucketIdTagKey;
    private final String bucketValueTagKey;
    private final String bucketValFmt;

    private final Set<MetricTag> commonTags;

    // TODO replace w/ an evicting queue
    private final BlockingQueue<SizedMetric> queue;

    // Executor service running processors flushing metrics to collectors
    private final ExecutorService executorService;

    private final ScheduledExecutorService scheduledExecutorService;

    private final Clock clock;

    // This is a synchronization barrier to make sure that reporter
    // is being shutdown only after all of its processor had done so
    private final CountDownLatch processorsShutdownLatch;

    // List of socket addresses for M3 collector endpoint
    private final SocketAddress[] collectorEndpointSockedAddresses;

    private final Processor[] processors;

    private final TProtocolFactory protocolFactory;

    private final AtomicBoolean isShutdown = new AtomicBoolean(false);

    // Use inner Builder class to construct an M3Reporter
    M3Reporter(Builder builder, TProtocolFactory thriftProtocolFactory) {
        payloadCapacity = calculatePayloadCapacity(builder.maxPacketSizeBytes, builder.metricTagSet);

        maxBufferingDelay = Duration.ofMillis(builder.maxProcessorWaitUntilFlushMillis);

        bucketIdTagKey = builder.histogramBucketIdName;
        bucketValueTagKey = builder.histogramBucketName;
        bucketValFmt = String.format("%%.%df", builder.histogramBucketTagPrecision);

        queue = new LinkedBlockingQueue<>(builder.maxQueueSize);

        ThreadFactory namedThreadFactory = createThreadFactory();

        executorService = builder.executor != null ? builder.executor : Executors.newFixedThreadPool(NUM_PROCESSORS, namedThreadFactory);
        scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(namedThreadFactory);

        clock = Clock.systemUTC();

        commonTags = builder.metricTagSet;

        protocolFactory = thriftProtocolFactory;

        processorsShutdownLatch = new CountDownLatch(NUM_PROCESSORS);

        collectorEndpointSockedAddresses = builder.endpointSocketAddresses;
        processors = new Processor[NUM_PROCESSORS];

        for (int i = 0; i < NUM_PROCESSORS; ++i) {
            processors[i] = bootProcessor(collectorEndpointSockedAddresses);
        }

        // Schedule regular heartbeat up-keeping processors up and running
        scheduledExecutorService.scheduleAtFixedRate(this::heartbeat, 1, 1, TimeUnit.SECONDS);
    }

    // NOTE: This method is not concurrent
    void heartbeat() {
        synchronized (this) {
            for (int i = 0; i < processors.length; ++i) {
                if (processors[i].getState() != ProcessorState.RUNNING) {
                    processors[i] = bootProcessor(collectorEndpointSockedAddresses);
                }
            }
        }
    }

    private int calculatePayloadCapacity(int maxPacketSizeBytes, Set<MetricTag> commonTags) {
        MetricBatch metricBatch = new MetricBatch();
        metricBatch.setCommonTags(commonTags);
        metricBatch.setMetrics(new ArrayList<>());

        int thriftRequestShellSize = PAYLOAD_SIZE_ESTIMATOR.get().evaluateThriftRequestWireSize(metricBatch);

        int payloadCapacity = maxPacketSizeBytes - (THRIFT_METADATA_PADDING + thriftRequestShellSize);
        if (payloadCapacity <= 0) {
            throw new IllegalArgumentException("Common tags serialized size exceeds packet size");
        }

        return payloadCapacity;
    }

    private static String getHostName() {
        try {
            return InetAddress.getLocalHost().getHostName();
        } catch (UnknownHostException e) {
            LOG.warn("Unable to determine hostname. Defaulting to: {}", DEFAULT_TAG_VALUE);
            return DEFAULT_TAG_VALUE;
        }
    }

    private Processor bootProcessor(SocketAddress[] endpointSocketAddresses) {
        try {
            Processor processor = new Processor(endpointSocketAddresses, protocolFactory);
            executorService.execute(processor);
            return processor;
        } catch (TTransportException | SocketException e) {
            LOG.error("Failed to boot processor", e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public Capabilities capabilities() {
        return CapableOf.REPORTING_TAGGING;
    }

    @Override
    public void flush() {
        if (isShutdown.get()) {
            return;
        }

        for (Processor processor : processors) {
            processor.scheduleFlush();
        }
    }

    // NOTE: This should only be used in tests
    void flushNow() throws TException {
        for (Processor processor : processors) {
            processor.flushBuffered();
        }
    }

    @Override
    public void close() {
        if (!isShutdown.compareAndSet(false, true)) {
            // Shutdown already
            return;
        }

        // Important to use `shutdownNow` instead of `shutdown` to interrupt processor
        // thread(s) or else they will block forever
        executorService.shutdownNow();

        try {
            // Wait a maximum of `MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS` for all processors
            // to complete
            if (!processorsShutdownLatch.await(MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS, TimeUnit.MILLISECONDS)) {
                LOG.warn(
                        "M3Reporter closing before Processors complete after waiting timeout of {}ms!",
                        MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS
                );
            }
        } catch (InterruptedException e) {
            LOG.warn("M3Reporter closing before Processors complete due to being interrupted!");
        }
    }

    private static Set<MetricTag> toMetricTagSet(Map<String, String> tags) {
        Set<MetricTag> metricTagSet = new HashSet<>();

        if (tags == null) {
            return metricTagSet;
        }

        for (Map.Entry<String, String> tag : tags.entrySet()) {<commit_msg>Shutdown scheduled executor service as well<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package com.uber.m3.tally.m3;

import com.uber.m3.tally.Buckets;
import com.uber.m3.tally.Capabilities;
import com.uber.m3.tally.CapableOf;
import com.uber.m3.tally.DurationBuckets;
import com.uber.m3.tally.StatsReporter;
import com.uber.m3.tally.ValueBuckets;
import com.uber.m3.tally.m3.thrift.TCalcTransport;
import com.uber.m3.tally.m3.thrift.TMultiUdpClient;
import com.uber.m3.tally.m3.thrift.TUdpClient;
import com.uber.m3.tally.m3.thrift.TUdpTransport;
import com.uber.m3.thrift.gen.CountValue;
import com.uber.m3.thrift.gen.GaugeValue;
import com.uber.m3.thrift.gen.M3;
import com.uber.m3.thrift.gen.Metric;
import com.uber.m3.thrift.gen.MetricBatch;
import com.uber.m3.thrift.gen.MetricTag;
import com.uber.m3.thrift.gen.MetricValue;
import com.uber.m3.thrift.gen.TimerValue;
import com.uber.m3.util.Duration;
import com.uber.m3.util.ImmutableMap;
import org.apache.http.annotation.NotThreadSafe;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TProtocolFactory;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetAddress;
import java.net.SocketAddress;
import java.net.SocketException;
import java.net.UnknownHostException;
import java.time.Clock;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 * An M3 implementation of a {@link StatsReporter}.
 */
public class M3Reporter implements StatsReporter, AutoCloseable {
    public static final String SERVICE_TAG = "service";
    public static final String ENV_TAG = "env";
    public static final String HOST_TAG = "host";
    public static final String DEFAULT_TAG_VALUE = "default";

    public static final String DEFAULT_HISTOGRAM_BUCKET_ID_NAME = "bucketid";
    public static final String DEFAULT_HISTOGRAM_BUCKET_NAME = "bucket";
    public static final int DEFAULT_HISTOGRAM_BUCKET_TAG_PRECISION = 6;

    /**
     * NOTE: DO NOT CHANGE THIS NUMBER!
     *       Reporter architecture is not suited for multi-processor setup and might cause some disruption
     *       to how metrics are processed and eventually submitted to M3 collectors;
     */
    static final int NUM_PROCESSORS = 1;

    private static final Logger LOG = LoggerFactory.getLogger(M3Reporter.class);

    private static final int MAX_DELAY_BEFORE_FLUSHING_MILLIS = 1_000;

    private static final int MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS = 5_000;

    private static final int DEFAULT_METRIC_SIZE = 100;
    private static final int DEFAULT_MAX_QUEUE_SIZE = 4096;
    private static final int DEFAULT_MAX_PACKET_SIZE = TUdpTransport.PACKET_DATA_PAYLOAD_MAX_SIZE;

    // NOTE: 256 bytes of overhead is reserved for Thrift metadata within UDP datagram payload
    private static final int THRIFT_METADATA_PADDING = 256;

    private static final int MIN_METRIC_BUCKET_ID_TAG_LENGTH = 4;

    private static final ThreadLocal<SerializedPayloadSizeEstimator> PAYLOAD_SIZE_ESTIMATOR =
            ThreadLocal.withInitial(SerializedPayloadSizeEstimator::new);

    private final Duration maxBufferingDelay;

    private final int payloadCapacity;

    private final String bucketIdTagKey;
    private final String bucketValueTagKey;
    private final String bucketValFmt;

    private final Set<MetricTag> commonTags;

    // TODO replace w/ an evicting queue
    private final BlockingQueue<SizedMetric> queue;

    // Executor service running processors flushing metrics to collectors
    private final ExecutorService executorService;

    private final ScheduledExecutorService scheduledExecutorService;

    private final Clock clock;

    // This is a synchronization barrier to make sure that reporter
    // is being shutdown only after all of its processor had done so
    private final CountDownLatch processorsShutdownLatch;

    // List of socket addresses for M3 collector endpoint
    private final SocketAddress[] collectorEndpointSockedAddresses;

    private final Processor[] processors;

    private final TProtocolFactory protocolFactory;

    private final AtomicBoolean isShutdown = new AtomicBoolean(false);

    // Use inner Builder class to construct an M3Reporter
    M3Reporter(Builder builder, TProtocolFactory thriftProtocolFactory) {
        payloadCapacity = calculatePayloadCapacity(builder.maxPacketSizeBytes, builder.metricTagSet);

        maxBufferingDelay = Duration.ofMillis(builder.maxProcessorWaitUntilFlushMillis);

        bucketIdTagKey = builder.histogramBucketIdName;
        bucketValueTagKey = builder.histogramBucketName;
        bucketValFmt = String.format("%%.%df", builder.histogramBucketTagPrecision);

        queue = new LinkedBlockingQueue<>(builder.maxQueueSize);

        ThreadFactory namedThreadFactory = createThreadFactory();

        executorService = builder.executor != null ? builder.executor : Executors.newFixedThreadPool(NUM_PROCESSORS, namedThreadFactory);
        scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(namedThreadFactory);

        clock = Clock.systemUTC();

        commonTags = builder.metricTagSet;

        protocolFactory = thriftProtocolFactory;

        processorsShutdownLatch = new CountDownLatch(NUM_PROCESSORS);

        collectorEndpointSockedAddresses = builder.endpointSocketAddresses;
        processors = new Processor[NUM_PROCESSORS];

        for (int i = 0; i < NUM_PROCESSORS; ++i) {
            processors[i] = bootProcessor(collectorEndpointSockedAddresses);
        }

        // Schedule regular heartbeat up-keeping processors up and running
        scheduledExecutorService.scheduleAtFixedRate(this::heartbeat, 1, 1, TimeUnit.SECONDS);
    }

    // NOTE: This method is not concurrent
    void heartbeat() {
        synchronized (this) {
            for (int i = 0; i < processors.length; ++i) {
                if (processors[i].getState() != ProcessorState.RUNNING) {
                    processors[i] = bootProcessor(collectorEndpointSockedAddresses);
                }
            }
        }
    }

    private int calculatePayloadCapacity(int maxPacketSizeBytes, Set<MetricTag> commonTags) {
        MetricBatch metricBatch = new MetricBatch();
        metricBatch.setCommonTags(commonTags);
        metricBatch.setMetrics(new ArrayList<>());

        int thriftRequestShellSize = PAYLOAD_SIZE_ESTIMATOR.get().evaluateThriftRequestWireSize(metricBatch);

        int payloadCapacity = maxPacketSizeBytes - (THRIFT_METADATA_PADDING + thriftRequestShellSize);
        if (payloadCapacity <= 0) {
            throw new IllegalArgumentException("Common tags serialized size exceeds packet size");
        }

        return payloadCapacity;
    }

    private static String getHostName() {
        try {
            return InetAddress.getLocalHost().getHostName();
        } catch (UnknownHostException e) {
            LOG.warn("Unable to determine hostname. Defaulting to: {}", DEFAULT_TAG_VALUE);
            return DEFAULT_TAG_VALUE;
        }
    }

    private Processor bootProcessor(SocketAddress[] endpointSocketAddresses) {
        try {
            Processor processor = new Processor(endpointSocketAddresses, protocolFactory);
            executorService.execute(processor);
            return processor;
        } catch (TTransportException | SocketException e) {
            LOG.error("Failed to boot processor", e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public Capabilities capabilities() {
        return CapableOf.REPORTING_TAGGING;
    }

    @Override
    public void flush() {
        if (isShutdown.get()) {
            return;
        }

        for (Processor processor : processors) {
            processor.scheduleFlush();
        }
    }

    // NOTE: This should only be used in tests
    void flushNow() throws TException {
        for (Processor processor : processors) {
            processor.flushBuffered();
        }
    }

    @Override
    public void close() {
        if (!isShutdown.compareAndSet(false, true)) {
            // Shutdown already
            return;
        }

        // Important to use `shutdownNow` instead of `shutdown` to interrupt processor
        // thread(s) or else they will block forever
        scheduledExecutorService.shutdownNow();
        executorService.shutdownNow();

        try {
            // Wait a maximum of `MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS` for all processors
            // to complete
            if (!processorsShutdownLatch.await(MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS, TimeUnit.MILLISECONDS)) {
                LOG.warn(
                        "M3Reporter closing before Processors complete after waiting timeout of {}ms!",
                        MAX_PROCESSOR_WAIT_ON_CLOSE_MILLIS
                );
            }
        } catch (InterruptedException e) {
            LOG.warn("M3Reporter closing before Processors complete due to being interrupted!");
        }
    }

    private static Set<MetricTag> toMetricTagSet(Map<String, String> tags) {
        Set<MetricTag> metricTagSet = new HashSet<>();

        if (tags == null) {
            return metricTagSet;
        }

        for (Map.Entry<String, String> tag : tags.entrySet()) {
{'fraction_non_alphanumeric': 0.07802058876648005, 'fraction_numerical': 0.006411414123171392, 'mean_word_length': 3.40337972166998, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '131727', 'n_tokens_mistral': 6232, 'n_tokens_neox': 5976, 'n_words': 1674}