SPY-160: Redistribute ops when waiting for authentication to complete.

Motivation
----------
When an operation gets added and the node needs to (re)connect, it waits
for a given amount of time (authWaitTime). If the auth does not complete
in the period of time, the operation is cancelled. This is suboptimal,
especially if the operation is not timed out yet, it could wait longer
until it gets timed out/cancelled.

Modifications
-------------
This changeset makes sure that an operation gets redistributed (assuming
FailureMode is NOT cancel) if the authLatch is not counted down in the
wait period.

Redistributing will only happen (further down the code stack) if the
op is 1) not cancelled and 2) not timed out already.

Result
------
Still valid operations will be transparently rescheduled until they
either go through validly or will eventually time out or get
cancelled.

Change-Id: I6f849314e11435a6ebc69a4bfc1eab705df850e8
Reviewed-on: http://review.couchbase.org/34852
Reviewed-by: Michael Nitschinger <michael.nitschinger@couchbase.com>
Tested-by: Michael Nitschinger <michael.nitschinger@couchbase.com>

Author: daschl

src/main/java/net/spy/memcached/BinaryConnectionFactory.java

@@ -70,7 +70,7 @@ public MemcachedNode createMemcachedNode(SocketAddress sa, SocketChannel c,
     return new BinaryMemcachedNodeImpl(sa, c, bufSize,
         createReadOperationQueue(), createWriteOperationQueue(),
         createOperationQueue(), getOpQueueMaxBlockTime(), doAuth,
-        getOperationTimeout(), getAuthWaitTime());
+        getOperationTimeout(), getAuthWaitTime(), this);
   }
 
   @Override

src/main/java/net/spy/memcached/DefaultConnectionFactory.java

@@ -177,7 +177,8 @@ public MemcachedNode createMemcachedNode(SocketAddress sa, SocketChannel c,
           createOperationQueue(),
           getOpQueueMaxBlockTime(),
           getOperationTimeout(),
-          getAuthWaitTime());
+          getAuthWaitTime(),
+          this);
     } else if (of instanceof BinaryOperationFactory) {
       boolean doAuth = false;
       if (getAuthDescriptor() != null) {
@@ -190,7 +191,8 @@ public MemcachedNode createMemcachedNode(SocketAddress sa, SocketChannel c,
           getOpQueueMaxBlockTime(),
           doAuth,
           getOperationTimeout(),
-          getAuthWaitTime());
+          getAuthWaitTime(),
+          this);
     } else {
       throw new IllegalStateException("Unhandled operation factory type " + of);
     }

src/main/java/net/spy/memcached/MemcachedConnection.java

