CommandPortService.java

/*
 * Copyright © 2014 Cask Data, Inc.
 *
 * 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 co.cask.cdap.common.service;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.LineReader;
import com.google.common.util.concurrent.AbstractExecutionThreadService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.Map;

/**
 * This class acts as a simple TCP server that accepts textual command and produce textual response.
 * The serving loop is single thread and can only serve one client at a time. {@link CommandHandler}s
 * binded to commands are invoked from the serving thread and is expected to return promptly to not
 * blocking the serving thread.
 *
 * Sample usage:
 * <pre>
 *   CommandPortService service = CommandPortService.builder("myservice")
 *                                                  .addCommandHandler("ruok", "Are you okay?", ruokHandler)
 *                                                  .build();
 *   service.startAndWait();
 * </pre>
 *
 * To stop the service, invoke {@link #stop()} or {@link #stopAndWait()}.
 */
public final class CommandPortService extends AbstractExecutionThreadService {

  private static final Logger LOG = LoggerFactory.getLogger(CommandPortService.class);

  /**
   * Stores binding from command to {@link CommandHandler}.
   */
  private final Map<String, CommandHandler> handlers;

  /**
   * The server socket for accepting incoming requests.
   */
  private ServerSocket serverSocket;

  /**
   * Port that the server socket binded to. It will only be set after the server socket is binded.
   */
  private int port;

  /**
   * Returns a {@link Builder} to build instance of this class.
   * @param serviceName Name of the service name to build
   * @return A {@link Builder}.
   */
  public static Builder builder(String serviceName) {
    return new Builder(serviceName);
  }

  private CommandPortService(int port, Map<String, CommandHandler> handlers) {
    this.port = port;
    this.handlers = handlers;
  }

  @Override
  protected void startUp() throws Exception {
    serverSocket = new ServerSocket(port, 0, InetAddress.getByName("localhost"));
    port = serverSocket.getLocalPort();
  }

  @Override
  protected void run() throws Exception {
    LOG.info("Running commandPortService at localhost:" + port);
    serve();
  }

  @Override
  protected void shutDown() throws Exception {
    // The serverSocket would never be null if this method is called (guaranteed by AbstractExecutionThreadService).
    serverSocket.close();
  }

  @Override
  protected void triggerShutdown() {
    // The serverSocket would never be null if this method is called (guaranteed by AbstractExecutionThreadService).
    try {
      if (serverSocket.isBound()) {
        serverSocket.close();
      }
    } catch (IOException e) {
      throw Throwables.propagate(e);
    }
  }

  /**
   * Returns the port that the server is binded to.
   *
   * @return An int represent the port number.
   */
  public int getPort() {
    Preconditions.checkState(isRunning());
    return port;
  }

  /**
   * Starts accepting incoming request. This method would block until this service is stopped.
   *
   * @throws IOException If any I/O error occurs on the socket connection.
   */
  private void serve() throws IOException {
    while (isRunning()) {
      try {
        Socket socket = serverSocket.accept();
        BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(socket.getOutputStream(), "UTF-8"));
        try {
          // Read the client command and dispatch
          String command = new LineReader(new InputStreamReader(socket.getInputStream(), "UTF-8")).readLine();
          CommandHandler handler = handlers.get(command);

          if (handler != null) {
            try {
              handler.handle(writer);
            } catch (Throwable t) {
              LOG.error(String.format("Exception thrown from CommandHandler for command %s", command), t);
            }
          }
        } finally {
          writer.flush();
          socket.close();
        }
      } catch (Throwable th) {
        // NOTE: catch any exception to keep the main service running
        // Trigger by serverSocket.close() through the call from stop().
        LOG.debug(th.getMessage(), th);
      }
    }
  }

  /**
   * Builder for creating {@link CommandPortService}.
   */
  public static final class Builder {
    private final ImmutableMap.Builder<String, CommandHandler> handlerBuilder;
    private final StringBuilder helpStringBuilder;
    private boolean hasHelp;
    private int port = 0;

    /**
     * Creates a builder for the give name.
     *
     * @param serviceName Name of the {@link CommandPortService} to build.
     */
    private Builder(String serviceName) {
      handlerBuilder = ImmutableMap.builder();
      helpStringBuilder = new StringBuilder(String.format("Help for %s command port service", serviceName));
    }

    /**
     * Adds a {@link CommandHandler} for a given command.
     *
     * @param command Name of the command handled by the handler.
     * @param desc A human readable description about the command.
     * @param handler The {@link CommandHandler} to invoke when the given command is received.
     * @return This {@link Builder}.
     */
    public Builder addCommandHandler(String command, String desc, CommandHandler handler) {
      hasHelp = "help".equals(command);
      handlerBuilder.put(command, handler);
      helpStringBuilder.append("\n").append(String.format("  %s : %s", command, desc));
      return this;
    }

    public Builder setPort(int port) {
      this.port = port;
      return this;
    }

    /**
     * Builds the {@link CommandPortService}.
     *
     * @return A {@link CommandPortService}.
     */
    public CommandPortService build() {
      if (!hasHelp) {
        final String helpString = helpStringBuilder.toString();
        handlerBuilder.put("help", new CommandHandler() {
          @Override
          public void handle(BufferedWriter respondWriter) throws IOException {
            respondWriter.write(helpString);
            respondWriter.newLine();
          }
        });
      }

      return new CommandPortService(port, handlerBuilder.build());
    }
  }

  /**
   * Interface for defining handler to react to a command.
   */
  public interface CommandHandler {

    /**
     * Invoked when the command that tied to this handler has been received
     * by the {@link CommandPortService}.
     *
     * @param respondWriter The {@link Writer} for writing response back to client
     * @throws IOException If I/O errors occurs when writing to the given writer.
     */
    void handle(BufferedWriter respondWriter) throws IOException;
  }
}