Delay VTGateConnection closure when rotating connections

java/client/src/main/java/io/vitess/client/VTGateConnection.java

@@ -252,4 +252,11 @@ public void close() throws IOException {
         client.close();
     }
 
+    @Override
+    public String toString() {
+        return String.format("[VTGateConnection-%s client=%s]",
+                Integer.toHexString(this.hashCode()),
+                client.toString()
+        );
+    }
 }

java/grpc-client/src/main/java/io/vitess/client/grpc/GrpcClient.java

@@ -21,6 +21,7 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import io.grpc.ManagedChannel;
 import io.grpc.StatusRuntimeException;
+import io.grpc.internal.WithLogId;
 import io.vitess.client.Context;
 import io.vitess.client.Proto;
 import io.vitess.client.RpcClient;
@@ -79,11 +80,17 @@
  */
 public class GrpcClient implements RpcClient {
   private final ManagedChannel channel;
+  private final String channelId;
   private final VitessStub asyncStub;
   private final VitessFutureStub futureStub;
 
   public GrpcClient(ManagedChannel channel) {
     this.channel = channel;
+    if (channel instanceof WithLogId) {
+      channelId = ((WithLogId) channel).getLogId().toString();
+    } else {
+      channelId = channel.toString();
+    }
     asyncStub = VitessGrpc.newStub(channel);
     futureStub = VitessGrpc.newFutureStub(channel);
   }
@@ -296,4 +303,12 @@ private VitessFutureStub getFutureStub(Context ctx) {
     }
     return futureStub.withDeadlineAfter(timeout.getMillis(), TimeUnit.MILLISECONDS);
   }
+
+  @Override
+  public String toString() {
+      return String.format("[GrpcClient-%s channel=%s]",
+              Integer.toHexString(this.hashCode()),
+              channelId
+      );
+  }
 }

java/jdbc/src/main/java/io/vitess/jdbc/ConnectionProperties.java

@@ -31,7 +31,6 @@
 import io.vitess.proto.Topodata;
 import io.vitess.util.Constants;
 import io.vitess.util.StringUtils;
