/*

  * #%L

  * AbstractSocketConnection.java - mongodb-async-driver - Allanbank Consulting, Inc.

  * %%

  * Copyright (C) 2011 - 2014 Allanbank Consulting, 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.

  * #L%

  */

 package com.allanbank.mongodb.client.connection.socket;

 import java.beans.PropertyChangeListener;

 import java.beans.PropertyChangeSupport;

 import java.io.BufferedOutputStream;

 import java.io.EOFException;

 import java.io.IOException;

 import java.io.InputStream;

 import java.io.InterruptedIOException;

 import java.io.StreamCorruptedException;

 import java.net.InetSocketAddress;

 import java.net.Socket;

 import java.net.SocketException;

 import java.net.SocketTimeoutException;

 import java.util.concurrent.Executor;

 import java.util.concurrent.TimeUnit;

 import java.util.concurrent.atomic.AtomicBoolean;

 import java.util.concurrent.atomic.AtomicInteger;

 import javax.net.SocketFactory;

 import com.allanbank.mongodb.MongoClientConfiguration;

 import com.allanbank.mongodb.MongoDbException;

 import com.allanbank.mongodb.Version;

 import com.allanbank.mongodb.bson.io.BsonInputStream;

 import com.allanbank.mongodb.bson.io.RandomAccessOutputStream;

 import com.allanbank.mongodb.bson.io.StringDecoderCache;

 import com.allanbank.mongodb.bson.io.StringEncoderCache;

 import com.allanbank.mongodb.client.Message;

 import com.allanbank.mongodb.client.Operation;

 import com.allanbank.mongodb.client.VersionRange;

 import com.allanbank.mongodb.client.callback.NoOpCallback;

 import com.allanbank.mongodb.client.callback.Receiver;

 import com.allanbank.mongodb.client.callback.ReplyCallback;

 import com.allanbank.mongodb.client.callback.ReplyHandler;

 import com.allanbank.mongodb.client.connection.Connection;

 import com.allanbank.mongodb.client.connection.SocketConnectionListener;

 import com.allanbank.mongodb.client.message.Delete;

 import com.allanbank.mongodb.client.message.GetMore;

 import com.allanbank.mongodb.client.message.Header;

 import com.allanbank.mongodb.client.message.Insert;

 import com.allanbank.mongodb.client.message.IsMaster;

 import com.allanbank.mongodb.client.message.KillCursors;

 import com.allanbank.mongodb.client.message.PendingMessage;

 import com.allanbank.mongodb.client.message.PendingMessageQueue;

 import com.allanbank.mongodb.client.message.Query;

 import com.allanbank.mongodb.client.message.Reply;

 import com.allanbank.mongodb.client.message.Update;

 import com.allanbank.mongodb.client.state.Server;

 import com.allanbank.mongodb.error.ConnectionLostException;

 import com.allanbank.mongodb.error.DocumentToLargeException;

 import com.allanbank.mongodb.error.ServerVersionException;

 import com.allanbank.mongodb.util.IOUtils;

 import com.allanbank.mongodb.util.log.Log;

 import com.allanbank.mongodb.util.log.LogFactory;

 /**

  * AbstractSocketConnection provides the basic functionality for a socket

  * connection that passes messages between the sender and receiver.

  * 

  * @api.no This class is <b>NOT</b> part of the drivers API. This class may be

  *         mutated in incompatible ways between any two releases of the driver.

  * @copyright 2013-2014, Allanbank Consulting, Inc., All Rights Reserved

  */

 public abstract class AbstractSocketConnection implements Connection, Receiver {

     /** The length of the message header in bytes. */

     public static final int HEADER_LENGTH = 16;

     /** The writer for BSON documents. Shares this objects {@link #myInput}. */

     protected final BsonInputStream myBsonIn;

     /** The connections configuration. */

     protected final MongoClientConfiguration myConfig;

     /** The cache for the encoding of strings. */

     protected final StringEncoderCache myEncoderCache;

     /** Support for emitting property change events. */

     protected final PropertyChangeSupport myEventSupport;

     /** The executor for the responses. */

     protected final Executor myExecutor;

     /** The buffered input stream. */

     protected final InputStream myInput;

     /** The logger for the connection. */

     protected final Log myLog;

     /** Holds if the connection is open. */

     protected final AtomicBoolean myOpen;

     /** The buffered output stream. */

     protected final BufferedOutputStream myOutput;

     /** The queue of messages sent but waiting for a reply. */

     protected final PendingMessageQueue myPendingQueue;

     /** The open socket. */

     protected final Server myServer;

     /** Set to true when the connection should be gracefully closed. */

     protected final AtomicBoolean myShutdown;

     /** The open socket. */

     protected final Socket mySocket;

     /** Tracks the number of sequential read timeouts. */

     private int myIdleTicks = 0;

     /** The {@link PendingMessage} used for the local cached copy. */

     private final PendingMessage myPendingMessage = new PendingMessage();

     /** Set to true when the sender discovers they are the receive thread. */

     private final AtomicInteger myReaderNeedsToFlush = new AtomicInteger(0);

     /**

      * Creates a new AbstractSocketConnection.

      * 

      * @param server

      *            The MongoDB server to connect to.

      * @param config

      *            The configuration for the Connection to the MongoDB server.

      * @param encoderCache

      *            Cache used for encoding strings.

      * @param decoderCache

      *            Cache used for decoding strings.

      * @throws SocketException

      *             On a failure connecting to the MongoDB server.

      * @throws IOException

      *             On a failure to read or write data to the MongoDB server.

      */

     public AbstractSocketConnection(final Server server,

             final MongoClientConfiguration config,

             final StringEncoderCache encoderCache,

             final StringDecoderCache decoderCache) throws SocketException,

             IOException {

         super();

         myServer = server;

         myConfig = config;

         myEncoderCache = encoderCache;

         myLog = LogFactory.getLog(getClass());

         myExecutor = config.getExecutor();

         myEventSupport = new PropertyChangeSupport(this);

         myOpen = new AtomicBoolean(false);

         myShutdown = new AtomicBoolean(false);

         mySocket = openSocket(server, config);

         updateSocketWithOptions(config);

         myOpen.set(true);

         myInput = mySocket.getInputStream();

         myBsonIn = new BsonInputStream(myInput, decoderCache);

         // Careful with the size of the buffer here. Seems Java likes to call

         // madvise(..., MADV_DONTNEED) for buffers over a certain size.

         // Net effect is that the performance of the system goes down the

         // drain. Some numbers using the

         // UnixDomainSocketAccepatanceTest.testMultiFetchiterator

         // 1M ==> More than a minute...

         // 512K ==> 24 seconds

         // 256K ==> 16.9 sec.

         // 128K ==> 17 sec.

         // 64K ==> 17 sec.

         // 32K ==> 16.5 sec.

         // Based on those numbers we set the buffer to 32K as larger does not

         // improve performance.

         myOutput = new BufferedOutputStream(mySocket.getOutputStream(),

                 32 * 1024);

         myPendingQueue = new PendingMessageQueue(

                 config.getMaxPendingOperationsPerConnection(),

                 config.getLockType());

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to add the listener to this connection.

      * </p>

      */

     @Override

     public void addPropertyChangeListener(final PropertyChangeListener listener) {

         myEventSupport.addPropertyChangeListener(listener);

     }

     /**

      * {@inheritDoc}

      */

     @Override

     public void flush() throws IOException {

         myReaderNeedsToFlush.set(0);

         myOutput.flush();

     }

     /**

      * {@inheritDoc}

      */

     @Override

     public int getPendingCount() {

         return myPendingQueue.size();

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to returns the server's name.

      * </p>

      */

     @Override

     public String getServerName() {

         return myServer.getCanonicalName();

     }

     /**

      * {@inheritDoc}

      * <p>

      * True if the connection is open and not shutting down.

      * </p>

      */

     @Override

     public boolean isAvailable() {

         return isOpen() && !isShuttingDown();

     }

     /**

      * {@inheritDoc}

      * <p>

      * True if the send and pending queues are empty.

      * </p>

      */

     @Override

     public boolean isIdle() {

         return myPendingQueue.isEmpty();

     }

     /**

      * {@inheritDoc}

      * <p>

      * True if the connection has not been closed.

      * </p>

      */

     @Override

     public boolean isOpen() {

         return myOpen.get();

     }

     /**

      * {@inheritDoc}

      */

     @Override

     public boolean isShuttingDown() {

         return myShutdown.get();

     }

     /**

      * {@inheritDoc}

      * <p>

      * Notifies the appropriate messages of the error.

      * </p>

      */

     @Override

     public void raiseErrors(final MongoDbException exception) {

         final PendingMessage message = new PendingMessage();

         while (myPendingQueue.poll(message)) {

             raiseError(exception, message.getReplyCallback());

         }

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to remove the listener from this connection.

      * </p>

      */

     @Override

     public void removePropertyChangeListener(

             final PropertyChangeListener listener) {

         myEventSupport.removePropertyChangeListener(listener);

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to mark the socket as shutting down and tickles the sender to

      * make sure that happens as soon as possible.

      * </p>

      */

     @Override

     public void shutdown(final boolean force) {

         // Mark

         myShutdown.set(true);

         if (force) {

             IOUtils.close(this);

         }

         else {

             if (isOpen()) {

                 // Force a message with a callback to wake the receiver up.

                 send(new IsMaster(), new NoOpCallback());

             }

         }

     }

     /**

      * Starts the connection.

      */

     public abstract void start();

     /**

      * Stops the socket connection by calling {@link #shutdown(boolean)

      * shutdown(false)}.

      */

     public void stop() {

         shutdown(false);

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return the socket information.

      * </p>

      */

     @Override

     public String toString() {

         return "MongoDB(" + mySocket.getLocalPort() + "-->"

                 + mySocket.getRemoteSocketAddress() + ")";

     }

     /**

      * {@inheritDoc}

      * <p>

      * If there is a pending flush then flushes.

      * </p>

      * <p>

      * If there is any available data then does a single receive.

      * </p>

      */

     @Override

     public void tryReceive() {

         try {

             doReceiverFlush();

             if ((myBsonIn.available() > 0) || (myInput.available() > 0)) {

                 doReceiveOne();

             }

         }

         catch (final IOException error) {

             myLog.info(

                     "Received an error when checking for pending messages: {}.",

                     error.getMessage());

         }

     }

     /**

      * {@inheritDoc}

      * <p>

      * Waits for the connections pending queues to empty.

      * </p>

      */

     @Override

     public void waitForClosed(final int timeout, final TimeUnit timeoutUnits) {

         long now = System.currentTimeMillis();

         final long deadline = now + timeoutUnits.toMillis(timeout);

         while (isOpen() && (now < deadline)) {

             try {

                 // A slow spin loop.

                 TimeUnit.MILLISECONDS.sleep(10);

             }

             catch (final InterruptedException e) {

                 // Ignore.

                 e.hashCode();

             }

             now = System.currentTimeMillis();

         }

     }

     /**

      * Receives a single message from the connection.

      * 

      * @return The {@link Message} received.

      * @throws MongoDbException

      *             On an error receiving the message.

      */

     protected Message doReceive() throws MongoDbException {

         try {

             int length;

             try {

                 length = readIntSuppressTimeoutOnNonFirstByte();

             }

             catch (final SocketTimeoutException ok) {

                 // This is OK. We check if we are still running and come right

                 // back.

                 return null;

             }

             myBsonIn.prefetch(length - 4);

             final int requestId = myBsonIn.readInt();

             final int responseId = myBsonIn.readInt();

             final int opCode = myBsonIn.readInt();

             final Operation op = Operation.fromCode(opCode);

             if (op == null) {

                 // Huh? Dazed and confused

                 throw new MongoDbException("Unexpected operation read '"

                         + opCode + "'.");

             }

             final Header header = new Header(length, requestId, responseId, op);

             Message message = null;

             switch (op) {

             case REPLY:

                 message = new Reply(header, myBsonIn);

                 break;

             case QUERY:

                 message = new Query(header, myBsonIn);

                 break;

             case UPDATE:

                 message = new Update(myBsonIn);

                 break;

             case INSERT:

                 message = new Insert(header, myBsonIn);

                 break;

             case GET_MORE:

                 message = new GetMore(myBsonIn);

                 break;

             case DELETE:

                 message = new Delete(myBsonIn);

                 break;

             case KILL_CURSORS:

                 message = new KillCursors(myBsonIn);

                 break;

             }

             myServer.incrementRepliesReceived();

             return message;

         }

         catch (final IOException ioe) {

             final MongoDbException error = new ConnectionLostException(ioe);

             shutdown(error, (ioe instanceof InterruptedIOException));

             throw error;

         }

     }

     /**

      * Receives and process a single message.

      */

     protected void doReceiveOne() {

         doReceiverFlush();

         final Message received = doReceive();

         if (received instanceof Reply) {

             myIdleTicks = 0;

             final Reply reply = (Reply) received;

             final int replyId = reply.getResponseToId();

             boolean took = false;

             // Keep polling the pending queue until we get to

             // message based on a matching replyId.

             try {

                 took = myPendingQueue.poll(myPendingMessage);

                 while (took && (myPendingMessage.getMessageId() != replyId)) {

                     final MongoDbException noReply = new MongoDbException(

                             "No reply received.");

                     // Note that this message will not get a reply.

                     raiseError(noReply, myPendingMessage.getReplyCallback());

                     // Keep looking.

                     took = myPendingQueue.poll(myPendingMessage);

                 }

                 if (took) {

                     // Must be the pending message's reply.

                     reply(reply, myPendingMessage);

                 }

                 else {

                     myLog.warn("Could not find the callback for reply '{}'.",

                             +replyId);

                 }

             }

             finally {

                 myPendingMessage.clear();

             }

         }

         else if (received != null) {

             myLog.warn("Received a non-Reply message: {}.", received);

             shutdown(new ConnectionLostException(new StreamCorruptedException(

                     "Received a non-Reply message: " + received)), false);

         }

         else {

             myIdleTicks += 1;

             if (myConfig.getMaxIdleTickCount() <= myIdleTicks) {

                 // Shutdown the connection., nicely.

                 shutdown(false);

             }

         }

     }

     /**

      * Sends a single message to the connection.

      * 

      * @param messageId

      *            The id to use for the message.

      * @param message

      *            The message to send.

      * @throws IOException

      *             On a failure sending the message.

      */

     protected void doSend(final int messageId,

             final RandomAccessOutputStream message) throws IOException {

         message.writeTo(myOutput);

         message.reset();

         myServer.incrementMessagesSent();

     }

     /**

      * Should be called when the send of a message happens on the receive

      * thread. The sender should not flush the {@link #myOutput}. Instead the

      * receive thread will {@link #flush()} once it has consumed all of the

      * pending messages to be received.

      */

     protected void markReaderNeedsToFlush() {

         myReaderNeedsToFlush.incrementAndGet();

     }

     /**

      * Updates to raise an error on the callback, if any.

      * 

      * @param exception

      *            The thrown exception.

      * @param replyCallback

      *            The callback for the reply to the message.

      */

     protected void raiseError(final Throwable exception,

             final ReplyCallback replyCallback) {

         ReplyHandler.raiseError(exception, replyCallback, myExecutor);

     }

     /**

      * Reads a little-endian 4 byte signed integer from the stream.

      * 

      * @return The integer value.

      * @throws EOFException

      *             On insufficient data for the integer.

      * @throws IOException

      *             On a failure reading the integer.

      */

     protected int readIntSuppressTimeoutOnNonFirstByte() throws EOFException,

             IOException {

         int read = 0;

         int eofCheck = 0;

         int result = 0;

         read = myBsonIn.read();

         eofCheck |= read;

         result += (read << 0);

         for (int i = Byte.SIZE; i < Integer.SIZE; i += Byte.SIZE) {

             try {

                 read = myBsonIn.read();

             }

             catch (final SocketTimeoutException ste) {

                 // Bad - Only the first byte should timeout.

                 throw new IOException(ste);

             }

             eofCheck |= read;

             result += (read << i);

         }

         if (eofCheck < 0) {

             throw new EOFException("Remote connection closed: "

                     + mySocket.getRemoteSocketAddress());

         }

         return result;

     }

     /**

      * Updates to set the reply for the callback, if any.

      * 

      * @param reply

      *            The reply.

      * @param pendingMessage

      *            The pending message.

      */

     protected void reply(final Reply reply, final PendingMessage pendingMessage) {

         final long latency = pendingMessage.latency();

         // Add the latency.

         if (latency > 0) {

             myServer.updateAverageLatency(latency);

         }

         final ReplyCallback callback = pendingMessage.getReplyCallback();

         ReplyHandler.reply(this, reply, callback, myExecutor);

     }

     /**

      * Sends a single message.

      * 

      * @param pendingMessage

      *            The message to be sent.

      * @param message

      *            The message that has already been encoded/serialized. This may

      *            be <code>null</code> in which case the message is streamed to

      *            the socket.

      * @throws InterruptedException

      *             If the thread is interrupted waiting for a message to send.

      * @throws IOException

      *             On a failure sending the message.

      */

     protected final void send(final PendingMessage pendingMessage,

             final RandomAccessOutputStream message)

             throws InterruptedException, IOException {

         final int messageId = pendingMessage.getMessageId();

         // Mark the timestamp.

         pendingMessage.timestampNow();

         // Make sure the message is on the queue before the

         // message is sent to ensure the receive thread can

         // assume an empty pending queue means that there is

         // no message for the reply.

         if ((pendingMessage.getReplyCallback() != null)

                 && !myPendingQueue.offer(pendingMessage)) {

             // Push what we have out before blocking.

             flush();

             myPendingQueue.put(pendingMessage);

         }

         doSend(messageId, message);

         // If shutting down then flush after each message.

         if (myShutdown.get()) {

             flush();

         }

     }

     /**

      * Shutsdown the connection on an error.

      * 

      * @param error

      *            The error causing the shutdown.

      * @param receiveError

      *            If true then the socket experienced a receive error.

      */

     protected void shutdown(final MongoDbException error,

             final boolean receiveError) {

         if (receiveError) {

             myServer.connectionTerminated();

         }

         // Have to assume all of the requests have failed that are pending.

         final PendingMessage message = new PendingMessage();

         while (myPendingQueue.poll(message)) {

             raiseError(error, message.getReplyCallback());

         }

         closeQuietly();

     }

     /**

      * Updates the socket with the configuration's socket options.

      * 

      * @param config

      *            The configuration to apply.

      * @throws SocketException

      *             On a failure setting the socket options.

      */

     protected void updateSocketWithOptions(final MongoClientConfiguration config)

             throws SocketException {

         mySocket.setKeepAlive(config.isUsingSoKeepalive());

         mySocket.setSoTimeout(config.getReadTimeout());

         try {

             mySocket.setTcpNoDelay(true);

         }

         catch (final SocketException seIgnored) {

             // The junixsocket implementation does not support TCP_NO_DELAY,

             // which makes sense but it throws an exception instead of silently

             // ignoring - ignore it here.

             if (!"AFUNIXSocketException".equals(seIgnored.getClass()

                     .getSimpleName())) {

                 throw seIgnored;

             }

         }

         mySocket.setPerformancePreferences(1, 5, 6);

     }

     /**

      * Ensures that the documents in the message do not exceed the maximum size

      * allowed by MongoDB.

      * 

      * @param message1

      *            The message to be sent to the server.

      * @param message2

      *            The second message to be sent to the server.

      * @throws DocumentToLargeException

      *             On a message being too large.

      * @throws ServerVersionException

      *             If one of the messages cannot be sent to the server version.

      */

     protected void validate(final Message message1, final Message message2)

             throws DocumentToLargeException, ServerVersionException {

         final Version serverVersion = myServer.getVersion();

         final int maxBsonSize = myServer.getMaxBsonObjectSize();

         message1.validateSize(maxBsonSize);

         validateVersion(message1, serverVersion);

         if (message2 != null) {

             message2.validateSize(maxBsonSize);

             validateVersion(message1, serverVersion);

         }

     }

     /**

      * Closes the connection to the server without allowing an exception to be

      * thrown.

      */

     private void closeQuietly() {

         try {

             close();

         }

         catch (final IOException e) {

             myLog.warn(e, "I/O exception trying to shutdown the connection.");

         }

     }

     /**

      * Check if the handler for a message dropped data in the send buffer that

      * it did not flush to avoid a deadlock with the server. If so then flush

      * that message.

      */

     private void doReceiverFlush() {

         try {

             final int unflushedMessages = myReaderNeedsToFlush.get();

             if ((unflushedMessages != 0)

                     && (myPendingQueue.size() <= unflushedMessages)) {

                 flush();

             }

         }

         catch (final IOException ignored) {

             myLog.warn("Error flushing data to the server: "

                     + ignored.getMessage());

         }

     }

     /**

      * Tries to open a connection to the server.

      * 

      * @param server

      *            The server to open the connection to.

      * @param config

      *            The configuration for attempting to open the connection.

      * @return The opened {@link Socket}.

      * @throws IOException

      *             On a failure opening a connection to the server.

      */

     private Socket openSocket(final Server server,

             final MongoClientConfiguration config) throws IOException {

         final SocketFactory factory = config.getSocketFactory();

         IOException last = null;

         Socket socket = null;

         for (final InetSocketAddress address : myServer.getAddresses()) {

             try {

                 socket = factory.createSocket();

                 socket.connect(address, config.getConnectTimeout());

                 // If the factory wants to know about the connection then let it

                 // know first.

                 if (factory instanceof SocketConnectionListener) {

                     ((SocketConnectionListener) factory).connected(address,

                             socket);

                 }

                 // Let the server know the working connection.

                 server.connectionOpened(address);

                 last = null;

                 break;

             }

             catch (final IOException error) {

                 last = error;

                 try {

                     if (socket != null) {

                         socket.close();

                     }

                 }

                 catch (final IOException ignore) {

                     myLog.info(

                             "Could not close the defunct socket connection: {}",

                             socket);

                 }

             }

         }

         if (last != null) {

             server.connectFailed();

             throw last;

         }

         return socket;

     }

     /**

      * Validates that the server we are about to send the message to knows how

      * to handle the message.

      * 

      * @param message

      *            The message to be sent.

      * @param serverVersion

      *            The server version.

      * @throws ServerVersionException

      *             If the messages cannot be sent to the server version.

      */

     private void validateVersion(final Message message,

             final Version serverVersion) throws ServerVersionException {

         final VersionRange range = message.getRequiredVersionRange();

         if ((range != null) && !range.contains(serverVersion)) {

             throw new ServerVersionException(message.getOperationName(), range,

                     serverVersion, message);

         }

     }

 }