+ make sure we reuse the connection only after request is fully sent

Author: oleksiys

.gitignore

@@ -18,3 +18,4 @@ test-output
 /META-INF/MANIFEST.MF
 work
 atlassian-ide-plugin.xml
+/nb-configuration.xml

src/main/java/com/ning/http/client/providers/grizzly/AhcEventFilter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2015 Sonatype, Inc. All rights reserved.
+ * Copyright (c) 2012-2016 Sonatype, Inc. All rights reserved.
  *
  * This program is licensed to you under the Apache License Version 2.0,
  * and you may not use this file except in compliance with the Apache License Version 2.0.
@@ -57,6 +57,7 @@
 
 import static com.ning.http.util.AsyncHttpProviderUtils.*;
 import static com.ning.http.util.MiscUtils.isNonEmpty;
+import org.glassfish.grizzly.EmptyCompletionHandler;
 /**
  * AHC {@link HttpClientFilter} implementation.
  * 
@@ -442,22 +443,25 @@ private static boolean isRedirectAllowed(final HttpTransactionContext ctx) {
     }
 
     private static void cleanup(final HttpContext httpContext) {
-        final HttpTransactionContext context =
-                HttpTransactionContext.cleanupTransaction(httpContext);
-        
-        if (!context.isReuseConnection()) {
-            final Connection c = (Connection) httpContext.getCloseable();
-            final ConnectionManager cm = context.provider.getConnectionManager();
-            if (!httpContext.getRequest().getProcessingState().isStayAlive()) {
-                if (notKeepAliveReason == null) {
-                    notKeepAliveReason =
-                            new IOException("HTTP keep-alive was disabled for this connection");
+        HttpTransactionContext.cleanupTransaction(httpContext,
+                new EmptyCompletionHandler<HttpTransactionContext>() {
+            @Override
+            public void completed(HttpTransactionContext context) {
+                if (!context.isReuseConnection()) {
+                    final Connection c = (Connection) httpContext.getCloseable();
+                    if (!httpContext.getRequest().getProcessingState().isStayAlive()) {
+                        if (notKeepAliveReason == null) {
+                            notKeepAliveReason
+                                    = new IOException("HTTP keep-alive was disabled for this connection");
+                        }
+                        c.closeWithReason(notKeepAliveReason);
+                    } else {
+                        final ConnectionManager cm = context.provider.getConnectionManager();
+                        cm.returnConnection(c);
+                    }
                 }
-                c.closeWithReason(notKeepAliveReason);
-            } else {
-                cm.returnConnection(c);
             }
-        }
+        });
     }
 
     private static boolean redirectCountExceeded(final HttpTransactionContext context) {

src/main/java/com/ning/http/client/providers/grizzly/AsyncHttpClientFilter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2015 Sonatype, Inc. All rights reserved.
+ * Copyright (c) 2012-2016 Sonatype, Inc. All rights reserved.
  *
  * This program is licensed to you under the Apache License Version 2.0,
  * and you may not use this file except in compliance with the Apache License Version 2.0.
@@ -151,7 +151,10 @@ private boolean sendAsGrizzlyRequest(final HttpTransactionContext httpTxCtx,
         }
 
         HttpRequestPacket requestPacket;
-        final PayloadGenerator payloadGenerator = isPayloadAllowed(method) ? PayloadGenFactory.getPayloadGenerator(ahcRequest) : null;
+        final PayloadGenerator payloadGenerator = isPayloadAllowed(method)
+                ? PayloadGenFactory.getPayloadGenerator(ahcRequest)
+                : null;
+        
         if (payloadGenerator != null) {
             final long contentLength = ahcRequest.getContentLength();
             if (contentLength >= 0) {
@@ -196,7 +199,13 @@ private boolean sendAsGrizzlyRequest(final HttpTransactionContext httpTxCtx,
         ctx.notifyDownstream(new SSLSwitchingEvent(connection, secure,
                 uri.getHost(), uri.getPort()));
 
-        return sendRequest(httpTxCtx, ctx, requestPacket, wrapWithExpectHandlerIfNeeded(payloadGenerator, requestPacket));
+       final boolean isFullySent = sendRequest(httpTxCtx, ctx, requestPacket,
+                wrapWithExpectHandlerIfNeeded(payloadGenerator, requestPacket));
+       if (isFullySent) {
+           httpTxCtx.onRequestFullySent();
+       }
+       
+       return isFullySent;
     }
 
     private boolean establishConnectTunnel(final ProxyServer proxy,
@@ -284,7 +293,9 @@ private PayloadGenerator wrapWithExpectHandlerIfNeeded(final PayloadGenerator pa
         // check if we need to wrap the PayloadGenerator with ExpectWrapper
         final MimeHeaders headers = requestPacket.getHeaders();
         final int expectHeaderIdx = headers.indexOf(Header.Expect, 0);
-        return expectHeaderIdx != -1 && headers.getValue(expectHeaderIdx).equalsIgnoreCase("100-Continue") ? PayloadGenFactory.wrapWithExpect(payloadGenerator) : payloadGenerator;
+        return expectHeaderIdx != -1 && headers.getValue(expectHeaderIdx).equalsIgnoreCase("100-Continue")
+                ? PayloadGenFactory.wrapWithExpect(payloadGenerator)
+                : payloadGenerator;
     }
 
     private boolean isPayloadAllowed(final Method method) {

src/main/java/com/ning/http/client/providers/grizzly/FeedableBodyGenerator.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2015 Sonatype, Inc. All rights reserved.
+ * Copyright (c) 2012-2016 Sonatype, Inc. All rights reserved.
  *
  * This program is licensed to you under the Apache License Version 2.0,
  * and you may not use this file except in compliance with the Apache License Version 2.0.
@@ -306,10 +306,10 @@ public interface Feeder {
      * an implementation for the contract defined by the {@link #feed} method.
      */
     public static abstract class BaseFeeder implements Feeder {
-
+        
         protected final FeedableBodyGenerator feedableBodyGenerator;
-
-
+        
+        private boolean wasLastSent;
         // -------------------------------------------------------- Constructors
 
 
@@ -329,15 +329,35 @@ public final synchronized void feed(final Buffer buffer, final boolean last)
                 throw new IllegalArgumentException(
                         "Buffer argument cannot be null.");
             }
