com.google.pubsub.clients.common.Task.java Source code

Java tutorial

Introduction

Here is the source code for com.google.pubsub.clients.common.Task.java

Source

// Copyright 2016 Google Inc.
//
// Licensed 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 com.google.pubsub.clients.common;

import com.google.common.base.Preconditions;
import com.google.common.base.Stopwatch;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.RateLimiter;
import com.google.protobuf.util.Timestamps;
import com.google.pubsub.flic.common.LoadtestProto.MessageIdentifier;
import com.google.pubsub.flic.common.LoadtestProto.StartRequest;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Each task is responsible for implementing its action and for creating {@link LoadTestRunner}.
 */
public abstract class Task implements Runnable {
    private final MetricsHandler metricsHandler;
    private AtomicInteger numMessages = new AtomicInteger(0);
    private final Map<MessageIdentifier, Long> identifiers = new HashMap<>();
    private final Map<MessageIdentifier, Long> identifiersToRecord = new HashMap<>();
    private final AtomicLong lastUpdateMillis = new AtomicLong(System.currentTimeMillis());
    private final long burnInTimeMillis;
    private final RateLimiter rateLimiter;
    private final Semaphore outstandingRequestLimiter;

    protected Task(StartRequest request, String type, MetricsHandler.MetricName metricName) {
        this.metricsHandler = new MetricsHandler(request.getProject(), type, metricName);
        this.burnInTimeMillis = Timestamps
                .toMillis(Timestamps.add(request.getStartTime(), request.getBurnInDuration()));
        rateLimiter = RateLimiter.create(request.getRequestRate());
        outstandingRequestLimiter = new Semaphore(request.getMaxOutstandingRequests(), false);
    }

    /** Stores the results of a single doRun call. */
    public static class RunResult {
        // If 'identifiers' and 'latencies' are both populated, the latencies[i] is the latency for
        // message identifiers[i]. If only 'latencies' is populated, we will not be able to do
        // deduplication or completeness checking, and will skip the deduplication logic. For example,
        // publisher tasks may choose to not populate 'identifiers'.
        public List<MessageIdentifier> identifiers = new ArrayList<>();
        public List<Long> latencies = new ArrayList<>();
        public int batchSize = 0;

        public void addMessageLatency(int clientId, int sequenceNumber, long latency) {
            identifiers.add(MessageIdentifier.newBuilder().setPublisherClientId(clientId)
                    .setSequenceNumber(sequenceNumber).build());
            latencies.add(latency);
        }

        public static RunResult fromBatchSize(int batchSize) {
            RunResult result = new RunResult();
            result.batchSize = batchSize;
            return result;
        }

        public static RunResult fromMessages(List<MessageIdentifier> identifiers, List<Long> latencies) {
            RunResult result = new RunResult();
            result.identifiers = identifiers;
            result.latencies = latencies;
            return result;
        }

        public static RunResult empty() {
            return new RunResult();
        }
    }

    public abstract ListenableFuture<RunResult> doRun();

    @Override
    public void run() {
        try {
            if (!outstandingRequestLimiter.tryAcquire(250, TimeUnit.MILLISECONDS)) {
                return;
            }
        } catch (InterruptedException e) {
            return;
        }
        if (!rateLimiter.tryAcquire(250, TimeUnit.MILLISECONDS)) {
            outstandingRequestLimiter.release();
            return;
        }

        Stopwatch stopwatch = Stopwatch.createStarted();
        Futures.addCallback(doRun(), new FutureCallback<RunResult>() {
            @Override
            public void onSuccess(RunResult result) {
                stopwatch.stop();
                outstandingRequestLimiter.release();
                if (result.batchSize > 0) {
                    recordBatchLatency(stopwatch.elapsed(TimeUnit.MILLISECONDS), result.batchSize);
                    return;
                }
                if (result.identifiers.isEmpty()) {
                    // Because result.identifiers is not populated, we are not checking for
                    // duplicates, so just record the latencies immediately.
                    result.latencies.forEach(l -> recordLatency(l));
                    return;
                }
                recordAllMessageLatencies(result.identifiers, result.latencies);
            }

            @Override
            public void onFailure(Throwable t) {
                outstandingRequestLimiter.release();
            }
        });
    }

    List<Long> getBucketValues() {
        return metricsHandler.flushBucketValues();
    }

    int getNumberOfMessages() {
        return numMessages.get();
    }

    private void recordLatency(long millis) {
        recordBatchLatency(millis, 1);
    }

    protected void recordBatchLatency(long millis, int batchSize) {
        lastUpdateMillis.set(System.currentTimeMillis());
        if (System.currentTimeMillis() < burnInTimeMillis) {
            return;
        }
        metricsHandler.recordBatchLatency(millis, batchSize);
        numMessages.getAndAdd(batchSize);
    }

    long getLastUpdateMillis() {
        return lastUpdateMillis.get();
    }

    protected synchronized void recordMessageLatency(int clientId, int sequenceNumber, long latency) {
        identifiers.put(MessageIdentifier.newBuilder().setPublisherClientId(clientId)
                .setSequenceNumber(sequenceNumber).build(), latency);
        lastUpdateMillis.set(System.currentTimeMillis());
    }

    protected synchronized void recordAllMessageLatencies(List<MessageIdentifier> identifiers,
            List<Long> latencies) {
        Preconditions.checkArgument(identifiers.size() == latencies.size(),
                "Identifiers and latencies must be the same size (%s != %s).", identifiers.size(),
                latencies.size());
        for (int i = 0; i < identifiers.size(); i++) {
            this.identifiers.put(identifiers.get(i), latencies.get(i));
        }
        lastUpdateMillis.set(System.currentTimeMillis());
    }

    synchronized List<MessageIdentifier> flushMessageIdentifiers(List<MessageIdentifier> duplicates) {
        identifiersToRecord.forEach((identifier, latency) -> {
            if (!duplicates.contains(identifier)) {
                recordLatency(latency);
            }
        });
        identifiersToRecord.clear();
        identifiersToRecord.putAll(identifiers);
        identifiers.clear();
        return new ArrayList<>(identifiersToRecord.keySet());
    }

    protected void shutdown() {
    }
}