com.linkedin.proxy.main.ProxyServer.java Source code

Java tutorial

Introduction

Here is the source code for com.linkedin.proxy.main.ProxyServer.java

Source

/**
 * Copyright 2014 LinkedIn Corp. 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.
 */

package com.linkedin.proxy.main;

import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.Channel;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.HashSet;
import java.util.Properties;
import java.util.Set;

import org.apache.log4j.Logger;
import org.rocksdb.RocksDB;

import com.linkedin.proxy.netty.MysqlInitializer;
import com.linkedin.proxy.netty.RocksdbInitializer;
import com.linkedin.proxy.pool.BlockingMysqlConnectionPool;
import com.linkedin.proxy.pool.BlockingRocksdbConnectionPool;
import com.linkedin.proxy.pool.ConnectionPool;

public class ProxyServer {
    public enum ProxyMode {
        MYSQL, ROCKSDB
    }

    private static final String FLAG_PROXY_MODE = "mode";
    private static final String FLAG_PROXY_RUNTIME = "runtime";
    private static final String FLAG_PROXY_THR = "thrPool";
    private static final String FLAG_PROXY_PORT = "port";

    private static final Logger m_log = Logger.getLogger(ProxyServer.class);

    private static String DB_PORT_MAP_FILE = "";
    private static String DB_SET_FILE = "";
    private static String PROP_FILE = "";

    private static EventLoopGroup bossGroup;
    private static EventLoopGroup workerGroup;
    private static ConnectionPool connPool;
    private static Channel ch;
    private static boolean isClosed = false;

    public static void main(String[] args) throws Exception {
        //process command line
        if (processCommandLine(args) == false) {
            m_log.fatal("Command line processing error. Closing...");
            return;
        }

        //process properties file
        Properties prop = getProperties(PROP_FILE);
        if (prop == null) {
            m_log.fatal("Error in processing properties file. Closing...");
            return;
        }

        //set mode
        ProxyMode runMode;
        String temp = prop.getProperty(FLAG_PROXY_MODE);
        if (temp == null) {
            m_log.fatal("Mode is missing. Closing...");
            return;
        } else if (temp.equals("mysql")) {
            runMode = ProxyMode.MYSQL;
        } else if (temp.equals("rocksdb")) {
            runMode = ProxyMode.ROCKSDB;
        } else {
            m_log.fatal("Unknown mode " + temp + ". Closing...");
            return;
        }

        //Process mode specific files
        Set<String> dbSet = new HashSet<String>();
        if (runMode == ProxyMode.ROCKSDB) {
            //for rocksdb, I need db names
            if (DB_SET_FILE.equals("")) {
                m_log.fatal("DB set file is missing. Closing...");
                return;
            } else if (processDbSet(dbSet, DB_SET_FILE) == false) {
                m_log.fatal("Error in processing Dbset file. Closing...");
                return;
            } else {
                m_log.info("DB set file is processed");
            }
        } else {
            m_log.fatal("Unknown mode " + runMode + ". Closing...");
            return;
        }

        //perform mode based initializations if any
        if (runMode == ProxyMode.ROCKSDB) {
            RocksDB.loadLibrary();
        }

        //get run time
        int runTime;
        temp = prop.getProperty(FLAG_PROXY_RUNTIME);
        if (temp == null) {
            runTime = 0;
        } else {
            runTime = Integer.parseInt(temp);
        }
        m_log.info("Runtime is " + runTime);

        //get thread pool size
        int thrSize;
        temp = prop.getProperty(FLAG_PROXY_THR);
        if (temp == null) {
            m_log.warn("Thread pool size parameter is missing. It is set to 10 by default");
            thrSize = 10;
        } else {
            thrSize = Integer.parseInt(temp);
        }

        //get listening port
        int port;
        temp = prop.getProperty(FLAG_PROXY_PORT);
        if (temp == null) {
            m_log.fatal("Listening port is not specified. Closing...");
            return;
        } else {
            port = Integer.parseInt(temp);
        }

        //init thread pools
        bossGroup = new NioEventLoopGroup(1);
        workerGroup = new NioEventLoopGroup(thrSize);

        //create connection pools
        if (runMode == ProxyMode.ROCKSDB) {
            connPool = new BlockingRocksdbConnectionPool(dbSet);
        } else if (runMode == ProxyMode.MYSQL) {
            connPool = new BlockingMysqlConnectionPool();
        } else {
            m_log.fatal("Unkown setup. Closing...");
            return;
        }

        //init connection pool
        if (connPool.init(prop) == false) {
            m_log.fatal("Cannot init conn pool. Closing...");
            return;
        }

        //if run time is specified, then start closing thread
        Thread closingThread = null;
        if (runTime > 0) {
            closingThread = new ClosingThread(runTime);
            closingThread.start();

            System.out.println("Closing in " + runTime + " seconds.");
        } else {
            System.out.println("Type \"close\" to close proxy.");
        }

        try {
            ServerBootstrap b = new ServerBootstrap();

            if (runMode == ProxyMode.MYSQL) {
                b.group(bossGroup, workerGroup).channel(NioServerSocketChannel.class)
                        .handler(new LoggingHandler(LogLevel.INFO))
                        .childHandler(new MysqlInitializer(prop, connPool));
            } else if (runMode == ProxyMode.ROCKSDB) {
                b.group(bossGroup, workerGroup).channel(NioServerSocketChannel.class)
                        .handler(new LoggingHandler(LogLevel.INFO))
                        .childHandler(new RocksdbInitializer(prop, connPool));
            }

            ch = b.bind(port).sync().channel();

            if (runTime > 0) {
                ch.closeFuture().sync();
            } else {
                BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
                while (true) {
                    String line = in.readLine();
                    m_log.debug("Got line: " + line);

                    if (line == null || "close".equals(line.toLowerCase())) {
                        break;
                    }
                }
            }
        } catch (Exception e) {
            m_log.error("Error..", e);
        } finally {
            close();
        }
    }

