Provide capability to capture the Call stack trace on leaked connections.

Also plugged some leaks discovered after running the tests with this
change in place.

https://github.com/square/okhttp/issues/2127

okhttp-tests/src/test/java/okhttp3/CallTest.java

@@ -45,7 +45,9 @@ import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.logging.Level;
 import java.util.logging.Logger;
+import java.util.logging.SimpleFormatter;
 import javax.net.ServerSocketFactory;
 import javax.net.ssl.SSLHandshakeException;
 import javax.net.ssl.SSLPeerUnverifiedException;
@@ -79,6 +81,7 @@ import org.junit.rules.TestRule;
 import org.junit.rules.Timeout;
 
 import static java.net.CookiePolicy.ACCEPT_ORIGINAL_SERVER;
+import static okhttp3.TestUtil.awaitGarbageCollection;
 import static okhttp3.TestUtil.defaultClient;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -342,6 +345,7 @@ public final class CallTest {
 
     Response response = client.newCall(request).execute();
     assertEquals(200, response.code());
+    response.body().close();
 
     RecordedRequest recordedRequest1 = server.takeRequest();
     assertEquals("POST", recordedRequest1.getMethod());
@@ -2569,6 +2573,78 @@ public final class CallTest {
         .assertBody("");
   }
 
+  @Test public void leakedResponseBodyLogsStackTrace() throws Exception {
+    server.enqueue(new MockResponse()
+        .setBody("This gets leaked."));
+
+    client = new OkHttpClient.Builder()
+        .connectionPool(new ConnectionPool(0, 10, TimeUnit.MILLISECONDS))
+        .build();
+
+    Request request = new Request.Builder()
+        .url(server.url("/"))
+        .build();
+
+    Level original = logger.getLevel();
+    logger.setLevel(Level.FINE);
+    logHandler.setFormatter(new SimpleFormatter());
+    try {
+      client.newCall(request).execute(); // Ignore the response so it gets leaked then GC'd.
+      awaitGarbageCollection();
+
+      String message = logHandler.take();
+      assertTrue(message.contains("WARNING: A connection to " + server.url("/") + " was leaked."
+          + " Did you forget to close a response body?"));
+      assertTrue(message.contains("okhttp3.RealCall.execute("));
+      assertTrue(message.contains("okhttp3.CallTest.leakedResponseBodyLogsStackTrace("));
+    } finally {
+      logger.setLevel(original);
+    }
+  }
+
+  @Test public void asyncLeakedResponseBodyLogsStackTrace() throws Exception {
+    server.enqueue(new MockResponse()
+        .setBody("This gets leaked."));
+
+    client = new OkHttpClient.Builder()
+        .connectionPool(new ConnectionPool(0, 10, TimeUnit.MILLISECONDS))
+        .build();
+
+    Request request = new Request.Builder()
+        .url(server.url("/"))
+        .build();
+
+    Level original = logger.getLevel();
+    logger.setLevel(Level.FINE);
+    logHandler.setFormatter(new SimpleFormatter());
+    try {
+      final CountDownLatch latch = new CountDownLatch(1);
+      client.newCall(request).enqueue(new Callback() {
+        @Override public void onFailure(Call call, IOException e) {
+          fail();
+        }
+
+        @Override public void onResponse(Call call, Response response) throws IOException {
+          // Ignore the response so it gets leaked then GC'd.
+          latch.countDown();
+        }
+      });
+      latch.await();
+      // There's some flakiness when triggering a GC for objects in a separate thread. Adding a
+      // small delay appears to ensure the objects will get GC'd.
+      Thread.sleep(200);
+      awaitGarbageCollection();
+
+      String message = logHandler.take();
+      assertTrue(message.contains("WARNING: A connection to " + server.url("/") + " was leaked."
+          + " Did you forget to close a response body?"));
+      assertTrue(message.contains("okhttp3.RealCall.enqueue("));
+      assertTrue(message.contains("okhttp3.CallTest.asyncLeakedResponseBodyLogsStackTrace("));
+    } finally {
+      logger.setLevel(original);
+    }
+  }
+
   private void makeFailingCall() {
     RequestBody requestBody = new RequestBody() {
       @Override public MediaType contentType() {

okhttp-tests/src/test/java/okhttp3/ConnectionPoolTest.java

@@ -28,6 +28,7 @@ import okhttp3.internal.connection.RealConnection;
 import okhttp3.internal.connection.StreamAllocation;
 import org.junit.Test;
 
+import static okhttp3.TestUtil.awaitGarbageCollection;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -82,7 +83,7 @@ public final class ConnectionPoolTest {
 
     RealConnection c1 = newConnection(pool, routeA1, 50L);
     synchronized (pool) {
-      StreamAllocation streamAllocation = new StreamAllocation(pool, addressA);
+      StreamAllocation streamAllocation = new StreamAllocation(pool, addressA, null);
       streamAllocation.acquire(c1);
     }
 
@@ -175,22 +176,11 @@ public final class ConnectionPoolTest {
   /** Use a helper method so there's no hidden reference remaining on the stack. */
   private void allocateAndLeakAllocation(ConnectionPool pool, RealConnection connection) {
     synchronized (pool) {
-      StreamAllocation leak = new StreamAllocation(pool, connection.route().address());
+      StreamAllocation leak = new StreamAllocation(pool, connection.route().address(), null);
       leak.acquire(connection);
     }
   }
 
-  /**
-   * See FinalizationTester for discussion on how to best trigger GC in tests.
-   * https://android.googlesource.com/platform/libcore/+/master/support/src/test/java/libcore/
-   * java/lang/ref/FinalizationTester.java
-   */
-  private void awaitGarbageCollection() throws InterruptedException {
-    Runtime.getRuntime().gc();
-    Thread.sleep(100);
-    System.runFinalization();
-  }
-
   private RealConnection newConnection(ConnectionPool pool, Route route, long idleAtNanos) {
     RealConnection connection = new RealConnection(route);
     connection.idleAtNanos = idleAtNanos;

okhttp-tests/src/test/java/okhttp3/TestLogHandler.java

@@ -19,6 +19,7 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.logging.Handler;
 import java.util.logging.LogRecord;
+import java.util.logging.SimpleFormatter;
 
 /**
  * A log handler that records which log messages were published so that a calling test can make
@@ -28,7 +29,11 @@ public final class TestLogHandler extends Handler {
   private final List<String> logs = new ArrayList<>();
 
   @Override public synchronized void publish(LogRecord logRecord) {
-    logs.add(logRecord.getLevel() + ": " + logRecord.getMessage());
+    if (getFormatter() == null) {
+      logs.add(logRecord.getLevel() + ": " + logRecord.getMessage());
+    } else {
+      logs.add(getFormatter().format(logRecord));
+    }
     notifyAll();
   }
 

okhttp-tests/src/test/java/okhttp3/TestUtil.java

@@ -43,4 +43,15 @@ public final class TestUtil {
     Arrays.fill(array, c);
     return new String(array);
   }
+
+  /**
+   * See FinalizationTester for discussion on how to best trigger GC in tests.
+   * https://android.googlesource.com/platform/libcore/+/master/support/src/test/java/libcore/
+   * java/lang/ref/FinalizationTester.java
+   */
+  public static void awaitGarbageCollection() throws InterruptedException {
+    Runtime.getRuntime().gc();
+    Thread.sleep(100);
+    System.runFinalization();
+  }
 }

okhttp-tests/src/test/java/okhttp3/URLConnectionTest.java

@@ -1549,6 +1549,7 @@ public final class URLConnectionTest {
     outputStream.write(body.getBytes("UTF-8"));
     outputStream.close();
     assertEquals(200, connection.getResponseCode());
+    connection.getInputStream().close();
 
     RecordedRequest recordedRequest1 = server.takeRequest();
     assertEquals("POST", recordedRequest1.getMethod());
@@ -2399,6 +2400,7 @@ public final class URLConnectionTest {
 
     assertEquals(408, connection.getResponseCode());
     assertEquals(1, server.getRequestCount());
+    connection.getErrorStream().close();
   }
 
   @Test public void readTimeouts() throws IOException {

okhttp/src/main/java/okhttp3/ConnectionPool.java

@@ -33,7 +33,6 @@ import okhttp3.internal.connection.StreamAllocation;
 import okhttp3.internal.platform.Platform;
 
 import static okhttp3.internal.Util.closeQuietly;
-import static okhttp3.internal.platform.Platform.WARN;
 
 /**
  * Manages reuse of HTTP and HTTP/2 connections for reduced network latency. HTTP requests that
@@ -246,8 +245,12 @@ public final class ConnectionPool {
       }
 
       // We've discovered a leaked allocation. This is an application bug.
-      Platform.get().log(WARN, "A connection to " + connection.route().address().url()
-          + " was leaked. Did you forget to close a response body?", null);
+      StreamAllocation.StreamAllocationReference streamAllocRef =
+          (StreamAllocation.StreamAllocationReference) reference;
+      String message = "A connection to " + connection.route().address().url()
+          + " was leaked. Did you forget to close a response body?";
+      Platform.get().logCloseableLeak(message, streamAllocRef.callStackTrace);
+
       references.remove(i);
       connection.noNewStreams = true;
 

okhttp/src/main/java/okhttp3/RealCall.java

@@ -55,6 +55,7 @@ final class RealCall implements Call {
       if (executed) throw new IllegalStateException("Already Executed");
       executed = true;
     }
+    captureCallStackTrace();
     try {
       client.dispatcher().executed(this);
       Response result = getResponseWithInterceptorChain();
@@ -65,6 +66,11 @@ final class RealCall implements Call {
     }
   }
 
+  private void captureCallStackTrace() {
+    Object callStackTrace = Platform.get().getStackTraceForCloseable("response.body().close()");
+    retryAndFollowUpInterceptor.setCallStackTrace(callStackTrace);
+  }
+
   synchronized void setForWebSocket() {
     if (executed) throw new IllegalStateException("Already Executed");
     this.retryAndFollowUpInterceptor.setForWebSocket(true);
@@ -75,6 +81,7 @@ final class RealCall implements Call {
       if (executed) throw new IllegalStateException("Already Executed");
       executed = true;
     }
+    captureCallStackTrace();
     client.dispatcher().enqueue(new AsyncCall(responseCallback));
   }
 

okhttp/src/main/java/okhttp3/internal/connection/StreamAllocation.java

@@ -73,6 +73,7 @@ public final class StreamAllocation {
   public final Address address;
   private Route route;
   private final ConnectionPool connectionPool;
+  private final Object callStackTrace;
 
   // State guarded by connectionPool.
   private final RouteSelector routeSelector;
@@ -82,10 +83,11 @@ public final class StreamAllocation {
   private boolean canceled;
   private HttpCodec codec;
 
-  public StreamAllocation(ConnectionPool connectionPool, Address address) {
+  public StreamAllocation(ConnectionPool connectionPool, Address address, Object callStackTrace) {
     this.connectionPool = connectionPool;
     this.address = address;
     this.routeSelector = new RouteSelector(address, routeDatabase());
+    this.callStackTrace = callStackTrace;
   }
 
   public HttpCodec newStream(OkHttpClient client, boolean doExtensiveHealthChecks) {
@@ -318,7 +320,7 @@ public final class StreamAllocation {
    */
   public void acquire(RealConnection connection) {
     assert (Thread.holdsLock(connectionPool));
-    connection.allocations.add(new WeakReference<>(this));
+    connection.allocations.add(new StreamAllocationReference(this, callStackTrace));
   }
 
   /** Remove this allocation from the connection's list of allocations. */
@@ -340,4 +342,17 @@ public final class StreamAllocation {
   @Override public String toString() {
     return address.toString();
   }
+
+  public static final class StreamAllocationReference extends WeakReference<StreamAllocation> {
+    /**
+     * Captures the stack trace at the time the Call is executed or enqueued. This is helpful for
+     * identifying the origin of connection leaks.
+     */
+    public final Object callStackTrace;
+
+    StreamAllocationReference(StreamAllocation referent, Object callStackTrace) {
+      super(referent);
+      this.callStackTrace = callStackTrace;
+    }
+  }
 }

okhttp/src/main/java/okhttp3/internal/http/RetryAndFollowUpInterceptor.java

@@ -64,6 +64,7 @@ public final class RetryAndFollowUpInterceptor implements Interceptor {
   private final OkHttpClient client;
   private StreamAllocation streamAllocation;
   private boolean forWebSocket;
+  private Object callStackTrace;
   private volatile boolean canceled;
 
   public RetryAndFollowUpInterceptor(OkHttpClient client) {
@@ -97,6 +98,10 @@ public final class RetryAndFollowUpInterceptor implements Interceptor {
     return forWebSocket;
   }
 
+  public void setCallStackTrace(Object callStackTrace) {
+    this.callStackTrace = callStackTrace;
+  }
+
   public StreamAllocation streamAllocation() {
     return streamAllocation;
   }
@@ -105,7 +110,7 @@ public final class RetryAndFollowUpInterceptor implements Interceptor {
     Request request = chain.request();
 
     streamAllocation = new StreamAllocation(
-        client.connectionPool(), createAddress(request.url()));
+        client.connectionPool(), createAddress(request.url()), callStackTrace);
 
     int followUpCount = 0;
     Response priorResponse = null;
@@ -164,13 +169,14 @@ public final class RetryAndFollowUpInterceptor implements Interceptor {
       }
 
       if (followUp.body() instanceof UnrepeatableRequestBody) {
+        streamAllocation.release();
         throw new HttpRetryException("Cannot retry streamed HTTP body", response.code());
       }
 
       if (!sameConnection(response, followUp.url())) {
         streamAllocation.release();
         streamAllocation = new StreamAllocation(
-            client.connectionPool(), createAddress(followUp.url()));
+            client.connectionPool(), createAddress(followUp.url()), callStackTrace);
       } else if (streamAllocation.codec() != null) {
         throw new IllegalStateException("Closing the body of " + response
             + " didn't close its backing stream. Bad interceptor?");

okhttp/src/main/java/okhttp3/internal/platform/AndroidPlatform.java

@@ -45,6 +45,8 @@ class AndroidPlatform extends Platform {
   private final OptionalMethod<Socket> getAlpnSelectedProtocol;
   private final OptionalMethod<Socket> setAlpnProtocols;
 
+  private final CloseGuard closeGuard = CloseGuard.get();
+
   public AndroidPlatform(Class<?> sslParametersClass, OptionalMethod<Socket> setUseSessionTickets,
       OptionalMethod<Socket> setHostname, OptionalMethod<Socket> getAlpnSelectedProtocol,
       OptionalMethod<Socket> setAlpnProtocols) {
@@ -132,6 +134,18 @@ class AndroidPlatform extends Platform {
     }
   }
 
+  @Override public Object getStackTraceForCloseable(String closer) {
+    return closeGuard.createAndOpen(closer);
+  }
+
+  @Override public void logCloseableLeak(String message, Object stackTrace) {
+    boolean reported = closeGuard.warnIfOpen(stackTrace);
+    if (!reported) {
+      // Unable to report via CloseGuard. As a last-ditch effort, send it to the logger.
+      log(WARN, message, null);
+    }
+  }
+
   @Override public boolean isCleartextTrafficPermitted(String hostname) {
     try {
       Class<?> networkPolicyClass = Class.forName("android.security.NetworkSecurityPolicy");
@@ -226,4 +240,63 @@ class AndroidPlatform extends Platform {
       }
     }
   }
+
+  /**
+   * Provides access to the internal dalvik.system.CloseGuard class. Android uses this in
+   * combination with android.os.StrictMode to report on leaked java.io.Closeable's. Available since
+   * Android API 11.
+   */
+  static final class CloseGuard {
+    private final Method getMethod;
+    private final Method openMethod;
+    private final Method warnIfOpenMethod;
+
+    CloseGuard(Method getMethod, Method openMethod, Method warnIfOpenMethod) {
+      this.getMethod = getMethod;
+      this.openMethod = openMethod;
+      this.warnIfOpenMethod = warnIfOpenMethod;
+    }
+
+    Object createAndOpen(String closer) {
+      if (getMethod != null) {
+        try {
+          Object closeGuardInstance = getMethod.invoke(null);
+          openMethod.invoke(closeGuardInstance, closer);
+          return closeGuardInstance;
+        } catch (Exception ignored) {
+        }
+      }
+      return null;
+    }
+
+    boolean warnIfOpen(Object closeGuardInstance) {
+      boolean reported = false;
+      if (closeGuardInstance != null) {
+        try {
+          warnIfOpenMethod.invoke(closeGuardInstance);
+          reported = true;
+        } catch (Exception ignored) {
+        }
+      }
+      return reported;
+    }
+
+    static CloseGuard get() {
+      Method getMethod;
+      Method openMethod;
+      Method warnIfOpenMethod;
+
+      try {
+        Class<?> closeGuardClass = Class.forName("dalvik.system.CloseGuard");
+        getMethod = closeGuardClass.getMethod("get");
+        openMethod = closeGuardClass.getMethod("open", String.class);
+        warnIfOpenMethod = closeGuardClass.getMethod("warnIfOpen");
+      } catch (Exception ignored) {
+        getMethod = null;
+        openMethod = null;
+        warnIfOpenMethod = null;
+      }
+      return new CloseGuard(getMethod, openMethod, warnIfOpenMethod);
+    }
+  }
 }

okhttp/src/main/java/okhttp3/internal/platform/Platform.java

@@ -133,6 +133,26 @@ public class Platform {
     return true;
   }
 
+  /**
+   * Returns an object that holds a stack trace created at the moment this method is executed. This
+   * should be used specifically for {@link java.io.Closeable} objects and in conjunction with
+   * {@link #logCloseableLeak(String, Object)}.
+   */
+  public Object getStackTraceForCloseable(String closer) {
+    if (logger.isLoggable(Level.FINE)) {
+      return new Throwable(closer); // These are expensive to allocate.
+    }
+    return null;
+  }
+
+  public void logCloseableLeak(String message, Object stackTrace) {
+    if (stackTrace == null) {
+      message += " To see where this was allocated, set the OkHttpClient logger level to FINE: "
+          + "Logger.getLogger(OkHttpClient.class.getName()).setLevel(Level.FINE);";
+    }
+    log(WARN, message, (Throwable) stackTrace);
+  }
+
   public static List<String> alpnProtocolNames(List<Protocol> protocols) {
     List<String> names = new ArrayList<>(protocols.size());
     for (int i = 0, size = protocols.size(); i < size; i++) {