+            
             if (!feedableBodyGenerator.asyncTransferInitiated) {
                 throw new IllegalStateException("Asynchronous transfer has not been initiated.");
             }
+            
+            if (wasLastSent) {
+                if (buffer.hasRemaining()) {
+                    throw new IOException("Last chunk was alredy written");
+                }
+                
+                return;
+            }
+            
             blockUntilQueueFree(feedableBodyGenerator.context.getConnection());
             final HttpContent content =
                     feedableBodyGenerator.contentBuilder.content(buffer).last(last).build();
             final CompletionHandler<WriteResult> handler =
                     ((last) ? new LastPacketCompletionHandler() : null);
             feedableBodyGenerator.context.write(content, handler);
+            
+            if (last) {
+                wasLastSent = true;
+                final HttpTransactionContext currentTransaction =
+                        HttpTransactionContext.currentTransaction(
+                                feedableBodyGenerator.requestPacket);
+                if (currentTransaction != null) {
+                    currentTransaction.onRequestFullySent();
+                }
+            }
         }
 
         /**

src/main/java/com/ning/http/client/providers/grizzly/HttpTransactionContext.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2015 Sonatype, Inc. All rights reserved.
+ * Copyright (c) 2012-2016 Sonatype, Inc. All rights reserved.
  *
  * This program is licensed to you under the Apache License Version 2.0,
  * and you may not use this file except in compliance with the Apache License Version 2.0.
@@ -21,9 +21,12 @@
 import com.ning.http.util.AsyncHttpProviderUtils;
 import com.ning.http.util.ProxyUtils;
 import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
 import org.glassfish.grizzly.CloseListener;
 import org.glassfish.grizzly.CloseType;
 import org.glassfish.grizzly.Closeable;
+import org.glassfish.grizzly.CompletionHandler;
 import org.glassfish.grizzly.Connection;
 import org.glassfish.grizzly.Grizzly;
 import org.glassfish.grizzly.attributes.Attribute;
@@ -80,6 +83,11 @@ public final class HttpTransactionContext {
     // the pool
     boolean isReuseConnection;
 
+    /**
+     * <tt>true</tt> if the request is fully sent, or <tt>false</tt>otherwise.
+     */
+    volatile boolean isRequestFullySent;
+    Set<CompletionHandler<HttpTransactionContext>> reqFullySentHandlers;
 
     private final CloseListener listener = new CloseListener<Closeable, CloseType>() {
         @Override
@@ -111,16 +119,30 @@ static void bind(final HttpContext httpCtx,
         REQUEST_STATE_ATTR.set(httpCtx, httpTxContext);
     }
 
-    static HttpTransactionContext cleanupTransaction(final HttpContext httpCtx) {
+    static void cleanupTransaction(final HttpContext httpCtx,
+            final CompletionHandler<HttpTransactionContext> completionHandler) {
         final HttpTransactionContext httpTxContext = currentTransaction(httpCtx);
-        if (httpTxContext != null) {
-            httpCtx.getCloseable().removeCloseListener(httpTxContext.listener);
-            REQUEST_STATE_ATTR.remove(httpCtx);
-        }
 
-        return httpTxContext;
+        assert httpTxContext != null;
+        
+        if (httpTxContext.isRequestFullySent) {
+            cleanupTransaction(httpCtx, httpTxContext);
+            completionHandler.completed(httpTxContext);
+        } else {
+            httpTxContext.addRequestSentCompletionHandler(completionHandler);
+            if (httpTxContext.isRequestFullySent &&
+                    httpTxContext.removeRequestSentCompletionHandler(completionHandler)) {
+                completionHandler.completed(httpTxContext);
+            }
+        }
     }
 
+    static void cleanupTransaction(final HttpContext httpCtx,
+            final HttpTransactionContext httpTxContext) {
+        httpCtx.getCloseable().removeCloseListener(httpTxContext.listener);
+        REQUEST_STATE_ATTR.remove(httpCtx);
+    }
+    
     static HttpTransactionContext currentTransaction(
             final HttpHeader httpHeader) {
         return currentTransaction(httpHeader.getProcessingState().getHttpContext());
@@ -171,7 +193,7 @@ Request getAhcRequest() {
 
     ProxyServer getProxyServer() {
         return proxyServer;
-    }    
+    }
 
     // ----------------------------------------------------- Private Methods
 
@@ -246,4 +268,45 @@ void touchConnection() {
     void closeConnection() {
         connection.closeSilently();
     }
+
+    private synchronized void addRequestSentCompletionHandler(
+            final CompletionHandler<HttpTransactionContext> completionHandler) {
+        if (reqFullySentHandlers == null) {
+            reqFullySentHandlers = new HashSet<>();
+        }
+        reqFullySentHandlers.add(completionHandler);
+    }
+
+    private synchronized boolean removeRequestSentCompletionHandler(
+            final CompletionHandler<HttpTransactionContext> completionHandler) {
+        return reqFullySentHandlers != null
+                ? reqFullySentHandlers.remove(completionHandler)
+                : false;
+    }
+    
+    boolean isRequestFullySent() {
+        return isRequestFullySent;
+    }
+    
+    @SuppressWarnings("unchecked")
+    void onRequestFullySent() {
+        this.isRequestFullySent = true;
+        
+        Object[] handlers = null;
+        synchronized (this) {
+            if (reqFullySentHandlers != null) {
+                handlers = reqFullySentHandlers.toArray();
+                reqFullySentHandlers = null;
+            }
+        }
+        
+        if (handlers != null) {
+            for (Object o : handlers) {
+                try {
+                    ((CompletionHandler<HttpTransactionContext>) o).completed(this);
+                } catch (Exception e) {
+                }
+            }
+        }
+    }
 } // END HttpTransactionContext