Close the response before throwing RateLimitedException

src/main/java/io/ipinfo/api/IPinfo.java

@@ -24,6 +24,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiConsumer;
 import javax.annotation.ParametersAreNonnullByDefault;
 import okhttp3.*;
@@ -277,6 +278,8 @@ private ConcurrentHashMap<String, Object> getBatchGeneric(
         CountDownLatch latch = new CountDownLatch(
             (int) Math.ceil(lookupUrls.size() / 1000.0)
         );
+        // set if any request is rate limited; checked once all are done.
+        AtomicBoolean rateLimited = new AtomicBoolean(false);
         Request.Builder reqCommon = new Request.Builder()
             .url(postUrl)
             .addHeader("Content-Type", "application/json")
@@ -300,32 +303,34 @@ private ConcurrentHashMap<String, Object> getBatchGeneric(
                 .connectTimeout(timeoutPerBatch, TimeUnit.SECONDS)
                 .readTimeout(timeoutPerBatch, TimeUnit.SECONDS)
                 .build();
-            chunkClient
-                .newCall(req)
-                .enqueue(
-                    new Callback() {
-                        @Override
-                        @ParametersAreNonnullByDefault
-                        public void onFailure(Call call, IOException e) {
-                            latch.countDown();
-                        }
+            chunkClient.newCall(req).enqueue(
+                new Callback() {
+                    @Override
+                    @ParametersAreNonnullByDefault
+                    public void onFailure(Call call, IOException e) {
+                        latch.countDown();
+                    }
 
-                        @Override
-                        @ParametersAreNonnullByDefault
-                        public void onResponse(Call call, Response response)
-                            throws IOException {
-                            if (
-                                response.body() == null ||
-                                response.code() == 429
-                            ) {
+                    @Override
+                    @ParametersAreNonnullByDefault
+                    public void onResponse(Call call, Response response)
+                        throws IOException {
+                        // the body is closed before the latch counts down,
+                        // whichever way we leave.
+                        try (ResponseBody body = response.body()) {
+                            if (response.code() == 429) {
+                                rateLimited.set(true);
+                                return;
+                            }
+                            if (body == null) {
                                 return;
                             }
 
                             Type respType = new TypeToken<
                                 HashMap<String, Object>
                             >() {}.getType();
                             HashMap<String, Object> localResult = gson.fromJson(
-                                response.body().string(),
+                                body.string(),
                                 respType
                             );
                             localResult.forEach(
@@ -356,11 +361,12 @@ public void accept(String k, Object v) {
                                     }
                                 }
                             );
-
+                        } finally {
                             latch.countDown();
                         }
                     }
-                );
+                }
+            );
         }
 
         // wait for all requests to finish.
@@ -388,6 +394,10 @@ public void accept(String k, Object v) {
             }
         }
 
+        if (rateLimited.get()) {
+            throw new RateLimitedException();
+        }
+
         // insert any new lookups into the cache:
         if (cache != null) {
             for (String url : lookupUrls) {

src/main/java/io/ipinfo/api/request/BaseRequest.java

@@ -42,6 +42,7 @@ public Response handleRequest(Request.Builder request)
         }
 
         if (response.code() == 429) {
+            response.close();
             throw new RateLimitedException();
         }
 

src/test/java/io/ipinfo/api/request/BaseRequestTest.java

@@ -0,0 +1,111 @@
+package io.ipinfo.api.request;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import io.ipinfo.api.IPinfo;
+import io.ipinfo.api.errors.ErrorResponseException;
+import io.ipinfo.api.errors.RateLimitedException;
+import java.time.Duration;
+import java.util.Arrays;
+import okhttp3.MediaType;
+import okhttp3.OkHttpClient;
+import okhttp3.Protocol;
+import okhttp3.Request;
+import okhttp3.Response;
+import okhttp3.ResponseBody;
+import okio.Buffer;
+import okio.BufferedSource;
+import okio.ForwardingSource;
+import okio.Okio;
+import org.junit.jupiter.api.Test;
+
+public class BaseRequestTest {
+
+    @Test
+    public void testRateLimitedResponseIsClosed() {
+        TrackingBody body = new TrackingBody("{\"error\":\"rate limited\"}");
+        BaseRequest<Void> request = requestReturning(429, body);
+
+        assertThrows(RateLimitedException.class, () ->
+            request.handleRequest(new Request.Builder().url("https://api.ipinfo.io/lookup/8.8.8.8"))
+        );
+        assertTrue(body.closed, "429 response body should be closed");
+    }
+
+    @Test
+    public void testErrorResponseIsClosed() {
+        TrackingBody body = new TrackingBody("{\"error\":\"forbidden\"}");
+        BaseRequest<Void> request = requestReturning(403, body);
+
+        assertThrows(ErrorResponseException.class, () ->
+            request.handleRequest(new Request.Builder().url("https://api.ipinfo.io/lookup/8.8.8.8"))
+        );
+        assertTrue(body.closed, "error response body should be closed");
+    }
+
+    @Test
+    public void testBatchRateLimitedResponseIsClosed() {
+        TrackingBody body = new TrackingBody("{\"error\":\"rate limited\"}");
+        IPinfo ii = new IPinfo.Builder().setClient(clientReturning(429, body)).build();
+
+        // Without a total timeout a batch that never completes blocks forever.
+        assertTimeoutPreemptively(Duration.ofSeconds(5), () ->
+            assertThrows(RateLimitedException.class, () -> ii.getBatch(Arrays.asList("8.8.8.8")))
+        );
+        assertTrue(body.closed, "429 batch response body should be closed");
+    }
+
+    private static BaseRequest<Void> requestReturning(int code, ResponseBody body) {
+        return new BaseRequest<Void>(clientReturning(code, body), "test_token") {
+            @Override
+            public Void handle() {
+                return null;
+            }
+        };
+    }
+
+    private static OkHttpClient clientReturning(int code, ResponseBody body) {
+        return new OkHttpClient.Builder()
+            .addInterceptor(chain -> new Response.Builder()
+                .request(chain.request())
+                .protocol(Protocol.HTTP_1_1)
+                .code(code)
+                .message("test")
+                .body(body)
+                .build())
+            .build();
+    }
+
+    /** Records whether the underlying source was closed, which is what releases the connection. */
+    private static final class TrackingBody extends ResponseBody {
+        private final Buffer buffer;
+        private final BufferedSource source;
+        volatile boolean closed;
+
+        TrackingBody(String content) {
+            this.buffer = new Buffer().writeUtf8(content);
+            this.source = Okio.buffer(new ForwardingSource(buffer) {
+                @Override
+                public void close() throws java.io.IOException {
+                    closed = true;
+                    super.close();
+                }
+            });
+        }
+
+        @Override
+        public MediaType contentType() {
+            return MediaType.get("application/json");
+        }
+
+        @Override
+        public long contentLength() {
+            return buffer.size();
+        }
+
+        @Override
+        public BufferedSource source() {
+            return source;
+        }
+    }
+}