/*
 * 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 org.apache.thrift.server;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.spi.SelectorProvider;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.thrift.TAsyncProcessor;
import org.apache.thrift.TByteArrayOutputStream;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TIOStreamTransport;
import org.apache.thrift.transport.TMemoryInputTransport;
import org.apache.thrift.transport.TNonblockingServerTransport;
import org.apache.thrift.transport.TNonblockingTransport;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.apache.thrift.transport.layered.TFramedTransport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Provides common methods and classes used by nonblocking TServer implementations. */
public abstract class AbstractNonblockingServer extends TServer {
  protected final Logger LOGGER = LoggerFactory.getLogger(getClass().getName());

  public abstract static class AbstractNonblockingServerArgs<
          T extends AbstractNonblockingServerArgs<T>>
      extends AbstractServerArgs<T> {
    public long maxReadBufferBytes = 256 * 1024 * 1024;

    public AbstractNonblockingServerArgs(TNonblockingServerTransport transport) {
      super(transport);
      transportFactory(new TFramedTransport.Factory());
    }
  }

  /**
   * The maximum amount of memory we will allocate to client IO buffers at a time. Without this
   * limit, the server will gladly allocate client buffers right into an out of memory exception,
   * rather than waiting.
   */
  final long MAX_READ_BUFFER_BYTES;

  /** How many bytes are currently allocated to read buffers. */
  final AtomicLong readBufferBytesAllocated = new AtomicLong(0);

  public AbstractNonblockingServer(AbstractNonblockingServerArgs<?> args) {
    super(args);
    MAX_READ_BUFFER_BYTES = args.maxReadBufferBytes;
  }

  /** Begin accepting connections and processing invocations. */
  public void serve() {
    // start any IO threads
    if (!startThreads()) {
      return;
    }

    // start listening, or exit
    if (!startListening()) {
      return;
    }

    setServing(true);

    // this will block while we serve
    waitForShutdown();

    setServing(false);

    // do a little cleanup
    stopListening();
  }

  /**
   * Starts any threads required for serving.
   *
   * @return true if everything went ok, false if threads could not be started.
   */
  protected abstract boolean startThreads();

  /** A method that will block until when threads handling the serving have been shut down. */
  protected abstract void waitForShutdown();

  /**
   * Have the server transport start accepting connections.
   *
   * @return true if we started listening successfully, false if something went wrong.
   */
  protected boolean startListening() {
    try {
      serverTransport_.listen();
      return true;
    } catch (TTransportException ttx) {
      LOGGER.error("Failed to start listening on server socket!", ttx);
      return false;
    }
  }

  /** Stop listening for connections. */
  protected void stopListening() {
    serverTransport_.close();
  }

  /**
   * Perform an invocation. This method could behave several different ways - invoke immediately
   * inline, queue for separate execution, etc.
   *
   * @return true if invocation was successfully requested, which is not a guarantee that invocation
   *     has completed. False if the request failed.
   */
  protected abstract boolean requestInvoke(FrameBuffer frameBuffer);

  /**
   * An abstract thread that handles selecting on a set of transports and {@link FrameBuffer
   * FrameBuffers} associated with selected keys corresponding to requests.
   */
  protected abstract class AbstractSelectThread extends Thread {
    protected Selector selector;

    // List of FrameBuffers that want to change their selection interests.
    protected final Set<FrameBuffer> selectInterestChanges = new HashSet<FrameBuffer>();

    public AbstractSelectThread() throws IOException {
      this.selector = SelectorProvider.provider().openSelector();
    }

    /** If the selector is blocked, wake it up. */
    public void wakeupSelector() {
      selector.wakeup();
    }

    /**
     * Add FrameBuffer to the list of select interest changes and wake up the selector if it's
     * blocked. When the select() call exits, it'll give the FrameBuffer a chance to change its
     * interests.
     */
    public void requestSelectInterestChange(FrameBuffer frameBuffer) {
      synchronized (selectInterestChanges) {
        selectInterestChanges.add(frameBuffer);
      }
      // wakeup the selector, if it's currently blocked.
      selector.wakeup();
    }

    /**
     * Check to see if there are any FrameBuffers that have switched their interest type from read
     * to write or vice versa.
     */
    protected void processInterestChanges() {
      synchronized (selectInterestChanges) {
        for (FrameBuffer fb : selectInterestChanges) {
          fb.changeSelectInterests();
        }
        selectInterestChanges.clear();
      }
    }

