private void scheduleRequestTimeout(RequestHeader requestHeader) {
UInteger requestHandle = requestHeader.getRequestHandle();
long timeoutHint =
requestHeader.getTimeoutHint() != null
? requestHeader.getTimeoutHint().longValue()
: DEFAULT_TIMEOUT_MS;
Timeout timeout =
wheelTimer.newTimeout(
t -> {
timeouts.remove(requestHandle);
if (!t.isCancelled()) {
CompletableFuture<UaResponseMessage> f = pending.remove(requestHandle);
if (f != null) {
String message = "request timed out after " + timeoutHint + "ms";
f.completeExceptionally(new UaException(StatusCodes.Bad_Timeout, message));
}
}
},
timeoutHint,
TimeUnit.MILLISECONDS);
timeouts.put(requestHandle, timeout);
}
private CompletableFuture<HRegionLocation> withTimeout(
CompletableFuture<HRegionLocation> future, long timeoutNs, Supplier<String> timeoutMsg) {
if (future.isDone() || timeoutNs <= 0) {
return future;
}
Timeout timeoutTask =
retryTimer.newTimeout(
t -> {
if (future.isDone()) {
return;
}
future.completeExceptionally(new TimeoutIOException(timeoutMsg.get()));
},
timeoutNs,
TimeUnit.NANOSECONDS);
return future.whenComplete(
(loc, error) -> {
if (error != null && error.getClass() != TimeoutIOException.class) {
// cancel timeout task if we are not completed by it.
timeoutTask.cancel();
}
});
}