AF: check host before closing

Author: rashtao

src/main/java/com/arangodb/internal/net/DirtyReadHostHandler.java

@@ -60,6 +60,11 @@ public void fail(Exception exception) {
         determineHostHandler().fail(exception);
     }
 
+    @Override
+    public void failIfNotMatch(HostDescription host, Exception exception) {
+        determineHostHandler().failIfNotMatch(host, exception);
+    }
+
     @Override
     public void reset() {
         determineHostHandler().reset();
@@ -81,6 +86,11 @@ public void closeCurrentOnError() {
         determineHostHandler().closeCurrentOnError();
     }
 
+    @Override
+    public void closeCurrentOnErrorIfNotMatch(HostDescription host) {
+        determineHostHandler().closeCurrentOnErrorIfNotMatch(host);
+    }
+
     @Override
     public void setJwt(String jwt) {
         master.setJwt(jwt);

src/main/java/com/arangodb/internal/net/FallbackHostHandler.java

@@ -79,6 +79,13 @@ public void fail(Exception exception) {
         lastFailExceptions.add(exception);
     }
 
+    @Override
+    public synchronized void failIfNotMatch(HostDescription host, Exception exception) {
+        if (!host.equals(current.getDescription())) {
+            fail(exception);
+        }
+    }
+
     @Override
     public void reset() {
         iterations = 0;
@@ -104,6 +111,13 @@ public void closeCurrentOnError() {
         current.closeOnError();
     }
 
+    @Override
+    public synchronized void closeCurrentOnErrorIfNotMatch(HostDescription host) {
+        if (!host.equals(current.getDescription())) {
+            closeCurrentOnError();
+        }
+    }
+
     @Override
     public void setJwt(String jwt) {
         hosts.setJwt(jwt);

src/main/java/com/arangodb/internal/net/HostHandler.java

@@ -33,6 +33,8 @@ public interface HostHandler {
 
     void fail(Exception exception);
 
+    void failIfNotMatch(HostDescription host, Exception exception);
+
     void reset();
 
     void confirm();
@@ -41,6 +43,8 @@ public interface HostHandler {
 
     void closeCurrentOnError();
 
+    void closeCurrentOnErrorIfNotMatch(HostDescription host);
+
     void setJwt(String jwt);
 
 }

src/main/java/com/arangodb/internal/net/RandomHostHandler.java

@@ -59,6 +59,13 @@ public void fail(Exception exception) {
         current = fallback.get(null, null);
     }
 
+    @Override
+    public synchronized void failIfNotMatch(HostDescription host, Exception exception) {
+        if (!host.equals(current.getDescription())) {
+            fail(exception);
+        }
+    }
+
     private Host getRandomHost(final boolean initial, final boolean closeConnections) {
         hosts = resolver.resolve(initial, closeConnections);
         final ArrayList<Host> hostList = new ArrayList<>(hosts.getHostsList());
@@ -85,6 +92,13 @@ public void closeCurrentOnError() {
         current.closeOnError();
     }
 
+    @Override
+    public synchronized void closeCurrentOnErrorIfNotMatch(HostDescription host) {
+        if (!host.equals(current.getDescription())) {
+            closeCurrentOnError();
+        }
+    }
+
     @Override
     public void setJwt(String jwt) {
         fallback.setJwt(jwt);

src/main/java/com/arangodb/internal/net/RoundRobinHostHandler.java

@@ -89,6 +89,11 @@ public void fail(Exception exception) {
         lastFailExceptions.add(exception);
     }
 
+    @Override
+    public void failIfNotMatch(HostDescription host, Exception exception) {
+        fail(exception);
+    }
+
     @Override
     public void reset() {
         fails = 0;
@@ -109,6 +114,11 @@ public void closeCurrentOnError() {
         currentHost.closeOnError();
     }
 
+    @Override
+    public void closeCurrentOnErrorIfNotMatch(HostDescription host) {
+        closeCurrentOnError();
+    }
+
     @Override
     public void setJwt(String jwt) {
         hosts.setJwt(jwt);

src/main/java/com/arangodb/internal/velocystream/VstCommunicationSync.java

@@ -151,8 +151,8 @@ protected Response execute(final Request request, final VstConnectionSync connec
             }
             final String location = e.getLocation();
             final HostDescription redirectHost = HostUtils.createFromLocation(location);
-            hostHandler.closeCurrentOnError();
-            hostHandler.fail(e);
+            hostHandler.closeCurrentOnErrorIfNotMatch(redirectHost);
+            hostHandler.failIfNotMatch(redirectHost, e);
             return execute(request, new HostHandle().setHost(redirectHost), attemptCount + 1);
         }
     }