gobblin.tunnel.MockServer.java Source code

Introduction

Here is the source code for gobblin.tunnel.MockServer.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 gobblin.tunnel;

import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.Vector;

/**
 * A convenient-to-use TCP server class to implement mock test servers.
 *
 * @author kkandekar@linkedin.com
 */
abstract class MockServer {
    protected static final Logger LOG = LoggerFactory.getLogger(MockServer.class);

    volatile boolean _serverRunning = true;
    ServerSocket _server;
    Set<EasyThread> _threads = Collections.synchronizedSet(new HashSet<EasyThread>());
    List<Socket> _sockets = new Vector<Socket>();
    int _serverSocketPort;
    int numConnects = 0; //only counted for proxy connects

    static void sleepQuietly(long millis) {
        try {
            Thread.sleep(millis);
        } catch (InterruptedException e) {
        }
    }

    public MockServer start() throws IOException {
        _server = new ServerSocket();
        _server.setSoTimeout(5000);
        _server.bind(new InetSocketAddress("localhost", 0));
        _serverSocketPort = _server.getLocalPort();
        _threads.add(new EasyThread() {
            @Override
            void runQuietly() throws Exception {
                runServer();
            }
        }.startThread());
        return this;
    }

    // accept thread
    public void runServer() {
        while (_serverRunning) {
            try {
                final Socket clientSocket = _server.accept();
                numConnects++;
                LOG.info("Accepted connection on " + getServerSocketPort());
                // client handler thread
                _threads.add(new EasyThread() {
                    @Override
                    void runQuietly() throws Exception {
                        try {
                            addSocket(clientSocket);
                            handleClientSocket(clientSocket);
                        } catch (IOException e) {
                            LOG.info("Exception in handleClientSocket, stopping server", e);
                            stopServer();
                        }
                    }
                }.startThread());
            } catch (IOException e) {
                LOG.info("Exception in runServer", e);
            }
        }
        try {
            _server.close();
        } catch (IOException e) {
            LOG.info("Exception while closing server", e);
        }
    }

    int getNumConnects() {
        return numConnects;
    }

    boolean isServerRunning() {
        return _serverRunning;
    }

    int getServerSocketPort() {
        return _serverSocketPort;
    }

    // need to keep track of socket because interrupting thread is not working
    void addSocket(Socket socket) {
        _sockets.add(socket);
    }

    abstract void handleClientSocket(Socket socket) throws IOException;

    public void stopServer() {
        _serverRunning = false;
        IOUtils.closeQuietly(_server);
        for (Socket socket : _sockets) {
            IOUtils.closeQuietly(socket);
        }
        for (EasyThread thread : _threads) {
            thread.interrupt();
        }
    }
}