com.jgoetsch.eventtrader.source.SocketIOWebSocketMsgSource.java Source code

Java tutorial

Introduction

Here is the source code for com.jgoetsch.eventtrader.source.SocketIOWebSocketMsgSource.java

Source

/*
 * Copyright (c) 2012 Jeremy Goetsch
 *
 * Licensed 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 com.jgoetsch.eventtrader.source;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicInteger;

import javax.annotation.PreDestroy;

import org.jboss.netty.bootstrap.ClientBootstrap;
import org.jboss.netty.channel.Channel;
import org.jboss.netty.channel.ChannelFuture;
import org.jboss.netty.channel.ChannelHandlerContext;
import org.jboss.netty.channel.ChannelPipeline;
import org.jboss.netty.channel.ChannelPipelineFactory;
import org.jboss.netty.channel.ChannelStateEvent;
import org.jboss.netty.channel.Channels;
import org.jboss.netty.channel.ExceptionEvent;
import org.jboss.netty.channel.MessageEvent;
import org.jboss.netty.channel.SimpleChannelUpstreamHandler;
import org.jboss.netty.channel.socket.nio.NioClientSocketChannelFactory;
import org.jboss.netty.handler.codec.http.HttpRequestEncoder;
import org.jboss.netty.handler.codec.http.HttpResponse;
import org.jboss.netty.handler.codec.http.HttpResponseDecoder;
import org.jboss.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import org.jboss.netty.handler.codec.http.websocketx.WebSocketClientHandshaker;
import org.jboss.netty.handler.codec.http.websocketx.WebSocketClientHandshakerFactory;
import org.jboss.netty.handler.codec.http.websocketx.WebSocketVersion;
import org.jboss.netty.util.ThreadNameDeterminer;
import org.jboss.netty.util.ThreadRenamingRunnable;
import org.json.simple.JSONArray;
import org.json.simple.JSONObject;
import org.json.simple.JSONValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.jgoetsch.eventtrader.source.parser.structured.StructuredMsgParser;

/**
    
 * Source for streaming messages from a Socket.IO server
 * using WebSockets as the transport.
 * 
 * As of March, 2012 this is the preferred source for Profiding alerts.
 */

public class SocketIOWebSocketMsgSource extends UrlBasedMsgSource {
    private Logger log = LoggerFactory.getLogger(getClass());
    private StructuredMsgParser msgParser;

    private ClientBootstrap bootstrap;
    private Channel ch;

    protected String getTokenUrl() throws IOException, URISyntaxException {
        URI base = new URI(getUrl());
        BufferedReader tokenReader = new BufferedReader(
                new InputStreamReader(new URI("http", base.getUserInfo(), base.getHost(), base.getPort(),
                        base.getPath(), base.getQuery(), base.getFragment()).toURL().openStream()));
        String token = tokenReader.readLine();
        tokenReader.close();

        String r[] = token.split(":");
        String comp[] = getUrl().split("\\?");
        if (!comp[0].endsWith("/"))
            comp[0] += '/';
        return comp[0] + "websocket/" + r[0] + (comp.length > 0 ? "?" + comp[1] : "");
    }