-import java.util.function.Function;
 
 public class ConnectionProperties {
 
@@ -194,6 +193,14 @@ public class ConnectionProperties {
         "refreshSeconds",
         "How often in seconds the driver will monitor for changes to the keystore and truststore files, when refreshConnection is enabled.",
         60);
+    private BooleanConnectionProperty refreshClosureDelayed = new BooleanConnectionProperty(
+            "refreshClosureDelayed",
+            "When enabled, the closing of the old connections will be delayed instead of happening instantly.",
+            false);
+    private LongConnectionProperty refreshClosureDelaySeconds = new LongConnectionProperty(
+            "refreshClosureDelaySeconds",
+            "How often in seconds the grace period before closing the old connections that had been recreated.",
+            300);
     private StringConnectionProperty keyStore = new StringConnectionProperty(
         Constants.Property.KEYSTORE,
         "The Java .JKS keystore file to use when TLS is enabled",
@@ -517,6 +524,22 @@ public void setRefreshSeconds(long refreshSeconds) {
         this.refreshSeconds.setValue(refreshSeconds);
     }
 
+    public boolean getRefreshClosureDelayed() {
+        return refreshClosureDelayed.getValueAsBoolean();
+    }
+
+    public void setRefreshClosureDelayed(boolean refreshClosureDelayed) {
+        this.refreshClosureDelayed.setValue(refreshClosureDelayed);
+    }
+
+    public long getRefreshClosureDelaySeconds() {
+        return refreshClosureDelaySeconds.getValueAsLong();
+    }
+
+    public void setRefreshClosureDelaySeconds(long refreshClosureDelaySeconds) {
+        this.refreshClosureDelaySeconds.setValue(refreshClosureDelaySeconds);
+    }
+
     public String getKeyStore() {
         return keyStore.getValueAsString();
     }

java/jdbc/src/main/java/io/vitess/jdbc/VitessVTGateManager.java

@@ -16,7 +16,6 @@
 
 package io.vitess.jdbc;
 
-import java.io.File;
 import java.io.IOException;
 import java.sql.SQLException;
 import java.util.ArrayList;
@@ -30,10 +29,7 @@
 import java.util.logging.Logger;
 import java.util.logging.Level;
 
-import com.google.common.io.Closeables;
-
 import io.vitess.client.Context;
-import io.vitess.client.RpcClient;
 import io.vitess.client.VTGateConnection;
 import io.vitess.client.RefreshableVTGateConnection;
 import io.vitess.client.grpc.GrpcClientFactory;
@@ -52,6 +48,8 @@ public class VitessVTGateManager {
     private static ConcurrentHashMap<String, VTGateConnection> vtGateConnHashMap =
         new ConcurrentHashMap<>();
     private static Timer vtgateConnRefreshTimer = null;
+    private static Timer vtgateClosureTimer = null;
+    private static long vtgateClosureDelaySeconds = 0L;
 
     /**
      * VTGateConnections object consist of vtGateIdentifire list and return vtGate object in round robin.
@@ -66,6 +64,7 @@ public static class VTGateConnections {
          * @param connection
          */
         public VTGateConnections(final VitessConnection connection) {
+            maybeStartClosureTimer(connection);
             for (final VitessJDBCUrl.HostInfo hostInfo : connection.getUrl().getHostInfos()) {
                 String identifier = getIdentifer(hostInfo.getHostname(), hostInfo.getPort(), connection.getUsername(), connection.getTarget());
                 synchronized (VitessVTGateManager.class) {
@@ -105,6 +104,17 @@ public VTGateConnection getVtGateConnInstance() {
 
     }
 
+    private static void maybeStartClosureTimer(VitessConnection connection) {
+        if (connection.getRefreshClosureDelayed() && vtgateClosureTimer == null) {
+            synchronized (VitessVTGateManager.class) {
+                if (vtgateClosureTimer == null) {
+                    vtgateClosureTimer = new Timer("vtgate-conn-closure", true);
+                    vtgateClosureDelaySeconds = connection.getRefreshClosureDelaySeconds();
+                }
+            }
+        }
+    }
+
     private static String getIdentifer(String hostname, int port, String userIdentifer, String keyspace) {
         return (hostname + port + userIdentifer + keyspace);
     }
@@ -130,11 +140,7 @@ private static void refreshUpdatedSSLConnections(VitessJDBCUrl.HostInfo hostInfo
                     if (existing.checkKeystoreUpdates()) {
                         updatedCount++;
                         VTGateConnection old = vtGateConnHashMap.replace(entry.getKey(), getVtGateConn(hostInfo, connection));
-                        try {
-                            old.close();
-                        } catch (IOException ioe) {
-                            logger.log(Level.WARNING, "Error closing VTGateConnection", ioe);
-                        }
+                        closeRefreshedConnection(old);
                     }
                 }
             }
@@ -144,6 +150,30 @@ private static void refreshUpdatedSSLConnections(VitessJDBCUrl.HostInfo hostInfo
         }
     }
 
+    private static void closeRefreshedConnection(final VTGateConnection old) {
+        if (vtgateClosureTimer != null) {
+            logger.info(String.format("%s Closing connection with a %s second delay", old, vtgateClosureDelaySeconds));
+            vtgateClosureTimer.schedule(new TimerTask() {
+                @Override
+                public void run() {
+                    actuallyCloseRefreshedConnection(old);
+                }
+            },
+            TimeUnit.SECONDS.toMillis(vtgateClosureDelaySeconds));
+        } else {
+            actuallyCloseRefreshedConnection(old);
+        }
+    }
+
+    private static void actuallyCloseRefreshedConnection(final VTGateConnection old) {
+        try {
+            logger.info(old + " Closing connection because it had been refreshed");
+            old.close();
+        } catch (IOException ioe) {
+            logger.log(Level.WARNING, String.format("Error closing VTGateConnection %s", old), ioe);
+        }
+    }
+
     /**
      * Create vtGateConn object with given identifier.
      *

java/jdbc/src/test/java/io/vitess/jdbc/ConnectionPropertiesTest.java

@@ -31,7 +31,7 @@
 
 public class ConnectionPropertiesTest {
 
-    private static final int NUM_PROPS = 37;
+    private static final int NUM_PROPS = 39;
 
     @Test
     public void testReflection() throws Exception {
@@ -154,8 +154,8 @@ public void testDriverPropertiesOutput() throws SQLException {
         Assert.assertEquals("grpcRetriesInitialBackoffMillis", infos[7].name);
         Assert.assertEquals("grpcRetriesMaxBackoffMillis", infos[8].name);
         Assert.assertEquals(Constants.Property.INCLUDED_FIELDS, infos[9].name);
-        Assert.assertEquals(Constants.Property.TABLET_TYPE, infos[19].name);
-        Assert.assertEquals(Constants.Property.TWOPC_ENABLED, infos[27].name);
+        Assert.assertEquals(Constants.Property.TABLET_TYPE, infos[21].name);
+        Assert.assertEquals(Constants.Property.TWOPC_ENABLED, infos[29].name);
     }
 
     @Test