/*

  * #%L

  * ReplicaSetReconnectStrategy.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.rs;

 import static java.util.concurrent.TimeUnit.MILLISECONDS;

 import java.io.IOException;

 import java.util.Collections;

 import java.util.HashMap;

 import java.util.List;

 import java.util.Map;

 import java.util.Set;

 import java.util.concurrent.ConcurrentHashMap;

 import java.util.concurrent.ExecutionException;

 import java.util.concurrent.Future;

 import java.util.concurrent.TimeUnit;

 import java.util.concurrent.TimeoutException;

 import java.util.logging.Level;

 import com.allanbank.mongodb.MongoClientConfiguration;

 import com.allanbank.mongodb.bson.Document;

 import com.allanbank.mongodb.bson.Element;

 import com.allanbank.mongodb.bson.element.StringElement;

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

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

 import com.allanbank.mongodb.client.connection.proxy.ConnectionInfo;

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

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

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

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

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

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

 import com.allanbank.mongodb.util.IOUtils;

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

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

 /**

  * ReplicaSetReconnectStrategy provides a {@link ReconnectStrategy} designed for

  * replica sets. The reconnect strategy attempts to locate the primary member of

  * the replica set by:

  * <ol>

  * <li>Querying each member of the replica set for the primary server.</li>

  * <li>Once a primary server has been identified by a member of the replica set

  * (the putative primary) the putative primary server is queried for the primary

  * server.</li>

  * <ol>

  * <li>If the putative primary concurs that it is the primary then the search

  * completes and the primary server's connection is used.</li>

  * <li>If the putative primary does not concur then the search continues

  * scanning each server in turn for the primary server.</li>

  * </ol>

  * 

  * @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 2012-2014, Allanbank Consulting, Inc., All Rights Reserved

  */

 public class ReplicaSetReconnectStrategy extends AbstractReconnectStrategy {

     /**

      * The initial amount of time to pause waiting for a server to take over as

      * the primary.

      */

     public static final int INITIAL_RECONNECT_PAUSE_TIME_MS = 10;

     /**

      * The Maximum amount of time to pause waiting for a server to take over as

      * the primary.

      */

     public static final int MAX_RECONNECT_PAUSE_TIME_MS = 1000;

     /** The logger for the {@link ReplicaSetReconnectStrategy}. */

     protected static final Log LOG = LogFactory

             .getLog(ReplicaSetReconnectStrategy.class);

     /** The set of servers we cannot connect to. */

     private final Set<Server> myDeadServers = Collections

             .newSetFromMap(new ConcurrentHashMap<Server, Boolean>());

     /**

      * Creates a new ReplicaSetReconnectStrategy.

      */

     public ReplicaSetReconnectStrategy() {

         super();

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to search for the primary server in the replica set. This will

      * only continue until the

      * {@link MongoClientConfiguration#getReconnectTimeout()} has expired.

      * </p>

      */

     @Override

     public ReplicaSetConnection reconnect(final Connection oldConnection) {

         final ConnectionInfo<Server> info = reconnectPrimary();

         if (info != null) {

             return new ReplicaSetConnection(info.getConnection(),

                     info.getConnectionKey(), getState(),

                     getConnectionFactory(), getConfig(), this);

         }

         return null;

     }

     /**

      * Overridden to search for the primary server in the replica set. This will

      * only continue until the

      * {@link MongoClientConfiguration#getReconnectTimeout()} has expired.

      * 

      * @return The information for the primary connection or null if the

      *         reconnect fails.

      */

     public synchronized ConnectionInfo<Server> reconnectPrimary() {

         LOG.debug("Trying replica set reconnect.");

         final Cluster state = getState();

         // Figure out a deadline for the reconnect.

         final int wait = getConfig().getReconnectTimeout();

         long now = System.currentTimeMillis();

         final long deadline = (wait <= 0) ? Long.MAX_VALUE : (now + wait);

         final Map<Server, Future<Reply>> answers = new HashMap<Server, Future<Reply>>();

         final Map<Server, Connection> connections = new HashMap<Server, Connection>();

         // Clear any interrupts

         final boolean interrupted = Thread.interrupted();

         try {

             // First try a simple reconnect.

             for (final Server writable : state.getWritableServers()) {

                 if (verifyPutative(answers, connections, writable, deadline)) {

                     LOG.debug("New primary for replica set: {}.",

                             writable.getCanonicalName());

                     return createReplicaSetConnection(connections, writable);

                 }

             }

             // How much time to pause for replies and waiting for a server

             // to become primary.

             int pauseTime = INITIAL_RECONNECT_PAUSE_TIME_MS;

             while (now < deadline) {

                 // Ask all of the servers who they think the primary is.

                 for (final Server server : state.getServers()) {

                     sendIsPrimary(answers, connections, server, false);

                     // Anyone replied yet?

                     final ConnectionInfo<Server> newConn = checkForReply(state,

                             answers, connections, deadline);

                     if (newConn != null) {

                         return newConn;

                     }

                     // Loop to the next server.

                 }

                 // Wait for a beat for a reply or a server to decide to be

                 // master.

                 sleep(pauseTime, MILLISECONDS);

                 pauseTime = Math.min(MAX_RECONNECT_PAUSE_TIME_MS, pauseTime

                         + pauseTime);

                 // Check again for replies before trying to reconnect.

                 final ConnectionInfo<Server> newConn = checkForReply(state,

                         answers, connections, deadline);

                 if (newConn != null) {

                     return newConn;

                 }

                 now = System.currentTimeMillis();

             }

         }

         finally {

             // Shut down the connections we created.

             for (final Connection conn : connections.values()) {

                 conn.shutdown(true);

             }

             if (interrupted) {

                 Thread.currentThread().interrupt();

             }

         }

         return null;

     }

     /**

      * Checks for a reply from a server. If one has been received then it tries

      * to confirm the primary server by asking it if it thinks it is the primary

      * server.

      * 

      * @param state

      *            The state of the cluster.

      * @param answers

      *            The pending ({@link Future}) answers from each server.

      * @param connections

      *            The connection to each server.

      * @param deadline

      *            The deadline for the reconnect attempt.

      * @return The new connection if there was a reply and that server confirmed

      *         it was the primary.

      */

     protected ConnectionInfo<Server> checkForReply(final Cluster state,

             final Map<Server, Future<Reply>> answers,

             final Map<Server, Connection> connections, final long deadline) {

         final Map<Server, Future<Reply>> copy = new HashMap<Server, Future<Reply>>(

                 answers);

         for (final Map.Entry<Server, Future<Reply>> entry : copy.entrySet()) {

             final Server server = entry.getKey();

             final Future<Reply> reply = entry.getValue();

             if (reply.isDone()) {

                 // Remove this reply.

                 answers.remove(server);

                 // Check the result.

                 final String putative = checkReply(reply, connections, server,

                         deadline);

                 // Phase2 - Verify the putative server.

                 if (putative != null) {

                     final Server putativeServer = getState().get(putative);

                     if (verifyPutative(answers, connections, putativeServer,

                             deadline)) {

                         // Phase 3 - Setup a new replica set connection to the

                         // primary and seed it with a secondary if there is a

                         // suitable server.

                         LOG.info("New primary for replica set: {}", putative);

                         updateUnknown(state, answers, connections);

                         return createReplicaSetConnection(connections,

                                 putativeServer);

                     }

                 }

             }

             else {

                 LOG.debug("No reply yet from {}.", server);

             }

         }

         return null;

     }

     /**

      * Extracts who the server thinks is the primary from the reply.

      * 

      * @param replyFuture

      *            The future to get the reply from.

      * @param connections

      *            The map of connections. The connection will be closed on an

      *            error.

      * @param server

      *            The server.

      * @param deadline

      *            The deadline for the reconnect attempt.

      * @return The name of the server the reply indicates is the primary, null

      *         if there is no primary or any error.

      */

     protected String checkReply(final Future<Reply> replyFuture,

             final Map<Server, Connection> connections, final Server server,

             final long deadline) {

         if (replyFuture != null) {

             try {

                 final Reply reply = replyFuture.get(

                         Math.max(0, deadline - System.currentTimeMillis()),

                         TimeUnit.MILLISECONDS);

                 final List<Document> results = reply.getResults();

                 if (!results.isEmpty()) {

                     final Document doc = results.get(0);

                     // Get the name of the primary server.

                     final Element primary = doc.get("primary");

                     if (primary instanceof StringElement) {

                         return ((StringElement) primary).getValue();

                     }

                 }

             }

             catch (final InterruptedException e) {

                 // Just ignore the reply.

             }

             catch (final TimeoutException e) {

                 // Kill the associated connection.

                 final Connection conn = connections.remove(server);

                 IOUtils.close(conn);

             }

             catch (final ExecutionException e) {

                 // Kill the associated connection.

                 final Connection conn = connections.remove(server);

                 IOUtils.close(conn);

             }

         }

         return null;

     }

     /**

      * Sends a command to the server to return what it thinks the state of the

      * cluster is. This method will not re-request the information from the

      * server if there is already an outstanding request.

      * 

      * @param answers

      *            The pending ({@link Future}) answers from each server.

      * @param connections

      *            The connection to each server.

      * @param server

      *            The server to send the request to.

      * @param isPrimary

      *            If true logs connection errors as warnings. Debug otherwise.

      * @return The future reply for the request sent to the server.

      */

     protected Future<Reply> sendIsPrimary(

             final Map<Server, Future<Reply>> answers,

             final Map<Server, Connection> connections, final Server server,

             final boolean isPrimary) {

         Future<Reply> reply = null;

         try {

             // Locate a connection to the server.

             Connection conn = connections.get(server);

             if ((conn == null) || !conn.isAvailable()) {

                 conn = getConnectionFactory().connect(server, getConfig());

                 connections.put(server, conn);

             }

             // Only send to the server if there is not an outstanding

             // request.

             reply = answers.get(server);

             if (reply == null) {

                 LOG.debug("Sending reconnect(rs) query to {}.",

                         server.getCanonicalName());

                 final ServerUpdateCallback replyCallback = new ServerUpdateCallback(

                         server);

                 conn.send(new IsMaster(), replyCallback);

                 reply = replyCallback;

                 answers.put(server, reply);

                 myDeadServers.remove(server);

             }

         }

         catch (final IOException e) {

             // Nothing to do for now. Log at a debug level if this is not the

             // primary. Warn if we think it is the primary (and have not warned

             // before)

             final Level level = (isPrimary && myDeadServers.add(server)) ? Level.WARNING

                     : Level.FINE;

             LOG.log(level, e, "Cannot create a connection to '{}'.", server);

         }

         return reply;

     }

     /**

      * Sleeps without throwing an exception.

      * 

      * @param sleepTime

      *            The amount of time to sleep.

      * @param units

      *            The untis for the amount of time to sleep.

      */

     protected void sleep(final int sleepTime, final TimeUnit units) {

         try {

             units.sleep(sleepTime);

         }

         catch (final InterruptedException e) {

             // Ignore.

         }

     }

     /**

      * Tries to verify that the suspected primary server is in fact the primary

      * server by asking it directly and synchronously.

      * 

      * @param answers

      *            The pending ({@link Future}) answers from each server.

      * @param connections

      *            The connection to each server.

      * @param putativePrimary

      *            The server we think is the primary.

      * @param deadline

      *            The deadline for the reconnect attempt.

      * @return True if the server concurs that it is the primary.

      */

     protected boolean verifyPutative(final Map<Server, Future<Reply>> answers,

             final Map<Server, Connection> connections,

             final Server putativePrimary, final long deadline) {

         LOG.debug("Verify putative server ({}) on reconnect(rs).",

                 putativePrimary);

         // Make sure we send a new request. The old reply might have been

         // before becoming the primary.

         answers.remove(putativePrimary);

         // If the primary agrees that they are the primary then it is

         // probably true.

         final Future<Reply> reply = sendIsPrimary(answers, connections,

                 putativePrimary, true);

         final String primary = checkReply(reply, connections, putativePrimary,

                 deadline);

         if (putativePrimary.getCanonicalName().equals(primary)) {

             return true;

         }

         return false;

     }

     /**

      * Creates the {@link ReplicaSetConnection} for the primary server.

      * 

      * @param connections

      *            The connection that are being managed.

      * @param primaryServer

      *            The primary server.

      * @return The {@link ReplicaSetConnection}.

      */

     private ConnectionInfo<Server> createReplicaSetConnection(

             final Map<Server, Connection> connections,

             final Server primaryServer) {

         final Connection primaryConn = connections.remove(primaryServer);

         return new ConnectionInfo<Server>(primaryConn, primaryServer);

     }

     /**

      * Tries to send messages to all of the members of the cluster in an

      * indeterminate state.

      * 

      * @param state

      *            The state of the cluster.

      * @param answers

      *            The pending responses.

      * @param connections

      *            The connection already created.

      */

     private void updateUnknown(final Cluster state,

             final Map<Server, Future<Reply>> answers,

             final Map<Server, Connection> connections) {

         for (final Server server : state.getServers()) {

             switch (server.getState()) {

             case UNKNOWN: // Fall through.

             case UNAVAILABLE: {

                 answers.remove(server);

                 sendIsPrimary(answers, connections, server, false);

                 break;

             }

             case READ_ONLY:

             case WRITABLE:

             default: {

                 // Known good.

                 break;

             }

             }

         }

     }

 }