gobblin.util.limiter.stressTest.RateComputingLimiterContainer.java Source code

Introduction

Here is the source code for gobblin.util.limiter.stressTest.RateComputingLimiterContainer.java

Source

/*
 * 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 gobblin.util.limiter.stressTest;

import java.io.Closeable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.commons.math3.stat.descriptive.DescriptiveStatistics;

import com.google.common.base.Function;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import gobblin.util.Decorator;
import gobblin.util.limiter.Limiter;
import gobblin.util.limiter.RestliServiceBasedLimiter;

import javax.annotation.Nullable;
import lombok.AccessLevel;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * Used to compute statistics on a set of {@link Limiter}s used in a throttling service stress test.
 */
@Slf4j
@RequiredArgsConstructor
public class RateComputingLimiterContainer {
    private final List<AtomicLong> subLimiterPermitCounts = Lists.newArrayList();
    private final Queue<Long> unusedPermitsCounts = new LinkedList<>();

    private Map<String, Long> lastReportTimes = Maps.newHashMap();

    /**
     * Decorate a {@link Limiter} to measure its permit rate for statistics computation.
     */
    public Limiter decorateLimiter(Limiter limiter) {
        AtomicLong localCount = new AtomicLong();
        return new RateComputingLimiterDecorator(limiter, localCount);
    }

    /**
     * A {@link Limiter} decorator that records all permits granted.
     */
    @RequiredArgsConstructor(access = AccessLevel.PRIVATE)
    public class RateComputingLimiterDecorator implements Limiter, Decorator {

        private final Limiter underlying;
        private final AtomicLong localPermitCount;

        @Override
        public Object getDecoratedObject() {
            return this.underlying;
        }

        @Override
        public void start() {
            this.underlying.start();
            RateComputingLimiterContainer.this.subLimiterPermitCounts.add(this.localPermitCount);
        }

        @Override
        public Closeable acquirePermits(long permits) throws InterruptedException {
            Closeable closeable = this.underlying.acquirePermits(permits);
            this.localPermitCount.addAndGet(permits);
            return closeable;
        }

        @Override
        public void stop() {
            this.underlying.stop();
            if (this.underlying instanceof RestliServiceBasedLimiter) {
                RestliServiceBasedLimiter restliLimiter = (RestliServiceBasedLimiter) this.underlying;
                RateComputingLimiterContainer.this.unusedPermitsCounts.add(restliLimiter.getUnusedPermits());
                log.info("Unused permits: " + restliLimiter.getUnusedPermits());
            }
            RateComputingLimiterContainer.this.subLimiterPermitCounts.remove(this.localPermitCount);
        }
    }

    /**
     * Get a {@link DescriptiveStatistics} object with the rate of permit granting for all {@link Limiter}s decorated
     * with this {@link RateComputingLimiterContainer}.
     */
    public @Nullable DescriptiveStatistics getRateStatsSinceLastReport() {
        return getNormalizedStatistics("seenQPS",
                Lists.transform(this.subLimiterPermitCounts, new Function<AtomicLong, Double>() {
                    @Override
                    public Double apply(AtomicLong atomicLong) {
                        return (double) atomicLong.getAndSet(0);
                    }
                }));
    }

    public @Nullable DescriptiveStatistics getUnusedPermitsSinceLastReport() {
        DescriptiveStatistics stats = getNormalizedStatistics("unusedPermits", this.unusedPermitsCounts);
        this.unusedPermitsCounts.clear();
        return stats;
    }

    private @Nullable DescriptiveStatistics getNormalizedStatistics(String key,
            Collection<? extends Number> values) {
        long now = System.currentTimeMillis();

        long deltaTime = 0;
        if (this.lastReportTimes.containsKey(key)) {
            deltaTime = now - this.lastReportTimes.get(key);
        }
        this.lastReportTimes.put(key, now);

        if (deltaTime == 0) {
            return null;
        }

        double[] normalizedValues = new double[values.size()];
        int i = 0;
        for (Number value : values) {
            normalizedValues[i++] = 1000 * value.doubleValue() / deltaTime;
        }

        return new DescriptiveStatistics(normalizedValues);
    }
}