Improve error handling in ThriftClientHandler

If transport fails, and there are no pending requests new requests
could potentially hang forever.

drift-transport-netty/src/main/java/io/airlift/drift/transport/netty/ThriftClientHandler.java

@@ -31,9 +31,9 @@
 
 import javax.annotation.concurrent.ThreadSafe;
 
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -95,27 +95,36 @@ private void sendMessage(ChannelHandlerContext context, ThriftRequest thriftRequ
             }
         }
 
+        // if this connection is failed, immediately fail the request
+        TException channelError = this.channelError.get();
+        if (channelError != null) {
+            thriftRequest.failed(channelError);
+            requestBuffer.release();
+            return;
+        }
+
         try {
             ChannelFuture sendFuture = context.write(requestBuffer, promise);
             sendFuture.addListener(future -> messageSent(context, sendFuture, requestHandler));
         }
         catch (Throwable t) {
-            onError(context, t);
+            onError(context, t, Optional.of(requestHandler));
+            requestBuffer.release();
         }
     }
 
     private void messageSent(ChannelHandlerContext context, ChannelFuture future, RequestHandler requestHandler)
     {
         try {
             if (!future.isSuccess()) {
-                onError(context, new TTransportException("Sending request failed", future.cause()));
+                onError(context, new TTransportException("Sending request failed", future.cause()), Optional.of(requestHandler));
                 return;
             }
 
             requestHandler.onRequestSent();
         }
         catch (Throwable t) {
-            onError(context, t);
+            onError(context, t, Optional.of(requestHandler));
         }
     }
 
@@ -135,21 +144,22 @@ public void channelRead(ChannelHandlerContext context, Object message)
 
     private void messageReceived(ChannelHandlerContext context, ByteBuf response)
     {
+        RequestHandler requestHandler = null;
         try {
             OptionalInt sequenceId = messageEncoding.extractResponseSequenceId(response.retainedDuplicate());
             if (!sequenceId.isPresent()) {
                 throw new TTransportException("Could not find sequenceId in Thrift message");
             }
 
-            RequestHandler requestHandler = pendingRequests.remove(sequenceId.getAsInt());
+            requestHandler = pendingRequests.remove(sequenceId.getAsInt());
             if (requestHandler == null) {
                 throw new TTransportException("Unknown sequence id in response: " + sequenceId.getAsInt());
             }
 
             requestHandler.onResponseReceived(response.retainedDuplicate());
         }
         catch (Throwable t) {
-            onError(context, t);
+            onError(context, t, Optional.ofNullable(requestHandler));
         }
         finally {
             response.release();
@@ -159,19 +169,17 @@ private void messageReceived(ChannelHandlerContext context, ByteBuf response)
     @Override
     public void exceptionCaught(ChannelHandlerContext context, Throwable cause)
     {
-        onError(context, cause);
+        onError(context, cause, Optional.empty());
     }
 
     @Override
     public void channelInactive(ChannelHandlerContext context)
             throws Exception
     {
-        if (!pendingRequests.isEmpty()) {
-            onError(context, new TTransportException("Client was disconnected by server"));
-        }
+        onError(context, new TTransportException("Client was disconnected by server"), Optional.empty());
     }
 
-    private void onError(ChannelHandlerContext context, Throwable throwable)
+    private void onError(ChannelHandlerContext context, Throwable throwable, Optional<RequestHandler> currentRequest)
     {
         TException thriftException;
         if (throwable instanceof TException) {
@@ -187,13 +195,20 @@ private void onError(ChannelHandlerContext context, Throwable throwable)
             return;
         }
 
+        // current request may have already been removed from pendingRequests, so notify it directly
+        currentRequest.ifPresent(request -> {
+            pendingRequests.remove(request.getSequenceId());
+            request.onChannelError(thriftException);
+        });
+
         // notify all pending requests of the error
+        // Note while loop should not be necessary since this class should be single
+        // threaded, but it is better to be safe in cleanup code
         while (!pendingRequests.isEmpty()) {
-            for (Iterator<RequestHandler> iterator = pendingRequests.values().iterator(); iterator.hasNext(); ) {
-                RequestHandler requestHandler = iterator.next();
-                iterator.remove();
-                requestHandler.onChannelError(thriftException);
-            }
+            pendingRequests.values().removeIf(request -> {
+                request.onChannelError(thriftException);
+                return true;
+            });
         }
 
         context.close();
@@ -258,6 +273,11 @@ public RequestHandler(ThriftRequest thriftRequest, int sequenceId)
             this.sequenceId = sequenceId;
         }
 
+        public int getSequenceId()
+        {
+            return sequenceId;
+        }
+
         void registerRequestTimeout(EventExecutor executor)
         {
             try {