    /**
     * Do the work required to read from a readable client. If the frame is fully read, then invoke
     * the method call.
     */
    protected void handleRead(SelectionKey key) {
      FrameBuffer buffer = (FrameBuffer) key.attachment();
      if (!buffer.read()) {
        cleanupSelectionKey(key);
        return;
      }

      // if the buffer's frame read is complete, invoke the method.
      if (buffer.isFrameFullyRead() && !requestInvoke(buffer)) {
        cleanupSelectionKey(key);
      }
    }

    /** Let a writable client get written, if there's data to be written. */
    protected void handleWrite(SelectionKey key) {
      FrameBuffer buffer = (FrameBuffer) key.attachment();
      if (!buffer.write()) {
        cleanupSelectionKey(key);
      }
    }

    /** Do connection-close cleanup on a given SelectionKey. */
    protected void cleanupSelectionKey(SelectionKey key) {
      // remove the records from the two maps
      FrameBuffer buffer = (FrameBuffer) key.attachment();
      if (buffer != null) {
        // close the buffer
        buffer.close();
      }
      // cancel the selection key
      key.cancel();
    }
  } // SelectThread

  /** Possible states for the FrameBuffer state machine. */
  private enum FrameBufferState {
    // in the midst of reading the frame size off the wire
    READING_FRAME_SIZE,
    // reading the actual frame data now, but not all the way done yet
    READING_FRAME,
    // completely read the frame, so an invocation can now happen
    READ_FRAME_COMPLETE,
    // waiting to get switched to listening for write events
    AWAITING_REGISTER_WRITE,
    // started writing response data, not fully complete yet
    WRITING,
    // another thread wants this framebuffer to go back to reading
    AWAITING_REGISTER_READ,
    // we want our transport and selection key invalidated in the selector
    // thread
    AWAITING_CLOSE
  }

  /**
   * Class that implements a sort of state machine around the interaction with a client and an
   * invoker. It manages reading the frame size and frame data, getting it handed off as wrapped
   * transports, and then the writing of response data back to the client. In the process it manages
   * flipping the read and write bits on the selection key for its client.
   */
  public class FrameBuffer {
    private final Logger LOGGER = LoggerFactory.getLogger(getClass().getName());

    // the actual transport hooked up to the client.
    protected final TNonblockingTransport trans_;

    // the SelectionKey that corresponds to our transport
    protected final SelectionKey selectionKey_;

    // the SelectThread that owns the registration of our transport
    protected final AbstractSelectThread selectThread_;

    // where in the process of reading/writing are we?
    protected FrameBufferState state_ = FrameBufferState.READING_FRAME_SIZE;

    // the ByteBuffer we'll be using to write and read, depending on the state
    protected ByteBuffer buffer_;

    protected final TByteArrayOutputStream response_;

    // the frame that the TTransport should wrap.
    protected final TMemoryInputTransport frameTrans_;

    // the transport that should be used to connect to clients
    protected final TTransport inTrans_;

    protected final TTransport outTrans_;

    // the input protocol to use on frames
    protected final TProtocol inProt_;

    // the output protocol to use on frames
    protected final TProtocol outProt_;

    // context associated with this connection
    protected final ServerContext context_;

    public FrameBuffer(
        final TNonblockingTransport trans,
        final SelectionKey selectionKey,
        final AbstractSelectThread selectThread)
        throws TTransportException {
      trans_ = trans;
      selectionKey_ = selectionKey;
      selectThread_ = selectThread;
      buffer_ = ByteBuffer.allocate(4);

      frameTrans_ = new TMemoryInputTransport();
      response_ = new TByteArrayOutputStream();
      inTrans_ = inputTransportFactory_.getTransport(frameTrans_);
      outTrans_ = outputTransportFactory_.getTransport(new TIOStreamTransport(response_));
      inProt_ = inputProtocolFactory_.getProtocol(inTrans_);
      outProt_ = outputProtocolFactory_.getProtocol(outTrans_);

      if (eventHandler_ != null) {
        context_ = eventHandler_.createContext(inProt_, outProt_);
      } else {
        context_ = null;
      }
    }

