com.jetdrone.vertx.mods.stomp.StompClient Maven / Gradle / Ivy

Go to download
package com.jetdrone.vertx.mods.stomp;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import org.vertx.java.core.buffer.Buffer;
import org.vertx.java.core.net.NetSocket;
import org.vertx.java.core.AsyncResult;
import org.vertx.java.core.AsyncResultHandler;
import org.vertx.java.core.Handler;
import org.vertx.java.core.Vertx;
import org.vertx.java.core.logging.Logger;
import org.vertx.java.core.net.NetClient;

import java.io.IOException;
import java.util.*;

/**
 * STOMP StompClient Class
 *
 *All STOMP protocol is exposed as methods of this class (`connect()`, `send()`, etc.)
 */
public class StompClient {

    private static class Heartbeat {
        int sx;
        int sy;

        static Heartbeat parse(String header) {
            String[] token = header.split(",");
            Heartbeat beat = new Heartbeat();
            beat.sx = Integer.parseInt(token[0]);
            beat.sy = Integer.parseInt(token[1]);

            return beat;
        }

        @Override
        public String toString() {
            return sx + "," + sy;
        }
    }

    private static final Frame ASYNCFRAME = new Frame("ASYNC");

    private final Heartbeat heartbeat = new Heartbeat();
    private final Queue> replies = new LinkedList<>();
    private final StompSubscriptions subscriptions;

    private long pinger;
    private long ponger;
    private long serverActivity;

    private final Vertx vertx;
    private final Logger logger;
    private final String host;
    private final int port;
    private final String login;
    private final String passcode;

    private NetSocket netSocket;

    private static enum State {
        DISCONNECTED,
        CONNECTING,
        CONNECTED
    }

    private State state = State.DISCONNECTED;

    private static String getSupportedVersions() {
        return Protocol.V1_2.version + "," + Protocol.V1_1.version + "," + Protocol.V1_0.version;
    }

    public StompClient(Vertx vertx, Logger logger, String host, int port, String login, String passcode, StompSubscriptions subscriptions) {
        this.vertx = vertx;
        this.logger = logger;
        this.host = host;
        this.port = port;
        this.login = login;
        this.passcode = passcode;
        this.subscriptions = subscriptions;
    }


    public void connect(final AsyncResultHandler resultHandler) {
        if (state == State.DISCONNECTED) {
            state = State.CONNECTING;
            NetClient client = vertx.createNetClient();
            client.connect(port, host, new AsyncResultHandler() {
                @Override
                public void handle(final AsyncResult asyncResult) {
                    if (asyncResult.failed()) {
                        logger.error("Net client error", asyncResult.cause());
                        if (resultHandler != null) {
                            resultHandler.handle(new AsyncResult() {

                                @Override
                                public Void result() {
                                    return null;
                                }

                                @Override
                                public Throwable cause() {
                                    return asyncResult.cause();
                                }

                                @Override
                                public boolean succeeded() {
                                    return asyncResult.succeeded();
                                }

                                @Override
                                public boolean failed() {
                                    return asyncResult.failed();
                                }
                            });
                        }
                        disconnect();
                    } else {
                        state = State.CONNECTED;
                        netSocket = asyncResult.result();
                        init(netSocket);
                        netSocket.exceptionHandler(new Handler() {
                            public void handle(Throwable e) {
                                logger.error("Socket client error", e);
                                disconnect();
                            }
                        });
                        netSocket.closeHandler(new Handler() {
                            public void handle(Void arg0) {
                                logger.info("Socket closed");
                                disconnect();
                            }
                        });
                        if (resultHandler != null) {
                            resultHandler.handle(new AsyncResult() {

                                @Override
                                public Void result() {
                                    return null;
                                }

                                @Override
                                public Throwable cause() {
                                    return null;
                                }

                                @Override
                                public boolean succeeded() {
                                    return true;
                                }

                                @Override
                                public boolean failed() {
                                    return false;
                                }
                            });
                        }
                    }
                }
            });
        }
    }

    /**
     * Clean up client resources when it is disconnected or the server did not
     * send heart beats in a timely fashion
     */
    private void disconnect() {
        state = State.DISCONNECTED;
        if (pinger != 0) {
            vertx.cancelTimer(pinger);
        }
        if (ponger != 0) {
            vertx.cancelTimer(ponger);
        }
        // make sure the socket is closed
        if (netSocket != null) {
            netSocket.close();
        }
    }

    void send(final Frame frame, final boolean async, final Handler replyHandler) {
        switch (state) {
            case CONNECTED:
                netSocket.write(frame.command);
                netSocket.write("\n");

                for (Map.Entry entry : frame.headers.entrySet()) {
                    String value = entry.getValue();
                    if (value != null) {
                        netSocket.write(entry.getKey());
                        netSocket.write(":");
                        netSocket.write(Frame.escape(entry.getValue()));
                        netSocket.write("\n");
                    }
                }

                if (frame.body != null) {
                    netSocket.write("content-length:");
                    netSocket.write(Integer.toString(frame.body.length()));
                    netSocket.write("\n");
                }

                netSocket.write("\n");
                if (frame.body != null) {
                    netSocket.write(frame.body);
                }

                netSocket.write("\0");
                if (async) {
                    replyHandler.handle(ASYNCFRAME);
                } else {
                    replies.offer(replyHandler);
                }
                break;
            case DISCONNECTED:
                logger.info("Got request when disconnected. Trying to connect.");
                connect(new AsyncResultHandler() {
                    public void handle(AsyncResult connection) {
                        if (connection.succeeded()) {
                            send(frame, async, replyHandler);
                        } else {
                            Frame error = new Frame("ERROR");
                            error.body = "Unable to connect";
                            replyHandler.handle(error);
                        }
                    }
                });
                break;
            case CONNECTING:
                logger.debug("Got send request while connecting. Will try again in a while.");
                vertx.setTimer(100, new Handler() {
                    public void handle(Long event) {
                        send(frame, async, replyHandler);
                    }
                });
        }
    }


