org.apache.nifi.stream.io.LeakyBucketStreamThrottler Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of nifi-utils Show documentation
Show all versions of nifi-utils Show documentation
This nifi-utils module should be a general purpose place to store widely
and generally useful functions that any component might want to leverage.
NO DEPENDENCIES should be added. This module is likely to be leveraged by
every extension and should not bring along any other dependencies. The only
dependency intended is the nifi-api and even this is expected to be already
provided in any case where it would be used. The typical place this util
would be found is within a nar and all nars already have nifi-api as a parent
dependency. The nifi-api can be thought of as a NiFi Application Container level
dependency.
/*
* 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 org.apache.nifi.stream.io;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
public class LeakyBucketStreamThrottler implements StreamThrottler {
private final int maxBytesPerSecond;
private final BlockingQueue requestQueue = new LinkedBlockingQueue();
private final ScheduledExecutorService executorService;
private final AtomicBoolean shutdown = new AtomicBoolean(false);
public LeakyBucketStreamThrottler(final int maxBytesPerSecond) {
this.maxBytesPerSecond = maxBytesPerSecond;
executorService = Executors.newSingleThreadScheduledExecutor();
final Runnable task = new Drain();
executorService.scheduleAtFixedRate(task, 0, 1000, TimeUnit.MILLISECONDS);
}
@Override
public void close() {
this.shutdown.set(true);
executorService.shutdown();
try {
// Should not take more than 2 seconds because we run every second. If it takes more than
// 2 seconds, it is because the Runnable thread is blocking on a write; in this case,
// we will just ignore it and return
executorService.awaitTermination(2, TimeUnit.SECONDS);
} catch (InterruptedException e) {
}
}
@Override
public OutputStream newThrottledOutputStream(final OutputStream toWrap) {
return new OutputStream() {
@Override
public void write(final int b) throws IOException {
write(new byte[]{(byte) b}, 0, 1);
}
@Override
public void write(byte[] b) throws IOException {
write(b, 0, b.length);
}
@Override
public void write(byte[] b, int off, int len) throws IOException {
final InputStream in = new ByteArrayInputStream(b, off, len);
LeakyBucketStreamThrottler.this.copy(in, toWrap);
}
@Override
public void close() throws IOException {
toWrap.close();
}
@Override
public void flush() throws IOException {
toWrap.flush();
}
};
}
@Override
public InputStream newThrottledInputStream(final InputStream toWrap) {
return new InputStream() {
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
@Override
public int read() throws IOException {
final ByteArrayOutputStream baos = new ByteArrayOutputStream(1);
LeakyBucketStreamThrottler.this.copy(toWrap, baos, 1L);
if (baos.size() < 1) {
return -1;
}
return baos.toByteArray()[0] & 0xFF;
}
@Override
public int read(final byte[] b) throws IOException {
if (b.length == 0) {
return 0;
}
return read(b, 0, b.length);
}
@Override
public int read(byte[] b, int off, int len) throws IOException {
if (len < 0) {
throw new IllegalArgumentException();
}
if (len == 0) {
return 0;
}
baos.reset();
final int copied = (int) LeakyBucketStreamThrottler.this.copy(toWrap, baos, len);
if (copied == 0) {
return -1;
}
System.arraycopy(baos.toByteArray(), 0, b, off, copied);
return copied;
}
@Override
public void close() throws IOException {
toWrap.close();
}
@Override
public int available() throws IOException {
return toWrap.available();
}
};
}
@Override
public long copy(final InputStream in, final OutputStream out) throws IOException {
return copy(in, out, -1);
}
@Override
public long copy(final InputStream in, final OutputStream out, final long maxBytes) throws IOException {
long totalBytesCopied = 0;
boolean finished = false;
while (!finished) {
final long requestMax = (maxBytes < 0) ? Long.MAX_VALUE : maxBytes - totalBytesCopied;
final Request request = new Request(in, out, requestMax);
boolean transferred = false;
while (!transferred) {
if (shutdown.get()) {
throw new IOException("Throttler shutdown");
}
try {
transferred = requestQueue.offer(request, 1000, TimeUnit.MILLISECONDS);
} catch (final InterruptedException e) {
throw new IOException("Interrupted", e);
}
}
final BlockingQueue responseQueue = request.getResponseQueue();
Response response = null;
while (response == null) {
try {
if (shutdown.get()) {
throw new IOException("Throttler shutdown");
}
response = responseQueue.poll(1000L, TimeUnit.MILLISECONDS);
} catch (InterruptedException e) {
throw new IOException("Interrupted", e);
}
}
if (!response.isSuccess()) {
throw response.getError();
}
totalBytesCopied += response.getBytesCopied();
finished = (response.getBytesCopied() == 0) || (totalBytesCopied >= maxBytes && maxBytes > 0);
}
return totalBytesCopied;
}
/**
* This class is responsible for draining water from the leaky bucket. I.e., it actually moves the data
*/
private class Drain implements Runnable {
private final byte[] buffer;
public Drain() {
final int bufferSize = Math.min(4096, maxBytesPerSecond);
buffer = new byte[bufferSize];
}
@Override
public void run() {
final long start = System.currentTimeMillis();
int bytesTransferred = 0;
while (bytesTransferred < maxBytesPerSecond) {
final long maxMillisToWait = 1000 - (System.currentTimeMillis() - start);
if (maxMillisToWait < 1) {
return;
}
try {
final Request request = requestQueue.poll(maxMillisToWait, TimeUnit.MILLISECONDS);
if (request == null) {
return;
}
final BlockingQueue responseQueue = request.getResponseQueue();
final OutputStream out = request.getOutputStream();
final InputStream in = request.getInputStream();
try {
final long requestMax = request.getMaxBytesToCopy();
long maxBytesToTransfer;
if (requestMax < 0) {
maxBytesToTransfer = Math.min(buffer.length, maxBytesPerSecond - bytesTransferred);
} else {
maxBytesToTransfer = Math.min(requestMax,
Math.min(buffer.length, maxBytesPerSecond - bytesTransferred));
}
maxBytesToTransfer = Math.max(1L, maxBytesToTransfer);
final int bytesCopied = fillBuffer(in, maxBytesToTransfer);
out.write(buffer, 0, bytesCopied);
final Response response = new Response(true, bytesCopied);
responseQueue.put(response);
bytesTransferred += bytesCopied;
} catch (final IOException e) {
final Response response = new Response(e);
responseQueue.put(response);
}
} catch (InterruptedException e) {
}
}
}
private int fillBuffer(final InputStream in, final long maxBytes) throws IOException {
int bytesRead = 0;
int len;
while (bytesRead < maxBytes && (len = in.read(buffer, bytesRead, (int) Math.min(maxBytes - bytesRead, buffer.length - bytesRead))) > 0) {
bytesRead += len;
}
return bytesRead;
}
}
private static class Response {
private final boolean success;
private final IOException error;
private final int bytesCopied;
public Response(final boolean success, final int bytesCopied) {
this.success = success;
this.bytesCopied = bytesCopied;
this.error = null;
}
public Response(final IOException error) {
this.success = false;
this.error = error;
this.bytesCopied = -1;
}
public boolean isSuccess() {
return success;
}
public IOException getError() {
return error;
}
public int getBytesCopied() {
return bytesCopied;
}
}
private static class Request {
private final OutputStream out;
private final InputStream in;
private final long maxBytesToCopy;
private final BlockingQueue responseQueue;
public Request(final InputStream in, final OutputStream out, final long maxBytesToCopy) {
this.out = out;
this.in = in;
this.maxBytesToCopy = maxBytesToCopy;
this.responseQueue = new LinkedBlockingQueue(1);
}
public BlockingQueue getResponseQueue() {
return this.responseQueue;
}
public OutputStream getOutputStream() {
return out;
}
public InputStream getInputStream() {
return in;
}
public long getMaxBytesToCopy() {
return maxBytesToCopy;
}
@Override
public String toString() {
return "Request[maxBytes=" + maxBytesToCopy + "]";
}
}
}