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

Go to download
package co.easimart.vertx.stream;

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

import java.util.Iterator;
import java.util.LinkedList;
import java.util.function.Function;

import io.vertx.core.Context;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.streams.ReadStream;

/**
 * Builder which build multiple ReadStream based on one input ReadStream.
 * Each data from the input ReadStream will be copied or sent to all created TeeReadStream.
 * It works like a T shape pipe.
 */
public class TeeReadStreamBuilder {
    private static final Logger logger = LoggerFactory.getLogger(TeeReadStreamBuilder.class);

    private final ReadStream inputStream;
    private final LinkedList> teeStreams;
    private Function copier;

    public TeeReadStreamBuilder(ReadStream inputStream) {
        this(inputStream, null);
    }

    public TeeReadStreamBuilder(ReadStream inputStream, Function copier) {
        this.inputStream = inputStream;
        this.teeStreams = new LinkedList<>();
        setCopier(copier);
        setupHandlers();
    }

    public void setCopier(Function copier) {
        if (copier == null) {
            this.copier = x -> {
                if (x instanceof Buffer) //noinspection unchecked
                    return (T) ((Buffer) x).copy();
                else return (T) x;
            };
        } else {
            this.copier = copier;
        }
    }

    private void setupHandlers() {
        this.inputStream.handler(data -> {
            Iterator> itr = teeStreams.descendingIterator();
            while (itr.hasNext()) {
                ControllableReadStream stream = itr.next();
                if (itr.hasNext()) {
                    stream.feed(this.copier.apply(data));
                } else {
                    stream.feed(data);
                }
            }
        });
        this.inputStream.exceptionHandler(cause -> teeStreams.forEach(stream -> stream.error(cause)));
        this.inputStream.endHandler(r -> teeStreams.forEach(ControllableReadStream::end));
    }

    public ReadStream createTee(Context context) {
        ControllableReadStream tee = new ControllableReadStream<>(context);
        teeStreams.add(tee);
        return tee;
    }

}