co.easimart.vertx.thrift.THttpVertxTransport 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.thrift;

import org.apache.thrift.TException;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URL;

import io.netty.buffer.ByteBuf;
import io.vertx.core.AsyncResult;
import io.vertx.core.Future;
import io.vertx.core.Handler;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.http.HttpClient;
import io.vertx.core.http.HttpClientRequest;

/**
 * Thrift HTTP transport which use Vertx asynchronous http client
 */
public class THttpVertxTransport extends TTransport {

    @FunctionalInterface
    public interface Send {
        void send() throws TException;
    }

    @FunctionalInterface
    public interface Receive {
        R receive() throws TException;
    }

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

    private final HttpClient client;
    private final String userAgent;
    private Buffer reqBuffer;
    private ByteBuf resBuffer;

    public THttpVertxTransport(HttpClient client, String userAgent) {
        this.client = client;
        this.userAgent = userAgent;
    }

    public  void execute(URL url, Send send, Receive receive, Handler> handler) {
        this.reqBuffer = Buffer.buffer(255);

        // Run send command.
        try {
            send.send();
        } catch (Exception e) {
            handler.handle(Future.failedFuture(e));
            return;
        }
        if (this.reqBuffer == null || this.reqBuffer.length() == 0) {
            throw new IllegalStateException("Thrift send operation did not call thrift send methods");
        }

        // Setup HTTP request
        HttpClientRequest request = client.postAbs(url.toString());
        String uri = request.uri();
        logger.debug("Sending thrift request to {}", uri);
        request.putHeader("Content-Type", "application/x-thrift");
        request.putHeader("Accept", "application/x-thrift");
        if (this.userAgent != null) {
            request.putHeader("User-Agent", this.userAgent);
        }
        Future future = Future.future();
        future.setHandler(handler);
        // Handler exception
        request.exceptionHandler(t -> {
            logger.debug("Exception from thrift request to {}", uri, t);
            if (future.isComplete()) return;
            future.fail(t);
        });
        // Handler response
        request.handler(res -> {
            if (res.statusCode() != 200) {
                // Failure response
                future.fail(new RuntimeException("Fail with code " + res.statusCode() + " - " + res.statusMessage()));
                if (logger.isDebugEnabled()) {
                    res.bodyHandler(buffer -> {
                        // debug logs
                        logger.debug("Thrift response from {} (status={}): {}", uri, res.statusCode(), buffer.toString());
                    });
                }
                return;
            }
            this.logger.debug("Headers of Thrift response from {} - {}", uri, res.headers().entries());

            // Success response
            Buffer body = Buffer.buffer();
            res.handler(body::appendBuffer);
            res.endHandler(r -> {
                this.resBuffer = body.getByteBuf();
                logger.debug("Thrift response from {} received and translating", uri);
                // Read from response buffer
                try {
                    future.complete(receive.receive());
                } catch (Exception e) {
                    if (future.isComplete()) return;
                    future.fail(e);
                }
            });
        });

        // Send with request buffer
        request.end(reqBuffer);
    }

    @Override
    public boolean isOpen() {
        return this.reqBuffer != null;
    }

    @Override
    public void open() throws TTransportException {
    }

    @Override
    public void close() {
    }

    @Override
    public int read(byte[] bytes, int off, int len) throws TTransportException {
        if (this.resBuffer.readableBytes() == 0) return -1;
        int canRead = Math.min(len, this.resBuffer.readableBytes());
        this.resBuffer.readBytes(bytes, off, canRead);
        return canRead;
    }

    @Override
    public void write(byte[] bytes, int off, int len) throws TTransportException {
        if (this.reqBuffer == null) {
            throw new TTransportException("Transport not yet executed");
        }
        this.reqBuffer.appendBytes(bytes, off, len);
    }
}