diff --git a/src/main/java/io/ipinfo/api/IPinfo.java b/src/main/java/io/ipinfo/api/IPinfo.java index cb4366a..d87ba69 100644 --- a/src/main/java/io/ipinfo/api/IPinfo.java +++ b/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 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,24 +303,26 @@ private ConcurrentHashMap 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; } @@ -325,7 +330,7 @@ public void onResponse(Call call, Response response) HashMap >() {}.getType(); HashMap 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) { diff --git a/src/main/java/io/ipinfo/api/request/BaseRequest.java b/src/main/java/io/ipinfo/api/request/BaseRequest.java index 481f4c7..5936eca 100644 --- a/src/main/java/io/ipinfo/api/request/BaseRequest.java +++ b/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(); } diff --git a/src/test/java/io/ipinfo/api/request/BaseRequestTest.java b/src/test/java/io/ipinfo/api/request/BaseRequestTest.java new file mode 100644 index 0000000..da7f58a --- /dev/null +++ b/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 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 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 requestReturning(int code, ResponseBody body) { + return new BaseRequest(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; + } + } +}