    public static void close() throws Exception {
        if (!isClosed) {
            if (ch != null)
                ch.close();
            bossGroup.shutdownGracefully();
            workerGroup.shutdownGracefully();
            connPool.closeAll();

            isClosed = true;
        }
    }

    private static String getFlags() {
        StringBuilder sb = new StringBuilder();

        sb.append("Allowed flags:\n");
        sb.append("-prop=<FILE_PATH>: Properties file for the execution\n");
        sb.append("-dbSet=<FILE_PATH>: File containing name of the databases\n");
        sb.append("-dbPortMap=<FILE_PATH>: Mapping from database name to host port.");

        return sb.toString();
    }

    /**
     * -propFile=FILE_PATH
     * -dbPortMap=FILE_PATH
     * -dbSet=DB_SET
     * @param args
     * @return true if all information is provided. false otherwise.
     */
    private static boolean processCommandLine(String[] args) {
        try {
            for (int a = 0; a < args.length; a++) {
                if (args[a].startsWith("-dbPortMap")) {
                    DB_PORT_MAP_FILE = args[a].split("=")[1];
                    m_log.info("DbPortMapFile=" + DB_PORT_MAP_FILE);
                } else if (args[a].startsWith("-dbSet")) {
                    DB_SET_FILE = args[a].split("=")[1];
                    m_log.info("DbSetFile=" + DB_SET_FILE);
                } else if (args[a].startsWith("-propFile")) {
                    PROP_FILE = args[a].split("=")[1];
                    m_log.info("RocksDbFolder=" + PROP_FILE);
                } else {
                    m_log.error("Unknown flag: " + args[a]);
                    return false;
                }
            }

            if (PROP_FILE.equals("")) {
                m_log.error("Properties file is not given.");
                return false;
            }

            return true;
        } catch (Exception e) {
            m_log.error("Error while processing command line", e);
            m_log.error(getFlags());

            return false;
        }
    }

    private static boolean processDbSet(Set<String> dbSet, String filename) {
        try {
            FileInputStream fs = new FileInputStream(filename);
            return processDbSet(dbSet, fs);
        } catch (Exception e) {
            m_log.fatal("Error in processing dbSet file", e);
            return false;
        }
    }

    private static boolean processDbSet(Set<String> dbSet, InputStream inStr) {
        BufferedReader reader = new BufferedReader(new InputStreamReader(inStr));

        try {
            String line = reader.readLine();
            while (line != null) {
                dbSet.add(line);

                line = reader.readLine();
            }

            reader.close();
            return true;
        } catch (Exception e) {
            m_log.fatal("Cannot process input stream", e);
            return false;
        }
    }

    /**
     * Tries opening the given properties file, and reads it.
     * @param filename File path to the properties
     * @return Properties object for the given file. Null if any error occurs.
     */
    private static Properties getProperties(String filename) {
        try {
            Properties values = new Properties();
            File f = new File(filename);
            values.load(new FileReader(f));
            return values;
        } catch (Exception e) {
            m_log.fatal("Error in processing properties file", e);
            return null;
        }
    }
}