    /**
     * Give this FrameBuffer a chance to read. The selector loop should have received a read event
     * for this FrameBuffer.
     *
     * @return true if the connection should live on, false if it should be closed
     */
    public boolean read() {
      if (state_ == FrameBufferState.READING_FRAME_SIZE) {
        // try to read the frame size completely
        if (!internalRead()) {
          return false;
        }

        // if the frame size has been read completely, then prepare to read the
        // actual frame.
        if (buffer_.remaining() == 0) {
          // pull out the frame size as an integer.
          int frameSize = buffer_.getInt(0);
          if (frameSize <= 0) {
            LOGGER.error(
                "Read an invalid frame size of "
                    + frameSize
                    + ". Are you using TFramedTransport on the client side?");
            return false;
          }

          // if this frame will always be too large for this server, log the
          // error and close the connection.
          if (frameSize > trans_.getMaxFrameSize()) {
            LOGGER.error(
                "Read a frame size of "
                    + frameSize
                    + ", which is bigger than the maximum allowable frame size "
                    + trans_.getMaxFrameSize()
                    + " for ALL connections.");
            return false;
          }

          // if this frame will push us over the memory limit, then return.
          // with luck, more memory will free up the next time around.
          if (readBufferBytesAllocated.get() + frameSize > MAX_READ_BUFFER_BYTES) {
            return true;
          }

          // increment the amount of memory allocated to read buffers
          readBufferBytesAllocated.addAndGet(frameSize + 4);

          // reallocate the readbuffer as a frame-sized buffer
          buffer_ = ByteBuffer.allocate(frameSize + 4);
          buffer_.putInt(frameSize);

          state_ = FrameBufferState.READING_FRAME;
        } else {
          // this skips the check of READING_FRAME state below, since we can't
          // possibly go on to that state if there's data left to be read at
          // this one.
          return true;
        }
      }

      // it is possible to fall through from the READING_FRAME_SIZE section
      // to READING_FRAME if there's already some frame data available once
      // READING_FRAME_SIZE is complete.

      if (state_ == FrameBufferState.READING_FRAME) {
        if (!internalRead()) {
          return false;
        }

        // since we're already in the select loop here for sure, we can just
        // modify our selection key directly.
        if (buffer_.remaining() == 0) {
          // get rid of the read select interests
          selectionKey_.interestOps(0);
          state_ = FrameBufferState.READ_FRAME_COMPLETE;
        }

        return true;
      }

      // if we fall through to this point, then the state must be invalid.
      LOGGER.error("Read was called but state is invalid (" + state_ + ")");
      return false;
    }

    /** Give this FrameBuffer a chance to write its output to the final client. */
    public boolean write() {
      if (state_ == FrameBufferState.WRITING) {
        try {
          if (trans_.write(buffer_) < 0) {
            return false;
          }
        } catch (TTransportException e) {
          LOGGER.warn("Got an Exception during write", e);
          return false;
        }

        // we're done writing. now we need to switch back to reading.
        if (buffer_.remaining() == 0) {
          prepareRead();
        }
        return true;
      }

      LOGGER.error("Write was called, but state is invalid (" + state_ + ")");
      return false;
    }

    /** Give this FrameBuffer a chance to set its interest to write, once data has come in. */
    public void changeSelectInterests() {
      switch (state_) {
        case AWAITING_REGISTER_WRITE:
          // set the OP_WRITE interest
          selectionKey_.interestOps(SelectionKey.OP_WRITE);
          state_ = FrameBufferState.WRITING;
          break;
        case AWAITING_REGISTER_READ:
          prepareRead();
          break;
        case AWAITING_CLOSE:
          close();
          selectionKey_.cancel();
          break;
        default:
          LOGGER.error("changeSelectInterest was called, but state is invalid ({})", state_);
      }
    }

    /** Shut the connection down. */
    public void close() {
      // if we're being closed due to an error, we might have allocated a
      // buffer that we need to subtract for our memory accounting.
      if (state_ == FrameBufferState.READING_FRAME
          || state_ == FrameBufferState.READ_FRAME_COMPLETE
          || state_ == FrameBufferState.AWAITING_CLOSE) {
        readBufferBytesAllocated.addAndGet(-buffer_.array().length);
      }
      trans_.close();
      if (eventHandler_ != null) {
        eventHandler_.deleteContext(context_, inProt_, outProt_);
      }
    }

    /** Check if this FrameBuffer has a full frame read. */
    public boolean isFrameFullyRead() {
      return state_ == FrameBufferState.READ_FRAME_COMPLETE;
    }

