gobblin.restli.throttling.LocalStressTest.java Source code

Java tutorial

Introduction

Here is the source code for gobblin.restli.throttling.LocalStressTest.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.restli.throttling;

import java.util.LinkedList;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;
import org.apache.commons.math3.stat.descriptive.DescriptiveStatistics;
import org.apache.hadoop.conf.Configuration;

import com.google.common.collect.Maps;
import com.google.inject.Key;
import com.google.inject.name.Names;
import com.typesafe.config.ConfigFactory;

import gobblin.broker.BrokerConfigurationKeyGenerator;
import gobblin.broker.iface.SharedResourcesBroker;
import gobblin.util.limiter.Limiter;
import gobblin.util.limiter.MockRequester;
import gobblin.util.limiter.RestliServiceBasedLimiter;
import gobblin.util.limiter.broker.SharedLimiterKey;
import gobblin.util.limiter.stressTest.RateComputingLimiterContainer;
import gobblin.util.limiter.stressTest.StressTestUtils;
import gobblin.util.limiter.stressTest.Stressor;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * A stress test for throttling service. It creates a number of threads, each one running a stressor using a mock
 * {@link RestliServiceBasedLimiter}.
 *
 * The mock {@link RestliServiceBasedLimiter} sends requests to an embedded {@link LimiterServerResource}, adding an
 * artificial latency to the requests representing the network latency.
 *
 * The stress test prints permit granting statistics every 15 seconds.
 */
@Slf4j
public class LocalStressTest {

    public static final Option STRESSOR_THREADS = new Option("stressorThreads", true, "Number of stressor threads");
    public static final Option PROCESSOR_THREADS = new Option("processorThreads", true,
            "Number of request processor threads.");
    public static final Option ARTIFICIAL_LATENCY = new Option("latency", true,
            "Artificial request latency in millis.");
    public static final Option QPS = new Option("qps", true, "Target qps.");

    public static final Options OPTIONS = StressTestUtils.OPTIONS.addOption(STRESSOR_THREADS)
            .addOption(PROCESSOR_THREADS);

    public static final int DEFAULT_STRESSOR_THREADS = 10;
    public static final int DEFAULT_PROCESSOR_THREADS = 10;
    public static final int DEFAULT_ARTIFICIAL_LATENCY = 100;
    public static final int DEFAULT_TARGET_QPS = 100;

