Add async call retry to resolve the transient ZK connection issue.

zookeeper-api/src/main/java/org/apache/helix/zookeeper/zkclient/ZkClient.java

@@ -27,7 +27,10 @@
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.helix.zookeeper.exception.ZkClientException;
 import org.apache.helix.zookeeper.zkclient.annotation.PreFetch;
+import org.apache.helix.zookeeper.zkclient.callback.ZkAsyncCallMonitorContext;
 import org.apache.helix.zookeeper.zkclient.callback.ZkAsyncCallbacks;
+import org.apache.helix.zookeeper.zkclient.callback.ZkAsyncRetryCallContext;
+import org.apache.helix.zookeeper.zkclient.callback.ZkAsyncRetryThread;
 import org.apache.helix.zookeeper.zkclient.exception.ZkBadVersionException;
 import org.apache.helix.zookeeper.zkclient.exception.ZkException;
 import org.apache.helix.zookeeper.zkclient.exception.ZkInterruptedException;
@@ -96,6 +99,10 @@ public class ZkClient implements Watcher {
   private PathBasedZkSerializer _pathBasedZkSerializer;
   private ZkClientMonitor _monitor;
 
+  // To automatically retry the async operation, we need a separate thread other than the
+  // ZkEventThread. Otherwise the retry request might block the normal event processing.
+  protected final ZkAsyncRetryThread _asyncCallRetryThread;
+
   private class IZkDataListenerEntry {
     final IZkDataListener _dataListener;
     final boolean _prefetchData;
@@ -183,6 +190,9 @@ protected ZkClient(IZkConnection zkConnection, int connectionTimeout, long opera
     _operationRetryTimeoutInMillis = operationRetryTimeout;
     _isNewSessionEventFired = false;
 
+    _asyncCallRetryThread = new ZkAsyncRetryThread(zkConnection.getServers());
+    _asyncCallRetryThread.start();
+
     connect(connectionTimeout, this);
 
     // initiate monitor
@@ -1736,15 +1746,23 @@ public void asyncCreate(final String path, Object datat, final CreateMode mode,
       data = (datat == null ? null : serialize(datat, path));
     } catch (ZkMarshallingError e) {
       cb.processResult(KeeperException.Code.MARSHALLINGERROR.intValue(), path,
-          new ZkAsyncCallbacks.ZkAsyncCallContext(_monitor, startT, 0, false), null);
+          new ZkAsyncCallMonitorContext(_monitor, startT, 0, false), null);
       return;
     }
+    doAsyncCreate(path, data, mode, startT, cb);
+  }
+
+  private void doAsyncCreate(final String path, final byte[] data, final CreateMode mode,
+      final long startT, final ZkAsyncCallbacks.CreateCallbackHandler cb) {
     retryUntilConnected(() -> {
       ((ZkConnection) getConnection()).getZookeeper()
-          .create(path, data, ZooDefs.Ids.OPEN_ACL_UNSAFE,
-              // Arrays.asList(DEFAULT_ACL),
-              mode, cb, new ZkAsyncCallbacks.ZkAsyncCallContext(_monitor, startT,
-                  data == null ? 0 : data.length, false));
+          .create(path, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, mode, cb,
+              new ZkAsyncRetryCallContext(_asyncCallRetryThread, cb, _monitor, startT, 0, false) {
+                @Override
+                protected void doRetry() {
+                  doAsyncCreate(path, data, mode, System.currentTimeMillis(), cb);
+                }
+              });
       return null;
     });
   }
@@ -1758,50 +1776,66 @@ public void asyncSetData(final String path, Object datat, final int version,
       data = serialize(datat, path);
     } catch (ZkMarshallingError e) {
       cb.processResult(KeeperException.Code.MARSHALLINGERROR.intValue(), path,
-          new ZkAsyncCallbacks.ZkAsyncCallContext(_monitor, startT, 0, false), null);
+          new ZkAsyncCallMonitorContext(_monitor, startT, 0, false), null);
       return;
     }
