com.linkedin.pinotdruidbenchmark.PinotResponseTime.java Source code

Introduction

Here is the source code for com.linkedin.pinotdruidbenchmark.PinotResponseTime.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.pinotdruidbenchmark;

import java.io.BufferedInputStream;
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.Arrays;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;

/**
 * Test single query response time for Pinot.
 */
public class PinotResponseTime {
    private PinotResponseTime() {
    }

    private static final byte[] BYTE_BUFFER = new byte[4096];

    public static void main(String[] args) throws Exception {
        if (args.length != 4 && args.length != 5) {
            System.err.println(
                    "4 or 5 arguments required: QUERY_DIR, RESOURCE_URL, WARM_UP_ROUNDS, TEST_ROUNDS, RESULT_DIR (optional).");
            return;
        }

        File queryDir = new File(args[0]);
        String resourceUrl = args[1];
        int warmUpRounds = Integer.parseInt(args[2]);
        int testRounds = Integer.parseInt(args[3]);
        File resultDir;
        if (args.length == 4) {
            resultDir = null;
        } else {
            resultDir = new File(args[4]);
            if (!resultDir.exists()) {
                if (!resultDir.mkdirs()) {
                    throw new RuntimeException("Failed to create result directory: " + resultDir);
                }
            }
        }

        File[] queryFiles = queryDir.listFiles();
        assert queryFiles != null;
        Arrays.sort(queryFiles);

        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
            HttpPost httpPost = new HttpPost(resourceUrl);

            for (File queryFile : queryFiles) {
                String query = new BufferedReader(new FileReader(queryFile)).readLine();
                httpPost.setEntity(new StringEntity("{\"pql\":\"" + query + "\"}"));

                System.out.println(
                        "--------------------------------------------------------------------------------");
                System.out.println("Running query: " + query);
                System.out.println(
                        "--------------------------------------------------------------------------------");

                // Warm-up Rounds
                System.out.println("Run " + warmUpRounds + " times to warm up...");
                for (int i = 0; i < warmUpRounds; i++) {
                    CloseableHttpResponse httpResponse = httpClient.execute(httpPost);
                    httpResponse.close();
                    System.out.print('*');
                }
                System.out.println();

                // Test Rounds
                System.out.println("Run " + testRounds + " times to get response time statistics...");
                long[] responseTimes = new long[testRounds];
                long totalResponseTime = 0L;
                for (int i = 0; i < testRounds; i++) {
                    long startTime = System.currentTimeMillis();
                    CloseableHttpResponse httpResponse = httpClient.execute(httpPost);
                    httpResponse.close();
                    long responseTime = System.currentTimeMillis() - startTime;
                    responseTimes[i] = responseTime;
                    totalResponseTime += responseTime;
                    System.out.print(responseTime + "ms ");
                }
                System.out.println();

                // Store result.
                if (resultDir != null) {
                    File resultFile = new File(resultDir, queryFile.getName() + ".result");
                    CloseableHttpResponse httpResponse = httpClient.execute(httpPost);
                    try (BufferedInputStream bufferedInputStream = new BufferedInputStream(
                            httpResponse.getEntity().getContent());
                            BufferedWriter bufferedWriter = new BufferedWriter(new FileWriter(resultFile))) {
                        int length;
                        while ((length = bufferedInputStream.read(BYTE_BUFFER)) > 0) {
                            bufferedWriter.write(new String(BYTE_BUFFER, 0, length));
                        }
                    }
                    httpResponse.close();
                }

                // Process response times.
                double averageResponseTime = (double) totalResponseTime / testRounds;
                double temp = 0;
                for (long responseTime : responseTimes) {
                    temp += (responseTime - averageResponseTime) * (responseTime - averageResponseTime);
                }
                double standardDeviation = Math.sqrt(temp / testRounds);
                System.out.println("Average response time: " + averageResponseTime + "ms");
                System.out.println("Standard deviation: " + standardDeviation);
            }
        }
    }
}