    @Override
    protected void receiveMsgs() {
        final String baseThreadName = Thread.currentThread().getName();
        ThreadRenamingRunnable.setThreadNameDeterminer(ThreadNameDeterminer.CURRENT);
        ThreadFactory threadFactory = new ThreadFactory() {
            private AtomicInteger n = new AtomicInteger();

            public Thread newThread(Runnable r) {
                return new Thread(r, baseThreadName + "-w-" + n.incrementAndGet());
            }
        };
        bootstrap = new ClientBootstrap(new NioClientSocketChannelFactory(
                Executors.newCachedThreadPool(threadFactory), Executors.newCachedThreadPool(threadFactory)));

        try {
            URI url = new URI(getTokenUrl());

            final WebSocketClientHandshaker handshaker = new WebSocketClientHandshakerFactory().newHandshaker(url,
                    WebSocketVersion.V13, null, false, null);

            bootstrap.setPipelineFactory(new ChannelPipelineFactory() {
                public ChannelPipeline getPipeline() throws Exception {
                    ChannelPipeline pipeline = Channels.pipeline();
                    pipeline.addLast("decoder", new HttpResponseDecoder());
                    pipeline.addLast("encoder", new HttpRequestEncoder());
                    pipeline.addLast("ws-handler", new WebSocketClientHandler(handshaker));
                    pipeline.addLast("sio-handler", new SocketIOClientHandler());
                    return pipeline;
                }
            });

            ChannelFuture future = bootstrap
                    .connect(new InetSocketAddress(url.getHost(), url.getPort() == -1 ? 80 : url.getPort()));
            future.syncUninterruptibly();
            ch = future.getChannel();
            handshaker.handshake(ch).syncUninterruptibly();
            ch.getCloseFuture().awaitUninterruptibly();
        } catch (URISyntaxException use) {
            log.error("Invalid URL: {}", getUrl(), use);
        } catch (Exception e) {
            log.error("Error getting token", e);
        } finally {
            if (ch != null)
                ch.close();
            bootstrap.releaseExternalResources();
        }
    }

    @PreDestroy
    public void shutdown() {
        if (ch != null)
            ch.close();
        joinMainThread(2000);
    }

    private class WebSocketClientHandler extends SimpleChannelUpstreamHandler {

        private final WebSocketClientHandshaker handshaker;

        public WebSocketClientHandler(WebSocketClientHandshaker handshaker) {
            this.handshaker = handshaker;
        }

        @Override
        public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Exception {
            if (!handshaker.isHandshakeComplete()) {
                handshaker.finishHandshake(ctx.getChannel(), (HttpResponse) e.getMessage());
                log.info("Connected to websocket host {}", ctx.getChannel().getRemoteAddress());
                return;
            }
            super.messageReceived(ctx, e);
        }

        @Override
        public void channelClosed(ChannelHandlerContext ctx, ChannelStateEvent e) throws Exception {
            log.info("Websocket connection to {} closed", ctx.getChannel().getRemoteAddress());
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, ExceptionEvent e) throws Exception {
            log.error("Websocket exception", e.getCause());
            e.getChannel().close();
        }
    }

    private class SocketIOClientHandler extends SimpleChannelUpstreamHandler {
        @Override
        public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Exception {
            if (e.getMessage() instanceof TextWebSocketFrame) {
                TextWebSocketFrame message = (TextWebSocketFrame) e.getMessage();
                if (message.getText().equals("2::")) {
                    log.trace("Heartbeat received");
                    ch.write(new TextWebSocketFrame("2::"));
                } else {
                    log.debug("Received: {}", message.getText());
                    try {
                        int idx = message.getText().indexOf('{');
                        if (idx >= 0) {
                            JSONObject json = (JSONObject) JSONValue.parse(message.getText().substring(idx));

                            Object data = json.get("args");
                            if (data != null && data.getClass().isAssignableFrom(JSONArray.class)) {
                                data = ((JSONArray) json.get("args")).get(0);
                                if (data.getClass().isAssignableFrom(JSONObject.class)) {
                                    if (!msgParser.parseData((String) json.get("name"), (JSONObject) data,
                                            SocketIOWebSocketMsgSource.this)) {
                                        ctx.getChannel().close();
                                    }
                                }
                            }
                        }
                    } catch (Exception ex) {
                        log.error("Exception occured parsing message: " + message.getText(), ex);
                    }
                }
            }
            super.messageReceived(ctx, e);
        }
    }

    public void setMsgParser(StructuredMsgParser msgParser) {
        this.msgParser = msgParser;
    }

    public StructuredMsgParser getMsgParser() {
        return msgParser;
    }
}