Add write timeout for post/put requests (#485)

* Introduce API to set writeTimeout and use it in the NetHttp request

* Surface the writeTimeout to HttpRequest

* remove debug statements

* Fix warnings

* Fix syntax

* DI for testing

* Make field constant field final

Author: chingor13

google-http-client/src/main/java/com/google/api/client/http/HttpRequest.java

@@ -158,6 +158,11 @@ static String executeAndGetValueOfSomeCustomHeader(HttpRequest request) {
    */
   private int readTimeout = 20 * 1000;
 
+  /**
+   * Timeout in milliseconds to set POST/PUT data or {@code 0} for an infinite timeout.
+   */
+  private int writeTimeout = 0;
+
   /** HTTP unsuccessful (non-2XX) response handler or {@code null} for none. */
   private HttpUnsuccessfulResponseHandler unsuccessfulResponseHandler;
 
@@ -493,6 +498,30 @@ public HttpRequest setReadTimeout(int readTimeout) {
     return this;
   }
 
+  /**
+   * Returns the timeout in milliseconds to send POST/PUT data or {@code 0} for an infinite timeout.
+   *
+   * <p>
+   * By default it is 0 (infinite).
+   * </p>
+   *
+   * @since 1.26
+   */
+  public int getWriteTimeout() {
+    return writeTimeout;
+  }
+
+  /**
+   * Sets the timeout in milliseconds to send POST/PUT data or {@code 0} for an infinite timeout.
+   *
+   * @since 1.26
+   */
+  public HttpRequest setWriteTimeout(int writeTimeout) {
+    Preconditions.checkArgument(writeTimeout >= 0);
+    this.writeTimeout = writeTimeout;
+    return this;
+  }
+
   /**
    * Returns the HTTP request headers.
    *
@@ -977,6 +1006,7 @@ public HttpResponse execute() throws IOException {
 
       // execute
       lowLevelHttpRequest.setTimeout(connectTimeout, readTimeout);
+      lowLevelHttpRequest.setWriteTimeout(writeTimeout);
       try {
         LowLevelHttpResponse lowLevelHttpResponse = lowLevelHttpRequest.execute();
         // Flag used to indicate if an exception is thrown before the response is constructed.

google-http-client/src/main/java/com/google/api/client/http/LowLevelHttpRequest.java

@@ -159,6 +159,21 @@ public final StreamingContent getStreamingContent() {
   public void setTimeout(int connectTimeout, int readTimeout) throws IOException {
   }
 
+  /**
+   * Sets the write timeout for POST/PUT requests.
+   *
+   * <p>
+   * Default implementation does nothing, but subclasses should normally override.
+   * </p>
+   *
+   * @param writeTimeout timeout in milliseconds to establish a connection or {@code 0} for an
+   *        infinite timeout
+   * @throws IOException I/O exception
+   * @since 1.26
+   */
+  public void setWriteTimeout(int writeTimeout) throws IOException {
+  }
+
   /** Executes the request and returns a low-level HTTP response object. */
   public abstract LowLevelHttpResponse execute() throws IOException;
 }

google-http-client/src/main/java/com/google/api/client/http/javanet/NetHttpRequest.java

@@ -18,22 +18,34 @@
 import com.google.api.client.http.LowLevelHttpResponse;
 import com.google.api.client.util.Preconditions;
 
+import com.google.api.client.util.StreamingContent;
+import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.net.HttpURLConnection;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 /**
  * @author Yaniv Inbar
  */
 final class NetHttpRequest extends LowLevelHttpRequest {
 
   private final HttpURLConnection connection;
+  private int writeTimeout;
 
   /**
    * @param connection HTTP URL connection
    */
   NetHttpRequest(HttpURLConnection connection) {
     this.connection = connection;
+    this.writeTimeout = 0;
     connection.setInstanceFollowRedirects(false);
   }
 
@@ -48,8 +60,32 @@ public void setTimeout(int connectTimeout, int readTimeout) {
     connection.setConnectTimeout(connectTimeout);
   }
 
+  @Override
+  public void setWriteTimeout(int writeTimeout) throws IOException {
+    this.writeTimeout = writeTimeout;
+  }
+
+  interface OutputWriter {
+    void write(OutputStream outputStream, StreamingContent content) throws IOException;
+  }
+
+  static class DefaultOutputWriter implements OutputWriter {
+    @Override
+    public void write(OutputStream outputStream, final StreamingContent content)
+        throws IOException {
+      content.writeTo(outputStream);
+    }
+  }
+
+  private static final OutputWriter DEFAULT_CONNECTION_WRITER = new DefaultOutputWriter();
+
   @Override
   public LowLevelHttpResponse execute() throws IOException {
+    return execute(DEFAULT_CONNECTION_WRITER);
+  }
+
+  @VisibleForTesting
+  LowLevelHttpResponse execute(final OutputWriter outputWriter) throws IOException {
     HttpURLConnection connection = this.connection;
     // write content
     if (getStreamingContent() != null) {
@@ -74,10 +110,12 @@ public LowLevelHttpResponse execute() throws IOException {
         } else {
           connection.setChunkedStreamingMode(0);
         }
-        OutputStream out = connection.getOutputStream();
+        final OutputStream out = connection.getOutputStream();
+
         boolean threw = true;
         try {
-          getStreamingContent().writeTo(out);
+          writeContentToOutputStream(outputWriter, out);
+
           threw = false;
         } finally {
           try {
@@ -111,4 +149,38 @@ public LowLevelHttpResponse execute() throws IOException {
       }
     }
   }
+
+  private void writeContentToOutputStream(final OutputWriter outputWriter, final OutputStream out)
+      throws IOException {
+    if (writeTimeout == 0) {
+      outputWriter.write(out, getStreamingContent());
+    } else {
+      // do it with timeout
+      final StreamingContent content = getStreamingContent();
+      final Callable<Boolean> writeContent = new Callable<Boolean>() {
+        @Override
+        public Boolean call() throws IOException {
+          outputWriter.write(out, content);
+          return Boolean.TRUE;
+        }
+      };
+
+      final ExecutorService executor = Executors.newSingleThreadExecutor();
+      final Future<Boolean> future = executor.submit(new FutureTask<Boolean>(writeContent), null);
+      executor.shutdown();
+
+      try {
+        future.get(writeTimeout, TimeUnit.MILLISECONDS);
+      } catch (InterruptedException e) {
+        throw new IOException("Socket write interrupted", e);
+      } catch (ExecutionException e) {
+        throw new IOException("Exception in socket write", e);
+      } catch (TimeoutException e) {
+        throw new IOException("Socket write timed out", e);
+      }
+      if (!executor.isTerminated()) {
+        executor.shutdown();
+      }
+    }
+  }
 }

google-http-client/src/test/java/com/google/api/client/http/javanet/NetHttpRequestTest.java

@@ -0,0 +1,85 @@
+package com.google.api.client.http.javanet;
+
+import com.google.api.client.http.HttpContent;
+import com.google.api.client.http.InputStreamContent;
+import com.google.api.client.http.javanet.NetHttpRequest.OutputWriter;
+import com.google.api.client.testing.http.HttpTesting;
+import com.google.api.client.testing.http.javanet.MockHttpURLConnection;
+import com.google.api.client.util.StreamingContent;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.URL;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import java.util.concurrent.TimeoutException;
+import org.junit.Test;
+
+public class NetHttpRequestTest {
+
+  static class SleepingOutputWriter implements OutputWriter {
+    private long sleepTimeInMs;
+    public SleepingOutputWriter(long sleepTimeInMs) {
+      this.sleepTimeInMs = sleepTimeInMs;
+    }
+    @Override
+    public void write(OutputStream outputStream, StreamingContent content) throws IOException {
+      try {
+        Thread.sleep(sleepTimeInMs);
+      } catch (InterruptedException e) {
+        throw new IOException("sleep interrupted", e);
+      }
+    }
+  }
+
+  @Test
+  public void testHangingWrite() throws InterruptedException {
+    Thread thread = new Thread() {
+      @Override
+      public void run() {
+        try {
+          postWithTimeout(0);
+        } catch (IOException e) {
+          // expected to be interrupted
+          assertEquals(e.getCause().getClass(), InterruptedException.class);
+          return;
+        } catch (Exception e) {
+          fail();
+        }
+        fail("should be interrupted before here");
+      }
+    };
+
+    thread.start();
+    Thread.sleep(1000);
+    assertTrue(thread.isAlive());
+    thread.interrupt();
+  }
+
+  @Test(timeout = 1000)
+  public void testOutputStreamWriteTimeout() throws Exception {
+    try {
+      postWithTimeout(100);
+      fail("should have timed out");
+    } catch (IOException e) {
+      assertEquals(e.getCause().getClass(), TimeoutException.class);
+    } catch (Exception e) {
+      fail("Expected an IOException not a " + e.getCause().getClass().getName());
+    }
+  }
+
+  private static void postWithTimeout(int timeout) throws Exception {
+    MockHttpURLConnection connection = new MockHttpURLConnection(new URL(HttpTesting.SIMPLE_URL));
+    connection.setRequestMethod("POST");
+    NetHttpRequest request = new NetHttpRequest(connection);
+    InputStream is = NetHttpRequestTest.class.getClassLoader().getResourceAsStream("file.txt");
+    HttpContent content = new InputStreamContent("text/plain", is);
+    request.setStreamingContent(content);
+    request.setWriteTimeout(timeout);
+    request.execute(new SleepingOutputWriter(5000L));
+  }
+
+}

google-http-client/src/test/resources/file.txt

@@ -0,0 +1 @@
+some sample file