kafka ProduceBenchWorker 源码

kafka ProduceBenchWorker 代码

文件路径：/trogdor/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchWorker.java

/*
 * 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.trogdor.workload;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.node.TextNode;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.producer.Callback;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.internals.KafkaFutureImpl;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.utils.ThreadUtils;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.trogdor.common.JsonUtil;
import org.apache.kafka.trogdor.common.Platform;
import org.apache.kafka.trogdor.common.WorkerUtils;
import org.apache.kafka.trogdor.task.TaskWorker;
import org.apache.kafka.trogdor.task.WorkerStatusTracker;
import org.apache.kafka.trogdor.workload.TransactionGenerator.TransactionAction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Optional;
import java.util.Properties;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

public class ProduceBenchWorker implements TaskWorker {
    private static final Logger log = LoggerFactory.getLogger(ProduceBenchWorker.class);
    
    private static final int THROTTLE_PERIOD_MS = 100;

    private final String id;

    private final ProduceBenchSpec spec;

    private final AtomicBoolean running = new AtomicBoolean(false);

    private ScheduledExecutorService executor;

    private WorkerStatusTracker status;

    private KafkaFutureImpl<String> doneFuture;

    public ProduceBenchWorker(String id, ProduceBenchSpec spec) {
        this.id = id;
        this.spec = spec;
    }

    @Override
    public void start(Platform platform, WorkerStatusTracker status,
                      KafkaFutureImpl<String> doneFuture) {
        if (!running.compareAndSet(false, true)) {
            throw new IllegalStateException("ProducerBenchWorker is already running.");
        }
        log.info("{}: Activating ProduceBenchWorker with {}", id, spec);
        // Create an executor with 2 threads.  We need the second thread so
        // that the StatusUpdater can run in parallel with SendRecords.
        this.executor = Executors.newScheduledThreadPool(2,
            ThreadUtils.createThreadFactory("ProduceBenchWorkerThread%d", false));
        this.status = status;
        this.doneFuture = doneFuture;
        executor.submit(new Prepare());
    }

    public class Prepare implements Runnable {
        @Override
        public void run() {
            try {
                Map<String, NewTopic> newTopics = new HashMap<>();
                HashSet<TopicPartition> active = new HashSet<>();
                for (Map.Entry<String, PartitionsSpec> entry :
                        spec.activeTopics().materialize().entrySet()) {
                    String topicName = entry.getKey();
                    PartitionsSpec partSpec = entry.getValue();
                    newTopics.put(topicName, partSpec.newTopic(topicName));
                    for (Integer partitionNumber : partSpec.partitionNumbers()) {
                        active.add(new TopicPartition(topicName, partitionNumber));
                    }
                }
                if (active.isEmpty()) {
                    throw new RuntimeException("You must specify at least one active topic.");
                }
                for (Map.Entry<String, PartitionsSpec> entry :
                        spec.inactiveTopics().materialize().entrySet()) {
                    String topicName = entry.getKey();
                    PartitionsSpec partSpec = entry.getValue();
                    newTopics.put(topicName, partSpec.newTopic(topicName));
                }
                status.update(new TextNode("Creating " + newTopics.keySet().size() + " topic(s)"));
                WorkerUtils.createTopics(log, spec.bootstrapServers(), spec.commonClientConf(),
                                         spec.adminClientConf(), newTopics, false);
                status.update(new TextNode("Created " + newTopics.keySet().size() + " topic(s)"));
                executor.submit(new SendRecords(active));
            } catch (Throwable e) {
                WorkerUtils.abort(log, "Prepare", e, doneFuture);
            }
        }
    }

    private static class SendRecordsCallback implements Callback {
        private final SendRecords sendRecords;
        private final long startMs;

        SendRecordsCallback(SendRecords sendRecords, long startMs) {
            this.sendRecords = sendRecords;
            this.startMs = startMs;
        }

        @Override
        public void onCompletion(RecordMetadata metadata, Exception exception) {
            long now = Time.SYSTEM.milliseconds();
            long durationMs = now - startMs;
            sendRecords.recordDuration(durationMs);
            if (exception != null) {
                log.error("SendRecordsCallback: error", exception);
            }
        }
    }

    /**
     * A subclass of Throttle which flushes the Producer right before the throttle injects a delay.
     * This avoids including throttling latency in latency measurements.
     */
    private static class SendRecordsThrottle extends Throttle {
        private final KafkaProducer<?, ?> producer;

        SendRecordsThrottle(int maxPerPeriod, KafkaProducer<?, ?> producer) {
            super(maxPerPeriod, THROTTLE_PERIOD_MS);
            this.producer = producer;
        }

        @Override
        protected synchronized void delay(long amount) throws InterruptedException {
            long startMs = time().milliseconds();
            producer.flush();
            long endMs = time().milliseconds();
            long delta = endMs - startMs;
            super.delay(amount - delta);
        }
    }

    public class SendRecords implements Callable<Void> {
        private final HashSet<TopicPartition> activePartitions;

        private final Histogram histogram;

        private final Future<?> statusUpdaterFuture;

        private final KafkaProducer<byte[], byte[]> producer;

        private final PayloadIterator keys;

        private final PayloadIterator values;

        private final Optional<TransactionGenerator> transactionGenerator;

        private final Throttle throttle;

        private Iterator<TopicPartition> partitionsIterator;
        private Future<RecordMetadata> sendFuture;
        private AtomicLong transactionsCommitted;
        private boolean enableTransactions;

        SendRecords(HashSet<TopicPartition> activePartitions) {
            this.activePartitions = activePartitions;
            this.partitionsIterator = activePartitions.iterator();
            this.histogram = new Histogram(10000);

            this.transactionGenerator = spec.transactionGenerator();
            this.enableTransactions = this.transactionGenerator.isPresent();
            this.transactionsCommitted = new AtomicLong();

            int perPeriod = WorkerUtils.perSecToPerPeriod(spec.targetMessagesPerSec(), THROTTLE_PERIOD_MS);
            this.statusUpdaterFuture = executor.scheduleWithFixedDelay(
                new StatusUpdater(histogram, transactionsCommitted), 30, 30, TimeUnit.SECONDS);

            Properties props = new Properties();
            props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, spec.bootstrapServers());
            if (enableTransactions)
                props.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "produce-bench-transaction-id-" + UUID.randomUUID());
            // add common client configs to producer properties, and then user-specified producer configs
            WorkerUtils.addConfigsToProperties(props, spec.commonClientConf(), spec.producerConf());
            this.producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer());
            this.keys = new PayloadIterator(spec.keyGenerator());
            this.values = new PayloadIterator(spec.valueGenerator());
            if (spec.skipFlush()) {
                this.throttle = new Throttle(perPeriod, THROTTLE_PERIOD_MS);
            } else {
                this.throttle = new SendRecordsThrottle(perPeriod, producer);
            }
        }

        @Override
        public Void call() throws Exception {
            long startTimeMs = Time.SYSTEM.milliseconds();
            try {
                try {
                    if (enableTransactions)
                        producer.initTransactions();

                    long sentMessages = 0;
                    while (sentMessages < spec.maxMessages()) {
                        if (enableTransactions) {
                            boolean tookAction = takeTransactionAction();
                            if (tookAction)
                                continue;
                        }
                        sendMessage();
                        sentMessages++;
                    }
                    if (enableTransactions)
                        takeTransactionAction(); // give the transactionGenerator a chance to commit if configured evenly
                } catch (Exception e) {
                    if (enableTransactions)
                        producer.abortTransaction();
                    throw e;
                } finally {
                    if (sendFuture != null) {
                        try {
                            sendFuture.get();
                        } catch (Exception e) {
                            log.error("Exception on final future", e);
                        }
                    }
                    producer.close();
                }
            } catch (Exception e) {
                WorkerUtils.abort(log, "SendRecords", e, doneFuture);
            } finally {
                statusUpdaterFuture.cancel(false);
                StatusData statusData = new StatusUpdater(histogram, transactionsCommitted).update();
                long curTimeMs = Time.SYSTEM.milliseconds();
                log.info("Sent {} total record(s) in {} ms.  status: {}",
                    histogram.summarize().numSamples(), curTimeMs - startTimeMs, statusData);
            }
            doneFuture.complete("");
            return null;
        }

        private boolean takeTransactionAction() {
            boolean tookAction = true;
            TransactionAction nextAction = transactionGenerator.get().nextAction();
            switch (nextAction) {
                case BEGIN_TRANSACTION:
                    log.debug("Beginning transaction.");
                    producer.beginTransaction();
                    break;
                case COMMIT_TRANSACTION:
                    log.debug("Committing transaction.");
                    producer.commitTransaction();
                    transactionsCommitted.getAndIncrement();
                    break;
                case ABORT_TRANSACTION:
                    log.debug("Aborting transaction.");
                    producer.abortTransaction();
                    break;
                case NO_OP:
                    tookAction = false;
                    break;
            }
            return tookAction;
        }

        private void sendMessage() throws InterruptedException {
            if (!partitionsIterator.hasNext())
                partitionsIterator = activePartitions.iterator();

            TopicPartition partition = partitionsIterator.next();
            ProducerRecord<byte[], byte[]> record;
            if (spec.useConfiguredPartitioner()) {
                record = new ProducerRecord<>(
                    partition.topic(), keys.next(), values.next());
            } else {
                record = new ProducerRecord<>(
                    partition.topic(), partition.partition(), keys.next(), values.next());
            }
            sendFuture = producer.send(record,
                new SendRecordsCallback(this, Time.SYSTEM.milliseconds()));
            throttle.increment();
        }

        void recordDuration(long durationMs) {
            histogram.add(durationMs);
        }
    }

    public class StatusUpdater implements Runnable {
        private final Histogram histogram;
        private final AtomicLong transactionsCommitted;

        StatusUpdater(Histogram histogram, AtomicLong transactionsCommitted) {
            this.histogram = histogram;
            this.transactionsCommitted = transactionsCommitted;
        }

        @Override
        public void run() {
            try {
                update();
            } catch (Exception e) {
                WorkerUtils.abort(log, "StatusUpdater", e, doneFuture);
            }
        }

        StatusData update() {
            Histogram.Summary summary = histogram.summarize(StatusData.PERCENTILES);
            StatusData statusData = new StatusData(summary.numSamples(), summary.average(),
                summary.percentiles().get(0).value(),
                summary.percentiles().get(1).value(),
                summary.percentiles().get(2).value(),
                transactionsCommitted.get());
            status.update(JsonUtil.JSON_SERDE.valueToTree(statusData));
            return statusData;
        }
    }

    public static class StatusData {
        private final long totalSent;
        private final float averageLatencyMs;
        private final int p50LatencyMs;
        private final int p95LatencyMs;
        private final int p99LatencyMs;
        private final long transactionsCommitted;

        /**
         * The percentiles to use when calculating the histogram data.
         * These should match up with the p50LatencyMs, p95LatencyMs, etc. fields.
         */
        final static float[] PERCENTILES = {0.5f, 0.95f, 0.99f};

        @JsonCreator
        StatusData(@JsonProperty("totalSent") long totalSent,
                   @JsonProperty("averageLatencyMs") float averageLatencyMs,
                   @JsonProperty("p50LatencyMs") int p50latencyMs,
                   @JsonProperty("p95LatencyMs") int p95latencyMs,
                   @JsonProperty("p99LatencyMs") int p99latencyMs,
                   @JsonProperty("transactionsCommitted") long transactionsCommitted) {
            this.totalSent = totalSent;
            this.averageLatencyMs = averageLatencyMs;
            this.p50LatencyMs = p50latencyMs;
            this.p95LatencyMs = p95latencyMs;
            this.p99LatencyMs = p99latencyMs;
            this.transactionsCommitted = transactionsCommitted;
        }

        @JsonProperty
        public long totalSent() {
            return totalSent;
        }

        @JsonProperty
        public long transactionsCommitted() {
            return transactionsCommitted;
        }

        @JsonProperty
        public float averageLatencyMs() {
            return averageLatencyMs;
        }

        @JsonProperty
        public int p50LatencyMs() {
            return p50LatencyMs;
        }

        @JsonProperty
        public int p95LatencyMs() {
            return p95LatencyMs;
        }

        @JsonProperty
        public int p99LatencyMs() {
            return p99LatencyMs;
        }
    }

    @Override
    public void stop(Platform platform) throws Exception {
        if (!running.compareAndSet(true, false)) {
            throw new IllegalStateException("ProduceBenchWorker is not running.");
        }
        log.info("{}: Deactivating ProduceBenchWorker.", id);
        doneFuture.complete("");
        executor.shutdownNow();
        executor.awaitTermination(1, TimeUnit.DAYS);
        this.executor = null;
        this.status = null;
        this.doneFuture = null;
    }
}

相关信息

kafka 源码目录

相关文章

kafka ConfigurableProducerSpec 源码

kafka ConfigurableProducerWorker 源码

kafka ConnectionStressSpec 源码

kafka ConnectionStressWorker 源码

kafka ConstantFlushGenerator 源码

kafka ConstantPayloadGenerator 源码

kafka ConstantThroughputGenerator 源码

kafka ConsumeBenchSpec 源码

kafka ConsumeBenchWorker 源码

kafka ExternalCommandSpec 源码