Incremental fix for https://issues.sonatype.org/browse/AHC-78

("Multi-Part Zero Copy in Memory with Large Files Causes Out of Memory Error")

Collaborative work with Gail Hernandez

Author: jfarcand

src/main/java/com/ning/http/client/providers/netty/NettyAsyncHttpProvider.java

@@ -44,6 +44,7 @@
 import com.ning.http.client.ntlm.NTLMEngine;
 import com.ning.http.client.ntlm.NTLMEngineException;
 import com.ning.http.client.providers.netty.spnego.SpnegoEngine;
+import com.ning.http.multipart.MultipartBody;
 import com.ning.http.multipart.MultipartRequestEntity;
 import com.ning.http.util.AsyncHttpProviderUtils;
 import com.ning.http.util.AuthenticatorUtils;
@@ -401,6 +402,7 @@ protected final <T> void writeRequest(final Channel channel,
 
             if (future.getAndSetWriteBody(true)) {
                 if (!future.getNettyRequest().getMethod().equals(HttpMethod.CONNECT)) {
+
                     if (future.getRequest().getFile() != null) {
                         final File file = future.getRequest().getFile();
                         long fileLength = 0;
@@ -412,12 +414,11 @@ protected final <T> void writeRequest(final Channel channel,
                             ChannelFuture writeFuture;
                             if (channel.getPipeline().get(SslHandler.class) != null) {
                                 writeFuture = channel.write(new ChunkedFile(raf, 0, fileLength, 8192));
-                                writeFuture.addListener(new ProgressListener(false, future.getAsyncHandler(), future));
                             } else {
                                 final FileRegion region = new OptimizedFileRegion(raf, 0, fileLength);
                                 writeFuture = channel.write(region);
-                                writeFuture.addListener(new ProgressListener(false, future.getAsyncHandler(), future));
                             }
+                            writeFuture.addListener(new ProgressListener(false, future.getAsyncHandler(), future));
                         } catch (IOException ex) {
                             if (raf != null) {
                                 try {
@@ -427,12 +428,23 @@ protected final <T> void writeRequest(final Channel channel,
                             }
                             throw ex;
                         }
-                    } else if (body != null) {
+                    } else if (body != null || future.getRequest().getParts() != null) {
+                        /**
+                         * TODO: AHC-78: SSL + zero copy isn't supported by the MultiPart class and pretty complex to implements.
+                         */
+                        if (future.getRequest().getParts() != null) {
+                            String boundary = future.getNettyRequest().getHeader("Content-Type");
+                            String length = future.getNettyRequest().getHeader("Content-Length");
+                            body = new MultipartBody(future.getRequest().getParts(), boundary, length);
+                        }
+
                         ChannelFuture writeFuture;
                         if (channel.getPipeline().get(SslHandler.class) == null && (body instanceof RandomAccessBody)) {
-                            writeFuture = channel.write(new BodyFileRegion((RandomAccessBody) body));
+                            BodyFileRegion bodyFileRegion = new BodyFileRegion((RandomAccessBody) body);
+                            writeFuture = channel.write(bodyFileRegion);
                         } else {
-                            writeFuture = channel.write(new BodyChunkedInput(body));
+                            BodyChunkedInput bodyChunkedInput = new BodyChunkedInput(body);
+                            writeFuture = channel.write(bodyChunkedInput);
                         }
 
                         final Body b = body;
@@ -694,10 +706,14 @@ private static HttpRequest construct(AsyncHttpClientConfig config,
                     nettyRequest.setHeader(HttpHeaders.Names.CONTENT_TYPE, mre.getContentType());
                     nettyRequest.setHeader(HttpHeaders.Names.CONTENT_LENGTH, String.valueOf(mre.getContentLength()));
 
-                    ChannelBuffer b = ChannelBuffers.dynamicBuffer(lenght);
-                    mre.writeRequest(new ChannelBufferOutputStream(b));
-                    nettyRequest.setContent(b);
-
+                    /**
+                     * TODO: AHC-78: SSL + zero copy isn't supported by the MultiPart class and pretty complex to implements.
+                     */
+                    if (uri.toString().startsWith("https")) {
+                        ChannelBuffer b = ChannelBuffers.dynamicBuffer(lenght);
+                        mre.writeRequest(new ChannelBufferOutputStream(b));
+                        nettyRequest.setContent(b);
+                    }
                 } else if (request.getEntityWriter() != null) {
                     int lenght = computeAndSetContentLength(request, nettyRequest);
 
@@ -1068,7 +1084,7 @@ public void messageReceived(final ChannelHandlerContext ctx, MessageEvent e) thr
                 if (realm != null && !future.getURI().getPath().equalsIgnoreCase(realm.getUri())) {
                     builder.setUrl(future.getURI().toString());
                 }
-                
+
                 if (statusCode == 401
                         && wwwAuth.size() > 0
                         && !future.getAndSetAuth(true)) {
@@ -1282,7 +1298,7 @@ private Realm kerberosChallenge(List<String> proxyAuth,
                                     FluentCaseInsensitiveStringsMap headers,
                                     Realm realm,
                                     NettyResponseFuture<?> future) throws NTLMEngineException {
-        
+
         URI uri = URI.create(request.getUrl());
         String host = request.getVirtualHost() == null ? uri.getHost() : request.getVirtualHost();
         String server = proxyServer == null ? host : proxyServer.getHost();

src/main/java/com/ning/http/multipart/MultipartBody.java

@@ -47,8 +47,6 @@ public class MultipartBody implements RandomAccessBody {
 
     enum FileLocation {NONE, START, MIDDLE, END}
 
-    ;
-
     public MultipartBody(List<com.ning.http.client.Part> parts, String boundary, String contentLength) {
         this.boundary = MultipartEncodingUtil.getAsciiBytes(boundary.substring("multipart/form-data; boundary=".length()));
         this.contentLength = Long.parseLong(contentLength);
@@ -194,7 +192,7 @@ public long read(ByteBuffer buffer) throws IOException {
             return overallLength;
 
         } catch (Exception e) {
-            e.printStackTrace();
+            logger.info("read exception", e);
             return 0;
         }
     }
@@ -312,7 +310,7 @@ public long transferTo(long position, long count, WritableByteChannel target)
         long overallLength = 0;
 
         if (startPart == parts.size()) {
-            return overallLength;
+            return contentLength;
         }
 
         int tempPart = startPart;
@@ -372,7 +370,7 @@ private FilePart generateClientFilePart(com.ning.http.client.Part part)
             throws FileNotFoundException {
         com.ning.http.client.FilePart currentPart = (com.ning.http.client.FilePart) part;
 
-        FilePart filePart = new FilePart(currentPart.getName(), currentPart.getFile());
+        FilePart filePart = new FilePart(currentPart.getName(), currentPart.getFile(), currentPart.getMimeType(), currentPart.getCharSet());
         return filePart;
     }
 
@@ -450,11 +448,22 @@ private long handleFilePart(WritableByteChannel target, FilePart filePart) throw
 
             FileChannel fc = raf.getChannel();
 
-            long fileLength = fc.transferTo(0, file.length(), target);
-
-            if (fileLength != file.length()) {
-                logger.info("Did not complete file.");
+            long l = file.length();
+            int fileLength = 0;
+            synchronized (fc) {
+                while (fileLength != l) {
+                    fileLength += fc.transferTo(fileLength, l, target);
+                    if (fileLength != l) {
+                        logger.info("Waiting for writing...");
+                        try {
+                            fc.wait(1000);
+                        } catch (InterruptedException e) {
+                            logger.trace(e.getMessage(), e);
+                        }
+                    }
+                }
             }
+            fc.close();
 
             length += handleFileEnd(target, filePart);
 
@@ -474,15 +483,20 @@ private long handlePartSource(WritableByteChannel target, FilePart filePart) thr
 
         InputStream stream = partSource.createInputStream();
 
-        int nRead = 0;
-        byte[] bytes = new byte[(int) partSource.getLength()];
-        while (nRead != -1) {
-            nRead = stream.read(bytes);
-            if (nRead > 0) {
-                ByteArrayOutputStream bos = new ByteArrayOutputStream(nRead);
-                bos.write(bytes, 0, nRead);
-                writeToTarget(target, bos);
+        try {
+            int nRead = 0;
+            while (nRead != -1) {
+                // Do not buffer the entire monster in memory.
+                byte[] bytes = new byte[8192];
+                nRead = stream.read(bytes);
+                if (nRead > 0) {
+                    ByteArrayOutputStream bos = new ByteArrayOutputStream(nRead);
+                    bos.write(bytes, 0, nRead);
+                    writeToTarget(target, bos);
+                }
             }
+        } finally {
+            stream.close();
         }
         length += handleFileEnd(target, filePart);
 
@@ -518,19 +532,24 @@ private long writeToTarget(WritableByteChannel target, ByteArrayOutputStream byt
             throws IOException {
 
         int written = 0;
+        int maxSpin = 0;
         synchronized (byteWriter) {
+            ByteBuffer message = ByteBuffer.wrap(byteWriter.toByteArray());
             while ((target.isOpen()) && (written < byteWriter.size())) {
-                ByteBuffer message = ByteBuffer.wrap(byteWriter.toByteArray());
-                written = target.write(message);
-                // TODO: This is dangerous to spin, we need to find another way to wait until 
-                //the byte channel is ready to receive the additional bytes or else data is lost.
-                if (written != byteWriter.size()) {
+                written += target.write(message);
+                if (written != byteWriter.size() && maxSpin++ < 10) {
                     logger.info("Waiting for writing...");
                     try {
                         byteWriter.wait(1000);
+                        maxSpin++;
                     } catch (InterruptedException e) {
                         logger.trace(e.getMessage(), e);
                     }
+                } else {
+                    if (maxSpin >= 10) {
+                        throw new IOException("Unable to write on channel " + target);
+                    }
+                    maxSpin = 0;
                 }
             }
         }

src/test/java/com/ning/http/client/async/AbstractBasicTest.java

@@ -19,7 +19,6 @@
 import com.ning.http.client.AsyncHandler;
 import com.ning.http.client.AsyncHttpClient;
 import com.ning.http.client.AsyncHttpClientConfig;
-import com.ning.http.client.AsyncHttpProvider;
 import com.ning.http.client.HttpResponseBodyPart;
 import com.ning.http.client.HttpResponseHeaders;
 import com.ning.http.client.HttpResponseStatus;
@@ -71,7 +70,7 @@ public void handle(String pathInContext,
             if (request.getMethod().equalsIgnoreCase("OPTIONS")) {
                 httpResponse.addHeader("Allow","GET,HEAD,POST,OPTIONS,TRACE");
             };
-            
+
             Enumeration<?> e = httpRequest.getHeaderNames();
             String param;
             while (e.hasMoreElements()) {

src/test/java/com/ning/http/client/async/FilePartLargeFileTest.java

@@ -0,0 +1,160 @@
+/*
+ * Copyright (c) 2010-2011 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.
+ * You may obtain a copy of the Apache License Version 2.0 at http://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the Apache License Version 2.0 is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the Apache License Version 2.0 for the specific language governing permissions and limitations there under.
+ */
+package com.ning.http.client.async;
+
+import com.ning.http.client.AsyncHttpClient;
+import com.ning.http.client.AsyncHttpClient.BoundRequestBuilder;
+import com.ning.http.client.AsyncHttpClientConfig;
+import com.ning.http.client.FilePart;
+import com.ning.http.client.Response;
+import org.eclipse.jetty.server.Request;
+import org.eclipse.jetty.server.handler.AbstractHandler;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.Test;
+
+import javax.servlet.ServletException;
+import javax.servlet.ServletInputStream;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.net.URL;
+import java.util.UUID;
+
+import static org.testng.FileAssert.fail;
+
+public abstract class FilePartLargeFileTest
+        extends AbstractBasicTest {
+
+    private File largeFile;
+
+    @Test(groups = {"standalone", "default_provider"}, enabled = true)
+    public void testPutImageFile()
+            throws Exception {
+        largeFile = getTestFile();
+        AsyncHttpClientConfig config = new AsyncHttpClientConfig.Builder().setRequestTimeoutInMs(100 * 6000).build();
+        AsyncHttpClient client = getAsyncHttpClient(config);
+        BoundRequestBuilder rb = client.preparePut(getTargetUrl());
+
+        rb.addBodyPart(new FilePart("test", largeFile, "application/octet-stream" , "UTF-8"));
+
+        Response response = rb.execute().get();
+        Assert.assertEquals(200, response.getStatusCode());
+
+        client.close();
+    }
+
+    @Test(groups = {"standalone", "default_provider"}, enabled = true)
+    public void testPutLargeTextFile()
+            throws Exception {
+        byte[] bytes = "RatherLargeFileRatherLargeFileRatherLargeFileRatherLargeFile".getBytes("UTF-16");
+        long repeats = (1024 * 1024 / bytes.length) + 1;
+        largeFile = createTempFile(bytes, (int) repeats);
+
+        AsyncHttpClientConfig config = new AsyncHttpClientConfig.Builder().build();
+        AsyncHttpClient client = getAsyncHttpClient(config);
+        BoundRequestBuilder rb = client.preparePut(getTargetUrl());
+
+        rb.setBody(largeFile);
+
+        Response response = rb.execute().get();
+        Assert.assertEquals(200, response.getStatusCode());
+        client.close();
+    }
+
+    private static File getTestFile() {
+        String testResource1 = "300k.png";
+
+        File testResource1File = null;
+        try {
+            ClassLoader cl = ChunkingTest.class.getClassLoader();
+            URL url = cl.getResource(testResource1);
+            testResource1File = new File(url.toURI());
+        } catch (Throwable e) {
+            // TODO Auto-generated catch block
+            fail("unable to find " + testResource1);
+        }
+
+        return testResource1File;
+    }
+
+    @AfterMethod
+    public void after() {
+        largeFile.delete();
+    }
+
+    @Override
+    public AbstractHandler configureHandler()
+            throws Exception {
+        return new AbstractHandler() {
+
+            public void handle(String arg0, Request arg1, HttpServletRequest req, HttpServletResponse resp)
+                    throws IOException, ServletException {
+
+                ServletInputStream in = req.getInputStream();
+                byte[] b = new byte[8192];
+
+                int count = -1;
+                int total = 0;
+                while ((count = in.read(b)) != -1) {
+                    b = new byte[8192];
+                    total += count;
+                }
+                System.err.println("consumed " + total + " bytes.");
+
+                resp.setStatus(200);
+                resp.addHeader("X-TRANFERED", String.valueOf(total));
+                resp.getOutputStream().flush();
+                resp.getOutputStream().close();
+
+                arg1.setHandled(true);
+
+            }
+        };
+    }
+
+    private static final File TMP = new File(System.getProperty("java.io.tmpdir"), "ahc-tests-"
+            + UUID.randomUUID().toString().substring(0, 8));
+
+    public static File createTempFile(byte[] pattern, int repeat)
+            throws IOException {
+        TMP.mkdirs();
+        TMP.deleteOnExit();
+        File tmpFile = File.createTempFile("tmpfile-", ".data", TMP);
+        tmpFile.deleteOnExit();
+        write(pattern, repeat, tmpFile);
+
+        return tmpFile;
+    }
+
+    public static void write(byte[] pattern, int repeat, File file)
+            throws IOException {
+        file.deleteOnExit();
+        file.getParentFile().mkdirs();
+        FileOutputStream out = null;
+        try {
+            out = new FileOutputStream(file);
+            for (int i = 0; i < repeat; i++) {
+                out.write(pattern);
+            }
+        }
+        finally {
+            if (out != null) {
+                out.close();
+            }
+        }
+    }
+
+}