+    doAsyncSetData(path, data, version, startT, cb);
+  }
+
+  private void doAsyncSetData(final String path, byte[] data, final int version, final long startT,
+      final ZkAsyncCallbacks.SetDataCallbackHandler cb) {
     retryUntilConnected(() -> {
       ((ZkConnection) getConnection()).getZookeeper().setData(path, data, version, cb,
-          new ZkAsyncCallbacks.ZkAsyncCallContext(_monitor, startT,
-              data == null ? 0 : data.length, false));
+          new ZkAsyncRetryCallContext(_asyncCallRetryThread, cb, _monitor, startT,
+              data == null ? 0 : data.length, false) {
+            @Override
+            protected void doRetry() {
+              doAsyncSetData(path, data, version, System.currentTimeMillis(), cb);
+            }
+          });
       return null;
     });
   }
 
   public void asyncGetData(final String path, final ZkAsyncCallbacks.GetDataCallbackHandler cb) {
     final long startT = System.currentTimeMillis();
-    retryUntilConnected(new Callable<Object>() {
-      @Override
-      public Object call() throws Exception {
-        ((ZkConnection) getConnection()).getZookeeper().getData(path, null, cb,
-            new ZkAsyncCallbacks.ZkAsyncCallContext(_monitor, startT, 0, true));
-        return null;
-      }
+    retryUntilConnected(() -> {
+      ((ZkConnection) getConnection()).getZookeeper().getData(path, null, cb,
+          new ZkAsyncRetryCallContext(_asyncCallRetryThread, cb, _monitor, startT, 0, true) {
+            @Override
+            protected void doRetry() {
+              asyncGetData(path, cb);
+            }
+          });
+      return null;
     });
   }
 
   public void asyncExists(final String path, final ZkAsyncCallbacks.ExistsCallbackHandler cb) {
     final long startT = System.currentTimeMillis();
-    retryUntilConnected(new Callable<Object>() {
-      @Override
-      public Object call() throws Exception {
-        ((ZkConnection) getConnection()).getZookeeper().exists(path, null, cb,
-            new ZkAsyncCallbacks.ZkAsyncCallContext(_monitor, startT, 0, true));
-        return null;
-      }
+    retryUntilConnected(() -> {
+      ((ZkConnection) getConnection()).getZookeeper().exists(path, null, cb,
+          new ZkAsyncRetryCallContext(_asyncCallRetryThread, cb, _monitor, startT, 0, true) {
+            @Override
+            protected void doRetry() {
+              asyncExists(path, cb);
+            }
+          });
+      return null;
     });
   }
 
   public void asyncDelete(final String path, final ZkAsyncCallbacks.DeleteCallbackHandler cb) {
     final long startT = System.currentTimeMillis();
-    retryUntilConnected(new Callable<Object>() {
-      @Override
-      public Object call() throws Exception {
-        ((ZkConnection) getConnection()).getZookeeper().delete(path, -1, cb,
-            new ZkAsyncCallbacks.ZkAsyncCallContext(_monitor, startT, 0, false));
-        return null;
-      }
+    retryUntilConnected(() -> {
+      ((ZkConnection) getConnection()).getZookeeper().delete(path, -1, cb,
+          new ZkAsyncRetryCallContext(_asyncCallRetryThread, cb, _monitor, startT, 0, false) {
+            @Override
+            protected void doRetry() {
+              asyncDelete(path, cb);
+            }
+          });
+      return null;
     });
   }
 
@@ -1955,6 +1989,10 @@ public void close() throws ZkInterruptedException {
         return;
       }
       setShutdownTrigger(true);
+      if (_asyncCallRetryThread != null) {
+        _asyncCallRetryThread.interrupt();
+        _asyncCallRetryThread.join(2000);
+      }
       _eventThread.interrupt();
       _eventThread.join(2000);
       if (isManagingZkConnection()) {

zookeeper-api/src/main/java/org/apache/helix/zookeeper/zkclient/callback/CancellableZkAsyncCallback.java

@@ -0,0 +1,8 @@
+package org.apache.helix.zookeeper.zkclient.callback;
+
+public interface CancellableZkAsyncCallback {
+  /**
+   * Notify all the callers that are waiting for the callback to cancel the wait.
+   */
+  void notifyCallers();
+}

zookeeper-api/src/main/java/org/apache/helix/zookeeper/zkclient/callback/ZkAsyncCallMonitorContext.java

@@ -0,0 +1,46 @@
+package org.apache.helix.zookeeper.zkclient.callback;
+
+import org.apache.helix.zookeeper.zkclient.metric.ZkClientMonitor;
+
+public class ZkAsyncCallMonitorContext {
+  private final long _startTimeMilliSec;
+  private final ZkClientMonitor _monitor;
+  private final boolean _isRead;
+  private int _bytes;
+
+  /**
+   * @param monitor           ZkClient monitor for update the operation result.
+   * @param startTimeMilliSec Operation initialization time.
+   * @param bytes             The data size in bytes that is involved in the operation.
+   * @param isRead            True if the operation is readonly.
+   */
+  public ZkAsyncCallMonitorContext(final ZkClientMonitor monitor, long startTimeMilliSec, int bytes,
+      boolean isRead) {
+    _monitor = monitor;
+    _startTimeMilliSec = startTimeMilliSec;
+    _bytes = bytes;
+    _isRead = isRead;
+  }
+
+  /**
+   * Update the operated data size in bytes.
+   * @param bytes
+   */
+  void setBytes(int bytes) {
+    _bytes = bytes;
+  }
+
+  /**
+   * Record the operation result into the specified ZkClient monitor.
+   * @param path
+   */
+  void recordAccess(String path) {
+    if (_monitor != null) {
+      if (_isRead) {
+        _monitor.record(path, _bytes, _startTimeMilliSec, ZkClientMonitor.AccessType.READ);
+      } else {
+        _monitor.record(path, _bytes, _startTimeMilliSec, ZkClientMonitor.AccessType.WRITE);
+      }
+    }
+  }
+}