com.splout.db.dnode.TCPStreamer.java Source code

Java tutorial

Introduction

Here is the source code for com.splout.db.dnode.TCPStreamer.java

Source

package com.splout.db.dnode;

/*
 * #%L
 * Splout SQL Server
 * %%
 * Copyright (C) 2012 - 2014 Datasalt Systems S.L.
 * %%
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 * #L%
 */

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.util.Arrays;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.splout.db.common.SploutConfiguration;
import com.splout.db.engine.EngineManager;
import com.splout.db.engine.EngineManager.EngineException;
import com.splout.db.engine.ResultSerializer;
import com.splout.db.engine.ResultSerializer.SerializationException;
import com.splout.db.engine.StreamingIterator;
import com.splout.db.thrift.DNodeException;

/**
 * An interface for the DNode to send streaming data through TCP.
 */
public class TCPStreamer {

    private final static Log log = LogFactory.getLog(TCPStreamer.class);

    private DNodeHandler dNode;
    private int tcpPort;

    private TCPServer server;

    public void start(SploutConfiguration config, DNodeHandler dNode) throws InterruptedException, IOException {
        this.dNode = dNode;
        this.tcpPort = config.getInt(DNodeProperties.STREAMING_PORT);

        server = new TCPServer();
        server.bind();

        Thread t = new Thread() {
            @Override
            public void run() {
                server.serve();
            }
        };
        t.start();

        while (!server.isServing()) {
            Thread.sleep(100);
        }
    }

    public void stop() {
        server.stop();
    }

    public int getTcpPort() {
        return tcpPort;
    }

    class TCPServer {

        protected ServerSocket serverSocket = null;
        protected boolean isStopped = false;

        private synchronized boolean isStopped() {
            return this.isStopped;
        }

        public synchronized void stop() {
            this.isStopped = true;
            try {
                this.serverSocket.close();
            } catch (IOException e) {
                throw new RuntimeException("Error closing server", e);
            }
        }

        public boolean isServing() {
            return !isStopped() && this.serverSocket != null && this.serverSocket.isBound();
        }

        public void bind() throws IOException {
            try {
                this.serverSocket = new ServerSocket(tcpPort);
            } catch (IOException e) {
                throw e;
            }
        }

        public void serve() {
            while (!isStopped()) {
                try {
                    new Thread(new WorkerRunnable(this.serverSocket.accept(), dNode)).start();
                } catch (IOException e) {
                    if (isStopped()) {
                        System.out.println("Server Stopped.");
                        return;
                    }
                    throw new RuntimeException("Error accepting client connection", e);
                }
            }
        }
    }

    static class WorkerRunnable implements Runnable {

        protected final Socket clientSocket;
        protected final DNodeHandler handler;

        public WorkerRunnable(Socket clientSocket, DNodeHandler handler) {
            this.clientSocket = clientSocket;
            this.handler = handler;
        }

        public void run() {
            try {
                InputStream input = clientSocket.getInputStream();
                OutputStream output = clientSocket.getOutputStream();

                final DataInputStream dis = new DataInputStream(new BufferedInputStream(input));
                // Get tablespace
                String tablespace = dis.readUTF();
                // Get version number
                long version = dis.readLong();
                // Get partition number
                int partition = dis.readInt();
                // Get query
                final String query = dis.readUTF();

                log.info("Got streaming request: " + tablespace + ", " + version + ", " + partition + ", " + query);
                EngineManager manager = handler.getManager(tablespace, version, partition);

                final DataOutputStream dos = new DataOutputStream(new BufferedOutputStream(output));

                manager.streamQuery(new StreamingIterator() {

                    @Override
                    public String getQuery() {
                        return query;
                    }

                    @Override
                    public void columns(String[] columns) {
                    }

                    @Override
                    public void collect(Object[] result) throws StreamingTerminationException {
                        byte[] bytes;
                        try {
                            bytes = ResultSerializer.serializeToByteArray(result);
                            dos.writeInt(bytes.length);
                            dos.write(bytes);
                        } catch (SerializationException e) {
                            throw new StreamingTerminationException("", e);
                        } catch (IOException e) {
                            log.error(e);
                            e.printStackTrace();
                            throw new StreamingTerminationException("", e);
                        }
                    }

                    @Override
                    public void endStreaming() {
                        try {
                            dos.writeInt(-1);
                            dos.flush();
                            log.info("Finished stream");
                        } catch (IOException e) {
                            log.error(e);
                            e.printStackTrace();
                        }
                    }
                });

                input.close();
                output.close();
                clientSocket.close();
            } catch (IOException e) {
                e.printStackTrace();
            } catch (EngineException e) {
                e.printStackTrace();
            } catch (DNodeException e1) {
                e1.printStackTrace();
            }
        }
    }

    /**
     * This main method can be used for testing the TCP interface directly to a
     * local DNode. Will ask for protocol input from Stdin and print output to
     * Stdout
     */
    public static void main(String[] args) throws UnknownHostException, IOException, SerializationException {
        SploutConfiguration config = SploutConfiguration.get();
        Socket clientSocket = new Socket("localhost", config.getInt(DNodeProperties.STREAMING_PORT));

        DataInputStream inFromServer = new DataInputStream(new BufferedInputStream(clientSocket.getInputStream()));
        DataOutputStream outToServer = new DataOutputStream(
                new BufferedOutputStream(clientSocket.getOutputStream()));

        BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
        System.out.println("Enter tablespace: ");
        String tablespace = reader.readLine();

        System.out.println("Enter version number: ");
        long versionNumber = Long.parseLong(reader.readLine());

        System.out.println("Enter partition: ");
        int partition = Integer.parseInt(reader.readLine());

        System.out.println("Enter query: ");
        String query = reader.readLine();

        outToServer.writeUTF(tablespace);
        outToServer.writeLong(versionNumber);
        outToServer.writeInt(partition);
        outToServer.writeUTF(query);

        outToServer.flush();

        byte[] buffer = new byte[0];
        boolean read;
        do {
            read = false;
            int nBytes = inFromServer.readInt();
            if (nBytes > 0) {
                buffer = new byte[nBytes];
                int inRead = inFromServer.read(buffer);
                if (inRead > 0) {
                    Object[] res = ResultSerializer.deserialize(ByteBuffer.wrap(buffer), Object[].class);
                    read = true;
                    System.out.println(Arrays.toString(res));
                }
            }
        } while (read);

        clientSocket.close();
    }
}