com.inmobi.conduit.distcp.tools.util.TestThrottledInputStream.java Source code

Introduction

Here is the source code for com.inmobi.conduit.distcp.tools.util.TestThrottledInputStream.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 com.inmobi.conduit.distcp.tools.util;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.io.IOUtils;
import org.junit.Assert;
import org.junit.Test;

import com.inmobi.conduit.distcp.tools.util.ThrottledInputStream;

import java.io.*;

public class TestThrottledInputStream {
    private static final Log LOG = LogFactory.getLog(TestThrottledInputStream.class);
    private static final int BUFF_SIZE = 1024;

    private enum CB {
        ONE_C, BUFFER, BUFF_OFFSET
    }

    // @Test
    public void testRead() {
        File tmpFile;
        File outFile;
        try {
            tmpFile = createFile(1024);
            outFile = createFile();

            tmpFile.deleteOnExit();
            outFile.deleteOnExit();

            long maxBandwidth = copyAndAssert(tmpFile, outFile, 0, 1, -1, CB.BUFFER);

            copyAndAssert(tmpFile, outFile, maxBandwidth, 20, 0, CB.BUFFER);
            /*
                  copyAndAssert(tmpFile, outFile, maxBandwidth, 10, 0, CB.BUFFER);
                  copyAndAssert(tmpFile, outFile, maxBandwidth, 50, 0, CB.BUFFER);
            */

            copyAndAssert(tmpFile, outFile, maxBandwidth, 20, 0, CB.BUFF_OFFSET);
            /*
                  copyAndAssert(tmpFile, outFile, maxBandwidth, 10, 0, CB.BUFF_OFFSET);
                  copyAndAssert(tmpFile, outFile, maxBandwidth, 50, 0, CB.BUFF_OFFSET);
            */

            copyAndAssert(tmpFile, outFile, maxBandwidth, 20, 0, CB.ONE_C);
            /*
                  copyAndAssert(tmpFile, outFile, maxBandwidth, 10, 0, CB.ONE_C);
                  copyAndAssert(tmpFile, outFile, maxBandwidth, 50, 0, CB.ONE_C);
            */
        } catch (IOException e) {
            LOG.error("Exception encountered ", e);
        }
    }

    private long copyAndAssert(File tmpFile, File outFile, long maxBandwidth, float factor, int sleepTime, CB flag)
            throws IOException {
        long bandwidth;
        ThrottledInputStream in;
        long maxBPS = (long) (maxBandwidth / factor);

        if (maxBandwidth == 0) {
            in = new ThrottledInputStream(new FileInputStream(tmpFile));
        } else {
            in = new ThrottledInputStream(new FileInputStream(tmpFile), maxBPS);
        }
        OutputStream out = new FileOutputStream(outFile);
        try {
            if (flag == CB.BUFFER) {
                copyBytes(in, out, BUFF_SIZE);
            } else if (flag == CB.BUFF_OFFSET) {
                copyBytesWithOffset(in, out, BUFF_SIZE);
            } else {
                copyByteByByte(in, out);
            }

            LOG.info(in);
            bandwidth = in.getBytesPerSec();
            Assert.assertEquals(in.getTotalBytesRead(), tmpFile.length());
            Assert.assertTrue(in.getBytesPerSec() > maxBandwidth / (factor * 1.2));
            Assert.assertTrue(in.getTotalSleepTime() > sleepTime || in.getBytesPerSec() <= maxBPS);
        } finally {
            IOUtils.closeStream(in);
            IOUtils.closeStream(out);
        }
        return bandwidth;
    }

    private static void copyBytesWithOffset(InputStream in, OutputStream out, int buffSize) throws IOException {

        byte buf[] = new byte[buffSize];
        int bytesRead = in.read(buf, 0, buffSize);
        while (bytesRead >= 0) {
            out.write(buf, 0, bytesRead);
            bytesRead = in.read(buf);
        }
    }

    private static void copyByteByByte(InputStream in, OutputStream out) throws IOException {

        int ch = in.read();
        while (ch >= 0) {
            out.write(ch);
            ch = in.read();
        }
    }

    private static void copyBytes(InputStream in, OutputStream out, int buffSize) throws IOException {

        byte buf[] = new byte[buffSize];
        int bytesRead = in.read(buf);
        while (bytesRead >= 0) {
            out.write(buf, 0, bytesRead);
            bytesRead = in.read(buf);
        }
    }

    private File createFile(long sizeInKB) throws IOException {
        File tmpFile = createFile();
        writeToFile(tmpFile, sizeInKB);
        return tmpFile;
    }

    private File createFile() throws IOException {
        return File.createTempFile("tmp", "dat");
    }

    private void writeToFile(File tmpFile, long sizeInKB) throws IOException {
        OutputStream out = new FileOutputStream(tmpFile);
        try {
            byte[] buffer = new byte[1024];
            for (long index = 0; index < sizeInKB; index++) {
                out.write(buffer);
            }
        } finally {
            IOUtils.closeStream(out);
        }
    }
}