Skip to content

Commit a026ff3

Browse files
ibauersachsibauersachs
committed
Fix initial request using recommended nanoTime calculation
#345
1 parent 265adac commit a026ff3

2 files changed

Lines changed: 96 additions & 141 deletions

File tree

‎src/main/java/org/xbill/DNS/DohResolver.java‎

Lines changed: 16 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import java.util.concurrent.Executor;
2424
import java.util.concurrent.ForkJoinPool;
2525
import java.util.concurrent.TimeUnit;
26-
import java.util.concurrent.atomic.AtomicBoolean;
2726
import java.util.concurrent.atomic.AtomicLong;
2827
import java.util.function.Function;
2928
import javax.net.ssl.HttpsURLConnection;
@@ -92,8 +91,6 @@ public final class DohResolver implements Resolver {
9291
private final AsyncSemaphore maxConcurrentRequests;
9392

9493
private final AtomicLong lastRequest = new AtomicLong(0);
95-
96-
private final AtomicBoolean initialRequestSentMark = new AtomicBoolean(false);
9794
private final AsyncSemaphore initialRequestLock = new AsyncSemaphore(1);
9895

9996
private static final String APPLICATION_DNS_MESSAGE = "application/dns-message";
@@ -177,6 +174,11 @@ public final class DohResolver implements Resolver {
177174
USE_HTTP_CLIENT = initSuccess;
178175
}
179176

177+
// package-visible for testing
178+
long getNanoTime() {
179+
return System.nanoTime();
180+
}
181+
180182
/**
181183
* Creates a new DoH resolver that performs lookups with HTTP GET and the default timeout (5s).
182184
*
@@ -318,7 +320,7 @@ public CompletionStage<Message> sendAsync(Message query, Executor executor) {
318320
private CompletionStage<Message> sendAsync8(final Message query, Executor executor) {
319321
byte[] queryBytes = prepareQuery(query).toWire();
320322
String url = getUrl(queryBytes);
321-
long startTime = System.nanoTime();
323+
long startTime = getNanoTime();
322324
return maxConcurrentRequests
323325
.acquire(timeout)
324326
.handleAsync(
@@ -366,7 +368,7 @@ private SendAndGetMessageBytesResponse sendAndGetMessageBytes(
366368
((HttpsURLConnection) conn).setSSLSocketFactory(sslSocketFactory);
367369
}
368370

369-
Duration remainingTimeout = timeout.minus(System.nanoTime() - startTime, ChronoUnit.NANOS);
371+
Duration remainingTimeout = timeout.minus(getNanoTime() - startTime, ChronoUnit.NANOS);
370372
conn.setConnectTimeout((int) remainingTimeout.toMillis());
371373
conn.setReadTimeout((int) remainingTimeout.toMillis());
372374
conn.setRequestMethod(usePost ? "POST" : "GET");
@@ -392,7 +394,7 @@ private SendAndGetMessageBytesResponse sendAndGetMessageBytes(
392394
int offset = 0;
393395
while ((r = is.read(responseBytes, offset, responseBytes.length - offset)) > 0) {
394396
offset += r;
395-
remainingTimeout = timeout.minus(System.nanoTime() - startTime, ChronoUnit.NANOS);
397+
remainingTimeout = timeout.minus(getNanoTime() - startTime, ChronoUnit.NANOS);
396398
if (remainingTimeout.isNegative()) {
397399
throw new SocketTimeoutException();
398400
}
@@ -406,7 +408,7 @@ private SendAndGetMessageBytesResponse sendAndGetMessageBytes(
406408
byte[] buffer = new byte[4096];
407409
int r;
408410
while ((r = is.read(buffer, 0, buffer.length)) > 0) {
409-
remainingTimeout = timeout.minus(System.nanoTime() - startTime, ChronoUnit.NANOS);
411+
remainingTimeout = timeout.minus(getNanoTime() - startTime, ChronoUnit.NANOS);
410412
if (remainingTimeout.isNegative()) {
411413
throw new SocketTimeoutException();
412414
}
@@ -435,7 +437,7 @@ private void discardStream(InputStream es) throws IOException {
435437
}
436438

437439
private CompletionStage<Message> sendAsync11(final Message query, Executor executor) {
438-
long startTime = System.nanoTime();
440+
long startTime = getNanoTime();
439441
byte[] queryBytes = prepareQuery(query).toWire();
440442
String url = getUrl(queryBytes);
441443

@@ -457,7 +459,7 @@ private CompletionStage<Message> sendAsync11(final Message query, Executor execu
457459
// check if this request needs to be done synchronously because of HttpClient's stupidity to
458460
// not use the connection pool for HTTP/2 until one connection is successfully established,
459461
// which could lead to hundreds of connections (and threads with the default executor)
460-
Duration remainingTimeout = timeout.minus(System.nanoTime() - startTime, ChronoUnit.NANOS);
462+
Duration remainingTimeout = timeout.minus(getNanoTime() - startTime, ChronoUnit.NANOS);
461463
return initialRequestLock
462464
.acquire(remainingTimeout)
463465
.handle(
@@ -472,34 +474,20 @@ private CompletionStage<Message> sendAsync11(final Message query, Executor execu
472474
.thenCompose(Function.identity());
473475
}
474476

475-
/**
476-
* Check whether current initiating DoH request is initial request of this {@link DohResolver}.
477-
*/
478-
private boolean checkInitialRequest() {
479-
// If initial request haven't been completed successfully yet, just return true.
480-
if (!initialRequestSentMark.get()) {
481-
return true;
482-
}
483-
484-
// Otherwise, check whether such request is happened
485-
// after last successful request plus idle connection timeout.
486-
long lastRequestTime = lastRequest.get();
487-
return (lastRequestTime + idleConnectionTimeout.toNanos() < System.nanoTime());
488-
}
489-
490477
private CompletionStage<Message> sendAsync11WithInitialRequestPermit(
491478
Message query,
492479
Executor executor,
493480
long startTime,
494481
Object requestBuilder,
495482
Permit initialRequestPermit) {
496-
boolean isInitialRequest = checkInitialRequest();
483+
long lastRequestTime = lastRequest.get();
484+
boolean isInitialRequest = idleConnectionTimeout.toNanos() > getNanoTime() - lastRequestTime;
497485
if (!isInitialRequest) {
498486
initialRequestPermit.release();
499487
}
500488

501489
// check if we already exceeded the query timeout while checking the initial connection
502-
Duration remainingTimeout = timeout.minus(System.nanoTime() - startTime, ChronoUnit.NANOS);
490+
Duration remainingTimeout = timeout.minus(getNanoTime() - startTime, ChronoUnit.NANOS);
503491
if (remainingTimeout.isNegative()) {
504492
if (isInitialRequest) {
505493
initialRequestPermit.release();
@@ -532,25 +520,6 @@ private CompletionStage<Message> sendAsync11WithInitialRequestPermit(
532520
.thenCompose(Function.identity());
533521
}
534522

535-
/**
536-
* Set last request time to {@link DohResolver#lastRequest}, which ensures only the largest timestamp could be accepted.
537-
*
538-
* @param startTime start time in nanos of a Doh request.
539-
*/
540-
private void setLastRequestTime(long startTime) {
541-
long current = lastRequest.get();
542-
// Only update value of 'lastRequest' if timestamp in 'lastRequest' is smaller than incoming 'startTime' value.
543-
if (current < startTime) {
544-
while (!lastRequest.compareAndSet(current, startTime)) {
545-
// CAS failed, re-verify the eligibility of timestamp in 'lastRequest' to be updated to the incoming 'startTime' value.
546-
current = lastRequest.get();
547-
if (current > startTime) {
548-
return;
549-
}
550-
}
551-
}
552-
}
553-
554523
private CompletionStage<Message> sendAsync11WithConcurrentRequestPermit(
555524
Message query,
556525
Executor executor,
@@ -560,7 +529,7 @@ private CompletionStage<Message> sendAsync11WithConcurrentRequestPermit(
560529
boolean isInitialRequest,
561530
Permit maxConcurrentRequestPermit) {
562531
// check if the stream lock acquisition took too long
563-
Duration remainingTimeout = timeout.minus(System.nanoTime() - startTime, ChronoUnit.NANOS);
532+
Duration remainingTimeout = timeout.minus(getNanoTime() - startTime, ChronoUnit.NANOS);
564533
if (remainingTimeout.isNegative()) {
565534
if (isInitialRequest) {
566535
initialRequestPermit.release();
@@ -583,12 +552,7 @@ private CompletionStage<Message> sendAsync11WithConcurrentRequestPermit(
583552
.whenComplete(
584553
(result, ex) -> {
585554
if (ex == null) {
586-
setLastRequestTime(startTime);
587-
if (isInitialRequest) {
588-
// initial request was completed successfully, so toggle initialRequestSentMark to true.
589-
// it's very safe to toggle initialRequestSentMark here, since this code had been guarded by initialRequestLock and its permit outside.
590-
initialRequestSentMark.compareAndSet(false, true);
591-
}
555+
lastRequest.set(startTime);
592556
}
593557
maxConcurrentRequestPermit.release();
594558
if (isInitialRequest) {

‎src/test/java/org/xbill/DNS/DohResolverTest.java‎

Lines changed: 80 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
// SPDX-License-Identifier: BSD-3-Clause
22
package org.xbill.DNS;
33

4+
import static org.junit.jupiter.api.Assertions.assertEquals;
5+
import static org.junit.jupiter.api.Assertions.assertTrue;
6+
import static org.mockito.Mockito.doAnswer;
7+
import static org.mockito.Mockito.spy;
8+
49
import io.netty.handler.codec.http.HttpHeaderNames;
510
import io.vertx.core.Future;
611
import io.vertx.core.Vertx;
@@ -18,16 +23,19 @@
1823
import java.util.Base64;
1924
import java.util.Collections;
2025
import java.util.concurrent.CompletionStage;
26+
import java.util.concurrent.TimeUnit;
2127
import java.util.concurrent.TimeoutException;
2228
import java.util.concurrent.atomic.AtomicBoolean;
2329
import java.util.concurrent.atomic.AtomicInteger;
30+
import java.util.concurrent.atomic.AtomicLong;
2431
import org.junit.jupiter.api.BeforeEach;
2532
import org.junit.jupiter.api.Test;
33+
import org.junit.jupiter.api.condition.EnabledForJreRange;
34+
import org.junit.jupiter.api.condition.JRE;
2635
import org.junit.jupiter.api.extension.ExtendWith;
2736
import org.junit.jupiter.params.ParameterizedTest;
2837
import org.junit.jupiter.params.provider.ValueSource;
29-
30-
import static org.junit.jupiter.api.Assertions.*;
38+
import org.mockito.stubbing.Answer;
3139

3240
@ExtendWith(VertxExtension.class)
3341
class DohResolverTest {
@@ -153,88 +161,6 @@ void initialRequestSlowResolve(Vertx vertx, VertxTestContext context) {
153161
});
154162
}
155163

156-
157-
@Test
158-
void initialRequestGuardIfIdleConnectionTimeIsLargerThanSystemNanoTime(Vertx vertx, VertxTestContext context) {
159-
if (isPreJava9()) {
160-
System.out.println("Current JVM is PreJava9, no need to run such test.");
161-
context.completeNow();
162-
return;
163-
}
164-
resolver = new DohResolver("http://localhost",
165-
2,
166-
// so long idleConnectionTimeout
167-
// in order to hack the condition for checking initial request in org.xbill.DNS.DohResolver.checkInitialRequest
168-
Duration.ofNanos(System.nanoTime() + Duration.ofSeconds(100L).toNanos()));
169-
resolver.setTimeout(Duration.ofSeconds(1));
170-
// Just add a 100ms delay before responding to the 1st call
171-
// to simulate a 'concurrent doh request' for the 2nd call,
172-
// then let the fake dns server respond to the 2nd call ASAP.
173-
allRequestsUseTimeout = false;
174-
175-
// idleConnectionTimeout = 2s, lastRequest = 0L
176-
// Ensure lastRequest + idleConnectionTimeout < System.nanoTime() (3s)
177-
178-
// Timeline:
179-
// |<-------- 100ms -------->|
180-
// ↑ ↑
181-
// 1st call sent response of 1st call
182-
// |20ms|<------ 80ms ------>|<------ few millis ------->|
183-
// ↑ wait until 1st call ↑ ↑
184-
// 2nd call begin 2nd call sent response of 2nd call
185-
186-
AtomicBoolean firstCallCompleted = new AtomicBoolean(false);
187-
188-
setupResolverWithServer(Duration.ofMillis(100L),
189-
200,
190-
2,
191-
vertx,
192-
context)
193-
.onSuccess(
194-
server -> {
195-
// First call
196-
CompletionStage<Message> firstCall = resolver.sendAsync(qm);
197-
// Ensure second call was made after first call.
198-
sleepNotThrown(20L);
199-
CompletionStage<Message> secondCall = resolver.sendAsync(qm);
200-
201-
Future.fromCompletionStage(firstCall)
202-
.onComplete(
203-
context.succeeding(
204-
result ->
205-
context.verify(
206-
() -> {
207-
assertEquals(Rcode.NOERROR, result.getHeader().getRcode());
208-
assertEquals(0, result.getHeader().getID());
209-
assertEquals(queryName, result.getQuestion().getName());
210-
firstCallCompleted.set(true);
211-
})));
212-
213-
Future.fromCompletionStage(secondCall)
214-
.onComplete(
215-
context.succeeding(
216-
result ->
217-
context.verify(
218-
() -> {
219-
assertTrue(firstCallCompleted.get());
220-
assertEquals(Rcode.NOERROR, result.getHeader().getRcode());
221-
assertEquals(0, result.getHeader().getID());
222-
assertEquals(queryName, result.getQuestion().getName());
223-
// Complete context after the 2nd call was completed.
224-
context.completeNow();
225-
})));
226-
}
227-
);
228-
}
229-
230-
private static void sleepNotThrown(long millis) {
231-
try {
232-
Thread.sleep(millis);
233-
} catch (InterruptedException e) {
234-
throw new RuntimeException(e);
235-
}
236-
}
237-
238164
@Test
239165
void initialRequestTimeoutResolve(Vertx vertx, VertxTestContext context) {
240166
resolver = new DohResolver("http://localhost", 2, Duration.ofMinutes(2));
@@ -275,10 +201,6 @@ void initialRequestTimeoutResolve(Vertx vertx, VertxTestContext context) {
275201
});
276202
}
277203

278-
private static boolean isPreJava9() {
279-
return System.getProperty("java.version").startsWith("1.");
280-
}
281-
282204
private Future<HttpServer> setupResolverWithServer(
283205
Duration responseDelay,
284206
int statusCode,
@@ -289,6 +211,75 @@ private Future<HttpServer> setupResolverWithServer(
289211
.onSuccess(server -> resolver.setUriTemplate("http://localhost:" + server.actualPort()));
290212
}
291213

214+
@EnabledForJreRange(
215+
min = JRE.JAVA_9,
216+
disabledReason = "Java 8 implementation doesn't have the initial request guard")
217+
@Test
218+
void initialRequestGuardIfIdleConnectionTimeIsLargerThanSystemNanoTime(
219+
Vertx vertx, VertxTestContext context) {
220+
AtomicLong startNanos = new AtomicLong(System.nanoTime());
221+
resolver = spy(new DohResolver("http://localhost", 2, Duration.ofMinutes(2)));
222+
resolver.setTimeout(Duration.ofSeconds(1));
223+
// Simulate a nanoTime value that is lower than the idle timeout
224+
doAnswer((Answer<Long>) invocationOnMock -> System.nanoTime() - startNanos.get())
225+
.when(resolver)
226+
.getNanoTime();
227+
228+
// Just add a 100ms delay before responding to the 1st call
229+
// to simulate a 'concurrent doh request' for the 2nd call,
230+
// then let the fake dns server respond to the 2nd call ASAP.
231+
allRequestsUseTimeout = false;
232+
233+
// idleConnectionTimeout = 2s, lastRequest = 0L
234+
// Ensure idleConnectionTimeout < System.nanoTime() - lastRequest (3s)
235+
236+
// Timeline:
237+
// |<-------- 100ms -------->|
238+
// ↑ ↑
239+
// 1st call sent response of 1st call
240+
// |20ms|<------ 80ms ------>|<------ few millis ------->|
241+
// ↑ wait until 1st call ↑ ↑
242+
// 2nd call begin 2nd call sent response of 2nd call
243+
244+
AtomicBoolean firstCallCompleted = new AtomicBoolean(false);
245+
246+
setupResolverWithServer(Duration.ofMillis(100L), 200, 2, vertx, context)
247+
.onSuccess(
248+
server -> {
249+
// First call
250+
CompletionStage<Message> firstCall = resolver.sendAsync(qm);
251+
// Ensure second call was made after first call and uses a different query
252+
startNanos.addAndGet(TimeUnit.MILLISECONDS.toNanos(20));
253+
CompletionStage<Message> secondCall = resolver.sendAsync(Message.newQuery(qr));
254+
255+
Future.fromCompletionStage(firstCall)
256+
.onComplete(
257+
context.succeeding(
258+
result ->
259+
context.verify(
260+
() -> {
261+
assertEquals(Rcode.NOERROR, result.getHeader().getRcode());
262+
assertEquals(0, result.getHeader().getID());
263+
assertEquals(queryName, result.getQuestion().getName());
264+
firstCallCompleted.set(true);
265+
})));
266+
267+
Future.fromCompletionStage(secondCall)
268+
.onComplete(
269+
context.succeeding(
270+
result ->
271+
context.verify(
272+
() -> {
273+
assertTrue(firstCallCompleted.get());
274+
assertEquals(Rcode.NOERROR, result.getHeader().getRcode());
275+
assertEquals(0, result.getHeader().getID());
276+
assertEquals(queryName, result.getQuestion().getName());
277+
// Complete context after the 2nd call was completed.
278+
context.completeNow();
279+
})));
280+
});
281+
}
282+
292283
private Future<HttpServer> setupServer(
293284
Message expectedDnsRequest,
294285
Message dnsResponse,
@@ -298,7 +289,7 @@ private Future<HttpServer> setupServer(
298289
VertxTestContext context,
299290
Vertx vertx) {
300291
HttpVersion version =
301-
isPreJava9()
292+
System.getProperty("java.version").startsWith("1.")
302293
? HttpVersion.HTTP_1_1
303294
: HttpVersion.HTTP_2;
304295
AtomicInteger requestCount = new AtomicInteger(0);

0 commit comments

Comments
 (0)