com.linkedin.pinot.transport.perf.ScatterGatherPerfServer.java Source code

Introduction

Here is the source code for com.linkedin.pinot.transport.perf.ScatterGatherPerfServer.java

Source

/**
 * Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com)
 *
 * 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.linkedin.pinot.transport.perf;

import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import io.netty.channel.ChannelHandlerContext;
import java.util.concurrent.CountDownLatch;

import io.netty.buffer.ByteBuf;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.Options;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.linkedin.pinot.common.metrics.AggregatedMetricsRegistry;
import com.linkedin.pinot.transport.netty.NettyServer;
import com.linkedin.pinot.transport.netty.NettyServer.RequestHandler;
import com.linkedin.pinot.transport.netty.NettyServer.RequestHandlerFactory;
import com.linkedin.pinot.transport.netty.NettyTCPServer;

/**
 *
 * This class is used for benchmarking the Scatter-Gather Layer
 *
 *
 */
public class ScatterGatherPerfServer {

    private static final Logger LOGGER = LoggerFactory.getLogger(ScatterGatherPerfServer.class);

    /*
    static
    {
      org.apache.log4j.Logger.getRootLogger().addAppender(new ConsoleAppender(
    new PatternLayout(PatternLayout.TTCC_CONVERSION_PATTERN), "System.out"));
      org.apache.log4j.Logger.getRootLogger().setLevel(Level.INFO);
    }
    */

    public static final String RESPONSE_SIZE_OPT_NAME = "response_size";
    public static final String SERVER_PORT_OPT_NAME = "server_port";

    private final int _serverPort;
    private final int _responseSize;
    private final long _responseLatencyMs;

    private byte[] _bakedResponse;

    private NettyTCPServer _server;

    public ScatterGatherPerfServer(int serverPort, int responseSize, long responseLatencyMs) {
        _serverPort = serverPort;
        _responseSize = responseSize;
        _responseLatencyMs = responseLatencyMs;
    }

    public void run() {
        AggregatedMetricsRegistry metricsRegistry = new AggregatedMetricsRegistry();
        _bakedResponse = new byte[_responseSize];
        for (int i = 0; i < _responseSize; i++)
            _bakedResponse[i] = 'a';

        MyRequestHandler handler = new MyRequestHandler(new String(_bakedResponse), null, _responseLatencyMs);
        MyRequestHandlerFactory handlerFactory = new MyRequestHandlerFactory(handler);
        _server = new NettyTCPServer(_serverPort, handlerFactory, metricsRegistry);
        Thread serverThread = new Thread(_server, "ServerMain");
        ShutdownHook shutdownHook = new ShutdownHook(_server);
        Runtime.getRuntime().addShutdownHook(shutdownHook);
        serverThread.start();
    }

    public void shutdown() {
        if (null != _server) {
            _server.shutdownGracefully();
            _server = null;
        }
    }

    private static Options buildCommandLineOptions() {
        Options options = new Options();
        options.addOption(SERVER_PORT_OPT_NAME, true, "Server Port for accepting queries from broker");
        options.addOption(RESPONSE_SIZE_OPT_NAME, true, "Response Size");
        return options;
    }

    /**
     * @param args
     */
    public static void main(String[] args) throws Exception {

        CommandLineParser cliParser = new GnuParser();
        Options cliOptions = buildCommandLineOptions();

        CommandLine cmd = cliParser.parse(cliOptions, args, true);

        if (!cmd.hasOption(RESPONSE_SIZE_OPT_NAME) || !cmd.hasOption(SERVER_PORT_OPT_NAME)) {
            System.err.println("Missing required arguments !!");
            System.err.println(cliOptions);
            throw new RuntimeException("Missing required arguments !!");
        }

        int responseSize = Integer.parseInt(cmd.getOptionValue(RESPONSE_SIZE_OPT_NAME));
        int serverPort = Integer.parseInt(cmd.getOptionValue(SERVER_PORT_OPT_NAME));

        ScatterGatherPerfServer server = new ScatterGatherPerfServer(serverPort, responseSize, 2); // 2ms latency
        server.run();
    }

    private static class MyRequestHandlerFactory implements RequestHandlerFactory {
        private final MyRequestHandler _requestHandler;

        public MyRequestHandlerFactory(MyRequestHandler requestHandler) {
            _requestHandler = requestHandler;
        }

        @Override
        public RequestHandler createNewRequestHandler() {
            return _requestHandler;
        }

    }

    private static class MyRequestHandler implements RequestHandler {
        private String _response;
        private String _request;
        private final CountDownLatch _responseHandlingLatch;
        private final long _responseLatencyMs;

        public MyRequestHandler(String response, CountDownLatch responseHandlingLatch, long responseLatencyMs) {
            _response = response;
            _responseHandlingLatch = responseHandlingLatch;
            _responseLatencyMs = responseLatencyMs;
        }

        @Override
        public ListenableFuture<byte[]> processRequest(ChannelHandlerContext channelHandlerContext,
                ByteBuf request) {
            byte[] b = new byte[request.readableBytes()];
            request.readBytes(b);
            if (null != _responseHandlingLatch) {
                try {
                    _responseHandlingLatch.await();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            _request = new String(b);

            if (_responseLatencyMs > 0) {
                try {
                    Thread.sleep(_responseLatencyMs);
                } catch (InterruptedException e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                }
            }
            //LOG.info("Server got the request (" + _request + ")");
            return Futures.immediateFuture(_response.getBytes());
        }

        public String getRequest() {
            return _request;
        }

        public String getResponse() {
            return _response;
        }

        public void setResponse(String response) {
            _response = response;
        }
    }

    public static class ShutdownHook extends Thread {
        private final NettyServer _server;

        public ShutdownHook(NettyServer server) {
            _server = server;
        }

        @Override
        public void run() {
            LOGGER.info("Running shutdown hook");
            if ((_server != null) && (!_server.isShutdownComplete())) {
                _server.shutdownGracefully();
            }
            LOGGER.info("Shutdown completed !!");
        }
    }
}