    private void init(NetSocket netSocket) {
        this.netSocket = netSocket;
        final StompDecoder stompDecoder = new StompDecoder();
        // send heartbeat every 10s by default (value is in ms)
        heartbeat.sx = 10000;
        // expect to receive server heartbeat at least every 10s by default (value in ms)
        heartbeat.sy = 10000;
        // setup the handlers
        netSocket.dataHandler(new Handler() {
            private ByteBuf read = null;

            @Override
            public void handle(Buffer buffer) {
//                System.out.println("<<<" + buffer.toString("UTF-8").replaceAll("\0", "^@"));
                serverActivity = System.currentTimeMillis();
                // Should only get one callback at a time, no sychronization necessary
                ByteBuf byteBuf = buffer.getByteBuf();

                if (read != null) {
                    // Merge the new buffer with the previous buffer
                    byteBuf = Unpooled.copiedBuffer(read, byteBuf);
                    read = null;
                }

                try {
                    // Attempt to decode a full reply from the channelbuffer
                    Frame receive = stompDecoder.receive(byteBuf);
                    // If successful, grab the matching handler
                    handleReply(receive);
                    // May be more to read
                    if (byteBuf.isReadable()) {
                        // More than one message in the buffer, need to be careful
                        handle(new Buffer(Unpooled.copiedBuffer(byteBuf)));
                    }
                } catch (IOException e) {
                    logger.error("Error receiving data", e);
                    disconnect();
                } catch (IndexOutOfBoundsException th) {
                    // Got to catch decoding fails and try it again
                    byteBuf.resetReaderIndex();
                    read = Unpooled.copiedBuffer(byteBuf);
                }
            }
        });
        // perform the connect command
        logger.debug("Socket Opened...");
        Frame connect = new Frame("CONNECT");
        connect.putHeader("accept-version", getSupportedVersions());
        connect.putHeader("heart-beat", heartbeat.toString());
        connect.putHeader("vhost", host);
        if (login != null) {
            connect.putHeader("login", login);
        }
        if (passcode != null) {
            connect.putHeader("passcode", passcode);
        }

        send(connect, false, new Handler() {
            @Override
            public void handle(Frame frame) {
                logger.debug("connected to server " + frame.headers.get("server"));
                // connected = true
                setupHeartbeat(frame.headers);
            }
        });
    }

    private void setupHeartbeat(Map headers) {
        if (headers.get("version").equals(Protocol.V1_0.version)) {
            return;
        }

        // heart-beat header received from the server looks like:
        //
        //     heart-beat: sx, sy
        Heartbeat heartbeat = Heartbeat.parse(headers.get("heart-beat"));

        if (this.heartbeat.sx != 0 && heartbeat.sy != 0) {
            final int ttl = Math.max(this.heartbeat.sx, heartbeat.sy);
            logger.debug("send PING every " + ttl + "ms");
            this.pinger = vertx.setPeriodic(ttl, new Handler() {
                @Override
                public void handle(Long event) {
                    if (state == State.CONNECTED) {
                        logger.debug("PING");
                        netSocket.write("\n");
                    }
                }
            });
        }

        if (this.heartbeat.sy != 0 && heartbeat.sx != 0) {
            final int ttl = Math.max(this.heartbeat.sy, heartbeat.sx);
            logger.debug("check PONG every " + ttl + "ms");
            ponger = vertx.setPeriodic(ttl, new Handler() {
                @Override
                public void handle(Long event) {
                    long delta = System.currentTimeMillis() - serverActivity;
                    // We wait twice the TTL to be flexible on window's setInterval calls
                    if (delta > ttl * 2) {
                        logger.debug("did not receive server activity for the last " + delta + "ms");
                        disconnect();
                    }
                }
            });
        }
    }

    void handleReply(Frame reply) throws IOException {
        if ("ERROR".equals(reply.command)) {
            logger.error(reply.body);
            disconnect();
            return;
        }

        Handler handler = replies.poll();

        if (handler != null) {
            // handler waits for this response
            handler.handle(reply);
            return;
        }

        // this is a subscribe message
        if ("MESSAGE".equals(reply.command)) {
            handler = subscriptions.getHandler(reply.headers.get("subscription"));
            if (handler != null) {
                // pub sub handler exists
                handler.handle(reply);
                return;
            }

        }

        System.out.println(reply.toJSON());

        throw new IOException("Received a non MESSAGE while in SUBSCRIBE mode");
    }
}