    /**
     * After the processor has processed the invocation, whatever thread is managing invocations
     * should call this method on this FrameBuffer so we know it's time to start trying to write
     * again. Also, if it turns out that there actually isn't any data in the response buffer, we'll
     * skip trying to write and instead go back to reading.
     */
    public void responseReady() {
      // the read buffer is definitely no longer in use, so we will decrement
      // our read buffer count. we do this here as well as in close because
      // we'd like to free this read memory up as quickly as possible for other
      // clients.
      readBufferBytesAllocated.addAndGet(-buffer_.array().length);

      if (response_.len() == 0) {
        // go straight to reading again. this was probably an oneway method
        state_ = FrameBufferState.AWAITING_REGISTER_READ;
        buffer_ = null;
      } else {
        buffer_ = ByteBuffer.wrap(response_.get(), 0, response_.len());

        // set state that we're waiting to be switched to write. we do this
        // asynchronously through requestSelectInterestChange() because there is
        // a possibility that we're not in the main thread, and thus currently
        // blocked in select(). (this functionality is in place for the sake of
        // the HsHa server.)
        state_ = FrameBufferState.AWAITING_REGISTER_WRITE;
      }
      requestSelectInterestChange();
    }

    /** Actually invoke the method signified by this FrameBuffer. */
    public void invoke() {
      frameTrans_.reset(buffer_.array());
      response_.reset();

      try {
        if (eventHandler_ != null) {
          eventHandler_.processContext(context_, inTrans_, outTrans_);
        }
        processorFactory_.getProcessor(inTrans_).process(inProt_, outProt_);
        responseReady();
        return;
      } catch (TException te) {
        LOGGER.warn("Exception while invoking!", te);
      } catch (Throwable t) {
        LOGGER.error("Unexpected throwable while invoking!", t);
      }
      // This will only be reached when there is a throwable.
      state_ = FrameBufferState.AWAITING_CLOSE;
      requestSelectInterestChange();
    }

    /**
     * Perform a read into buffer.
     *
     * @return true if the read succeeded, false if there was an error or the connection closed.
     */
    private boolean internalRead() {
      try {
        return trans_.read(buffer_) >= 0;
      } catch (TTransportException e) {
        LOGGER.warn("Got an Exception in internalRead", e);
        return false;
      }
    }

    /** We're done writing, so reset our interest ops and change state accordingly. */
    private void prepareRead() {
      // we can set our interest directly without using the queue because
      // we're in the select thread.
      selectionKey_.interestOps(SelectionKey.OP_READ);
      // get ready for another go-around
      buffer_ = ByteBuffer.allocate(4);
      state_ = FrameBufferState.READING_FRAME_SIZE;
    }

    /**
     * When this FrameBuffer needs to change its select interests and execution might not be in its
     * select thread, then this method will make sure the interest change gets done when the select
     * thread wakes back up. When the current thread is this FrameBuffer's select thread, then it
     * just does the interest change immediately.
     */
    protected void requestSelectInterestChange() {
      if (Thread.currentThread() == this.selectThread_) {
        changeSelectInterests();
      } else {
        this.selectThread_.requestSelectInterestChange(this);
      }
    }
  } // FrameBuffer

  public class AsyncFrameBuffer extends FrameBuffer {
    public AsyncFrameBuffer(
        TNonblockingTransport trans, SelectionKey selectionKey, AbstractSelectThread selectThread)
        throws TTransportException {
      super(trans, selectionKey, selectThread);
    }

    public TProtocol getInputProtocol() {
      return inProt_;
    }

    public TProtocol getOutputProtocol() {
      return outProt_;
    }

    public void invoke() {
      frameTrans_.reset(buffer_.array());
      response_.reset();

      try {
        if (eventHandler_ != null) {
          eventHandler_.processContext(context_, inTrans_, outTrans_);
        }
        ((TAsyncProcessor) processorFactory_.getProcessor(inTrans_)).process(this);
        return;
      } catch (TException te) {
        LOGGER.warn("Exception while invoking!", te);
      } catch (Throwable t) {
        LOGGER.error("Unexpected throwable while invoking!", t);
      }
      // This will only be reached when there is a throwable.
      state_ = FrameBufferState.AWAITING_CLOSE;
      requestSelectInterestChange();
    }
  }
}