    public static void main(String[] args) throws Exception {

        CommandLine cli = StressTestUtils.parseCommandLine(OPTIONS, args);

        int stressorThreads = Integer.parseInt(
                cli.getOptionValue(STRESSOR_THREADS.getOpt(), Integer.toString(DEFAULT_STRESSOR_THREADS)));
        int processorThreads = Integer.parseInt(
                cli.getOptionValue(PROCESSOR_THREADS.getOpt(), Integer.toString(DEFAULT_PROCESSOR_THREADS)));
        int artificialLatency = Integer.parseInt(
                cli.getOptionValue(ARTIFICIAL_LATENCY.getOpt(), Integer.toString(DEFAULT_ARTIFICIAL_LATENCY)));
        long targetQps = Integer.parseInt(cli.getOptionValue(QPS.getOpt(), Integer.toString(DEFAULT_TARGET_QPS)));

        Configuration configuration = new Configuration();
        StressTestUtils.populateConfigFromCli(configuration, cli);

        String resourceLimited = LocalStressTest.class.getSimpleName();

        Map<String, String> configMap = Maps.newHashMap();

        ThrottlingPolicyFactory factory = new ThrottlingPolicyFactory();
        SharedLimiterKey res1key = new SharedLimiterKey(resourceLimited);
        configMap.put(BrokerConfigurationKeyGenerator.generateKey(factory, res1key, null,
                ThrottlingPolicyFactory.POLICY_KEY), QPSPolicy.FACTORY_ALIAS);
        configMap.put(BrokerConfigurationKeyGenerator.generateKey(factory, res1key, null, QPSPolicy.QPS),
                Long.toString(targetQps));

        ThrottlingGuiceServletConfig guiceServletConfig = new ThrottlingGuiceServletConfig();
        guiceServletConfig.initialize(ConfigFactory.parseMap(configMap));
        LimiterServerResource limiterServer = guiceServletConfig.getInjector()
                .getInstance(LimiterServerResource.class);

        RateComputingLimiterContainer limiterContainer = new RateComputingLimiterContainer();

        Class<? extends Stressor> stressorClass = configuration.getClass(StressTestUtils.STRESSOR_CLASS,
                StressTestUtils.DEFAULT_STRESSOR_CLASS, Stressor.class);

        ExecutorService executorService = Executors.newFixedThreadPool(stressorThreads);

        SharedResourcesBroker broker = guiceServletConfig.getInjector().getInstance(
                Key.get(SharedResourcesBroker.class, Names.named(LimiterServerResource.BROKER_INJECT_NAME)));
        ThrottlingPolicy policy = (ThrottlingPolicy) broker.getSharedResource(new ThrottlingPolicyFactory(),
                new SharedLimiterKey(resourceLimited));
        ScheduledExecutorService reportingThread = Executors.newSingleThreadScheduledExecutor();
        reportingThread.scheduleAtFixedRate(new Reporter(limiterContainer, policy), 0, 15, TimeUnit.SECONDS);

        Queue<Future<?>> futures = new LinkedList<>();
        MockRequester requester = new MockRequester(limiterServer, artificialLatency, processorThreads);

        requester.start();
        for (int i = 0; i < stressorThreads; i++) {
            RestliServiceBasedLimiter restliLimiter = RestliServiceBasedLimiter.builder()
                    .resourceLimited(resourceLimited).requestSender(requester).serviceIdentifier("stressor" + i)
                    .build();

            Stressor stressor = stressorClass.newInstance();
            stressor.configure(configuration);
            futures.add(executorService
                    .submit(new StressorRunner(limiterContainer.decorateLimiter(restliLimiter), stressor)));
        }
        int stressorFailures = 0;
        for (Future<?> future : futures) {
            try {
                future.get();
            } catch (ExecutionException ee) {
                stressorFailures++;
            }
        }
        requester.stop();

        executorService.shutdownNow();

        if (stressorFailures > 0) {
            log.error("There were " + stressorFailures + " failed stressor threads.");
        }
        System.exit(stressorFailures);
    }

    @RequiredArgsConstructor
    private static class StressorRunner implements Runnable {
        private final Limiter limiter;
        private final Stressor stressor;

        @Override
        public void run() {
            try {
                this.limiter.start();
                this.stressor.run(this.limiter);
                this.limiter.stop();
            } catch (InterruptedException ie) {
                log.error("Error: ", ie);
            }
        }
    }

    @RequiredArgsConstructor
    private static class Reporter implements Runnable {
        private final RateComputingLimiterContainer limiter;
        private final ThrottlingPolicy policy;

        @Override
        public void run() {
            DescriptiveStatistics stats = limiter.getRateStatsSinceLastReport();
            if (stats != null) {
                log.info(String.format(
                        "Requests rate stats: count: %d, min: %f, max: %f, mean: %f, std: %f, sum: %f",
                        stats.getN(), stats.getMin(), stats.getMax(), stats.getMean(), stats.getStandardDeviation(),
                        stats.getSum()));
            }

            stats = limiter.getUnusedPermitsSinceLastReport();
            if (stats != null) {
                log.info(String.format(
                        "Unused permits rate stats: count: %d, min: %f, max: %f, mean: %f, std: %f, sum: %f",
                        stats.getN(), stats.getMin(), stats.getMax(), stats.getMean(), stats.getStandardDeviation(),
                        stats.getSum()));
            }

            if (this.policy instanceof QPSPolicy) {
                QPSPolicy qpsPolicy = (QPSPolicy) this.policy;
                DynamicTokenBucket dynamicTokenBucket = qpsPolicy.getTokenBucket();
                TokenBucket tokenBucket = dynamicTokenBucket.getTokenBucket();
                log.info("Stored tokens: " + tokenBucket.getStoredTokens());
            }
        }
    }
}