@@ -313,10 +313,10 @@ protected List<MemcachedNode> createConnections(
     for (SocketAddress sa : addrs) {
       SocketChannel ch = SocketChannel.open();
       ch.configureBlocking(false);
-      MemcachedNode qa =
-          this.connectionFactory.createMemcachedNode(sa, ch, bufSize);
+      MemcachedNode qa = connectionFactory.createMemcachedNode(sa, ch, bufSize);
+      qa.setConnection(this);
       int ops = 0;
-      ch.socket().setTcpNoDelay(!this.connectionFactory.useNagleAlgorithm());
+      ch.socket().setTcpNoDelay(!connectionFactory.useNagleAlgorithm());
 
       try {
         if (ch.connect(sa)) {
@@ -968,37 +968,56 @@ private void cancelOperations(final Collection<Operation> ops) {
    *
    * @param ops the operations to redistribute.
    */
-  private void redistributeOperations(final Collection<Operation> ops) {
+  public void redistributeOperations(final Collection<Operation> ops) {
     for (Operation op : ops) {
-      if (op.isCancelled() || op.isTimedOut()) {
-        continue;
-      }
+      redistributeOperation(op);
+    }
+  }
 
-      if (op instanceof MultiGetOperationImpl) {
-        for (String key : ((MultiGetOperationImpl) op).getRetryKeys()) {
-          addOperation(key, opFact.get(key,
-            (GetOperation.Callback) op.getCallback()));
-        }
-      } else if (op instanceof KeyedOperation) {
-        KeyedOperation ko = (KeyedOperation) op;
-        int added = 0;
-        for (Operation newop : opFact.clone(ko)) {
-          if (newop instanceof KeyedOperation) {
-            KeyedOperation newKeyedOp = (KeyedOperation) newop;
-            for (String k : newKeyedOp.getKeys()) {
-              addOperation(k, newop);
-            }
-          } else {
-            newop.cancel();
-            getLogger().warn("Could not redistribute cloned non-keyed " +
-              "operation", newop);
+  /**
+   * Redistribute the given operation to (potentially) other nodes.
+   *
+   * Note that operations can only be redistributed if they have not been
+   * cancelled already, timed out already or do not have definite targets
+   * (a key).
+   *
+   * @param op the operation to redistribute.
+   */
+  public void redistributeOperation(Operation op) {
+    if (op.isCancelled() || op.isTimedOut()) {
+      return;
+    }
+
+    // The operation gets redistributed but has never been actually written,
+    // it we just straight re-add it without cloning.
+    if (op.getState() == OperationState.WRITE_QUEUED) {
+      addOperation(op.getHandlingNode(), op);
+    }
+
+    if (op instanceof MultiGetOperationImpl) {
+      for (String key : ((MultiGetOperationImpl) op).getRetryKeys()) {
+        addOperation(key, opFact.get(key,
+          (GetOperation.Callback) op.getCallback()));
+      }
+    } else if (op instanceof KeyedOperation) {
+      KeyedOperation ko = (KeyedOperation) op;
+      int added = 0;
+      for (Operation newop : opFact.clone(ko)) {
+        if (newop instanceof KeyedOperation) {
+          KeyedOperation newKeyedOp = (KeyedOperation) newop;
+          for (String k : newKeyedOp.getKeys()) {
+            addOperation(k, newop);
           }
-          added++;
+        } else {
+          newop.cancel();
+          getLogger().warn("Could not redistribute cloned non-keyed " +
+            "operation", newop);
         }
-        assert added > 0 : "Didn't add any new operations when redistributing";
-      } else {
-        op.cancel();
+        added++;
       }
+      assert added > 0 : "Didn't add any new operations when redistributing";
+    } else {
+      op.cancel();
     }
   }
 

src/main/java/net/spy/memcached/MemcachedNode.java

@@ -238,4 +238,8 @@ public interface MemcachedNode {
   void setContinuousTimeout(boolean timedOut);
 
   int getContinuousTimeout();
+
+  MemcachedConnection getConnection();
+
+  void setConnection(MemcachedConnection connection);
 }

src/main/java/net/spy/memcached/MemcachedNodeROImpl.java

@@ -190,4 +190,14 @@ public void completedRead() {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public MemcachedConnection getConnection() {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void setConnection(MemcachedConnection connection) {
+    throw new UnsupportedOperationException();
+  }
+
 }

src/main/java/net/spy/memcached/protocol/TCPMemcachedNodeImpl.java

@@ -35,6 +35,9 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import net.spy.memcached.ConnectionFactory;
+import net.spy.memcached.FailureMode;
+import net.spy.memcached.MemcachedConnection;
 import net.spy.memcached.MemcachedNode;
 import net.spy.memcached.compat.SpyObject;
 import net.spy.memcached.ops.Operation;
@@ -56,6 +59,7 @@ public abstract class TCPMemcachedNodeImpl extends SpyObject implements
   private final BlockingQueue<Operation> inputQueue;
   private final long opQueueMaxBlockTime;
   private final long authWaitTime;
+  private final ConnectionFactory connectionFactory;
   private AtomicInteger reconnectAttempt = new AtomicInteger(1);
   private SocketChannel channel;
   private int toWrite = 0;
@@ -66,14 +70,15 @@ public abstract class TCPMemcachedNodeImpl extends SpyObject implements
   private ArrayList<Operation> reconnectBlocked;
   private long defaultOpTimeout;
   private long lastReadTimestamp = System.currentTimeMillis();
+  private MemcachedConnection connection;
 
   // operation Future.get timeout counter
   private final AtomicInteger continuousTimeout = new AtomicInteger(0);
 
   public TCPMemcachedNodeImpl(SocketAddress sa, SocketChannel c, int bufSize,
       BlockingQueue<Operation> rq, BlockingQueue<Operation> wq,
       BlockingQueue<Operation> iq, long opQueueMaxBlockTime,
-      boolean waitForAuth, long dt, long authWaitTime) {
+      boolean waitForAuth, long dt, long authWaitTime, ConnectionFactory fact) {
     super();
     assert sa != null : "No SocketAddress";
     assert c != null : "No SocketChannel";
@@ -82,6 +87,7 @@ public TCPMemcachedNodeImpl(SocketAddress sa, SocketChannel c, int bufSize,
     assert wq != null : "No operation write queue";
     assert iq != null : "No input queue";
     socketAddress = sa;
+    connectionFactory = fact;
     this.authWaitTime = authWaitTime;
     setChannel(c);
     // Since these buffers are allocated rarely (only on client creation
@@ -336,12 +342,20 @@ public final boolean hasWriteOp() {
   public final void addOp(Operation op) {
     try {
       if (!authLatch.await(authWaitTime, TimeUnit.MILLISECONDS)) {
-        op.cancel();
-        getLogger().warn("Operation canceled because authentication "
-          + "or reconnection and authentication has "
-          + "taken more than " + authWaitTime + " milliseconds to "
-          + "complete.");
-        getLogger().debug("Canceled operation %s", op.toString());
+        FailureMode mode = connectionFactory.getFailureMode();
+        if (mode == FailureMode.Redistribute || mode == FailureMode.Retry) {
+          getLogger().debug("Redistributing Operation " + op + " because auth "
+            + "latch taken longer than " + authWaitTime + " milliseconds to "
+            + "complete on node " + getSocketAddress());
+          connection.redistributeOperation(op);
+        } else {
+          op.cancel();
+          getLogger().warn("Operation canceled because authentication "
+            + "or reconnection and authentication has "
+            + "taken more than " + authWaitTime + " milliseconds to "
+            + "complete on node " + this);
+          getLogger().debug("Canceled operation %s", op.toString());
+        }
         return;
       }
       if (!inputQueue.offer(op, opQueueMaxBlockTime, TimeUnit.MILLISECONDS)) {
@@ -633,4 +647,13 @@ public void completedRead() {
     lastReadTimestamp = System.currentTimeMillis();
   }
 
+  @Override
+  public MemcachedConnection getConnection() {
+    return connection;
+  }
+
+  @Override
+  public void setConnection(MemcachedConnection connection) {
+    this.connection = connection;
+  }
 }

src/main/java/net/spy/memcached/protocol/ascii/AsciiMemcachedNodeImpl.java

@@ -27,6 +27,7 @@
 import java.nio.channels.SocketChannel;
 import java.util.concurrent.BlockingQueue;
 
+import net.spy.memcached.ConnectionFactory;
 import net.spy.memcached.ops.GetOperation;
 import net.spy.memcached.ops.Operation;
 import net.spy.memcached.ops.OperationState;
@@ -41,9 +42,9 @@ public final class AsciiMemcachedNodeImpl extends TCPMemcachedNodeImpl {
   public AsciiMemcachedNodeImpl(SocketAddress sa, SocketChannel c, int bufSize,
       BlockingQueue<Operation> rq, BlockingQueue<Operation> wq,
       BlockingQueue<Operation> iq, Long opQueueMaxBlockTimeNs, long dt,
-      long at) {
+      long at, ConnectionFactory fa) {
     // ASCII never does auth
-    super(sa, c, bufSize, rq, wq, iq, opQueueMaxBlockTimeNs, false, dt, at);
+    super(sa, c, bufSize, rq, wq, iq, opQueueMaxBlockTimeNs, false, dt, at, fa);
   }
 
   @Override

src/main/java/net/spy/memcached/protocol/binary/BinaryMemcachedNodeImpl.java

@@ -27,6 +27,7 @@
 import java.nio.channels.SocketChannel;
 import java.util.concurrent.BlockingQueue;
 
+import net.spy.memcached.ConnectionFactory;
 import net.spy.memcached.ops.CASOperation;
 import net.spy.memcached.ops.GetOperation;
 import net.spy.memcached.ops.Operation;
@@ -47,9 +48,9 @@ public class BinaryMemcachedNodeImpl extends TCPMemcachedNodeImpl {
   public BinaryMemcachedNodeImpl(SocketAddress sa, SocketChannel c,
       int bufSize, BlockingQueue<Operation> rq, BlockingQueue<Operation> wq,
       BlockingQueue<Operation> iq, Long opQueueMaxBlockTimeNs,
-      boolean waitForAuth, long dt, long at) {
+      boolean waitForAuth, long dt, long at, ConnectionFactory fa) {
     super(sa, c, bufSize, rq, wq, iq, opQueueMaxBlockTimeNs, waitForAuth, dt,
-      at);
+      at, fa);
   }
 
   @Override

src/test/java/net/spy/memcached/MockMemcachedNode.java

@@ -211,4 +211,14 @@ public void completedRead() {
     // noop
   }
 
+  @Override
+  public MemcachedConnection getConnection() {
+    return null;
+  }
+
+  @Override
+  public void setConnection(MemcachedConnection connection) {
+
+  }
+
 }