co.easimart.vertx.stream.InputStreamPump Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of vertx-util Show documentation
Library provides utility classes for Vertx project development.
The newest version!
package co.easimart.vertx.stream;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.vertx.core.AsyncResult;
import io.vertx.core.Future;
import io.vertx.core.Handler;
import io.vertx.core.Vertx;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.streams.WriteStream;

import static co.easimart.vertx.util.VertxHelper.ifSucceeded;

/**
 * Class which pump data from input stream to WriteStream<Buffer>
 */
public class InputStreamPump {

    private static final Logger logger = LoggerFactory.getLogger(InputStreamPump.class);

    private final Vertx vertx;
    private final WriteStream output;
    private final int MAX_BUFFER_SIZE = 4 * 1024;
    private long totalPumpedSize = 0;
    private Throwable canceled = null;

    public InputStreamPump(Vertx vertx, WriteStream output) {
        this.vertx = vertx;
        this.output = output;
        this.output.exceptionHandler(t -> this.canceled = t);
    }

    private void close(InputStream input) {
        try {
            input.close();
        } catch (IOException e) {
            logger.warn("Fail to close input stream after pump to write stream", e);
        }
    }

    public void pump(InputStream input, Handler> handler) {
        if (canceled != null) { // if has been canceled
            handler.handle(Future.failedFuture(this.canceled));

            // close the input stream
            this.vertx.executeBlocking(f -> close(input), false, res -> {
            });
            return;
        } else if (this.output.writeQueueFull()) { // if write stream is full
            // waiting for drain handler to pump again
            this.output.drainHandler(r -> pump(input, handler));
            return;
        }

        // Try read from input stream
        this.vertx.executeBlocking((Future future) -> {
            try {
                // Read to buffer
                ByteBuf buf = Unpooled.buffer(MAX_BUFFER_SIZE);
                int hasRead = buf.setBytes(0, input, MAX_BUFFER_SIZE);
                if (hasRead < 0) {
                    //end of input stream
                    future.complete(null);
                    close(input);
                } else {
                    buf.setIndex(0, hasRead);
                    future.complete(Buffer.buffer(buf));
                }
            } catch (IOException e) {
                future.fail(e);
                // close the input stream
                close(input);
            }
        }, false, ifSucceeded(buf -> {
            if (buf == null) {
                handler.handle(Future.succeededFuture(this.totalPumpedSize));
            } else {
                this.totalPumpedSize += buf.getByteBuf().readableBytes();
                this.output.write(buf);
                pump(input, handler);
            }
        }, handler));
    }
}