private synchronized void scheduleUpdate() {
// don't update if the task hasn't been started yet or if it is already finished
if (!needsUpdate.get() || taskInfo.get().getState().isDone()) {
return;
}
// if we have an old request outstanding, cancel it
if (currentRequest != null
&& Duration.nanosSince(currentRequestStartNanos).compareTo(new Duration(2, SECONDS)) >= 0) {
needsUpdate.set(true);
currentRequest.cancel(true);
currentRequest = null;
currentRequestStartNanos = 0;
}
// if there is a request already running, wait for it to complete
if (this.currentRequest != null && !this.currentRequest.isDone()) {
return;
}
// if throttled due to error, asynchronously wait for timeout and try again
ListenableFuture<?> errorRateLimit = updateErrorTracker.acquireRequestPermit();
if (!errorRateLimit.isDone()) {
errorRateLimit.addListener(this::scheduleUpdate, executor);
return;
}
List<TaskSource> sources = getSources();
TaskUpdateRequest updateRequest =
new TaskUpdateRequest(
session.toSessionRepresentation(), planFragment, sources, outputBuffers.get());
Request request =
preparePost()
.setUri(uriBuilderFrom(taskInfo.get().getSelf()).addParameter("summarize").build())
.setHeader(HttpHeaders.CONTENT_TYPE, MediaType.JSON_UTF_8.toString())
.setBodyGenerator(jsonBodyGenerator(taskUpdateRequestCodec, updateRequest))
.build();
updateErrorTracker.startRequest();
ListenableFuture<JsonResponse<TaskInfo>> future =
httpClient.executeAsync(request, createFullJsonResponseHandler(taskInfoCodec));
currentRequest = future;
currentRequestStartNanos = System.nanoTime();
// The needsUpdate flag needs to be set to false BEFORE adding the Future callback since
// callback might change the flag value
// and does so without grabbing the instance lock.
needsUpdate.set(false);
Futures.addCallback(
future,
new SimpleHttpResponseHandler<>(new UpdateResponseHandler(sources), request.getUri()),
executor);
}
private synchronized void sendDelete() {
HttpResponseFuture<StatusResponse> resultFuture =
httpClient.executeAsync(
prepareDelete().setUri(location).build(), createStatusResponseHandler());
future = resultFuture;
Futures.addCallback(
resultFuture,
new FutureCallback<StatusResponse>() {
@Override
public void onSuccess(@Nullable StatusResponse result) {
checkNotHoldsLock();
synchronized (HttpPageBufferClient.this) {
closed = true;
if (future == resultFuture) {
future = null;
errorDelayMillis = 0;
}
lastUpdate = DateTime.now();
}
requestsCompleted.incrementAndGet();
clientCallback.clientFinished(HttpPageBufferClient.this);
}
@Override
public void onFailure(Throwable t) {
checkNotHoldsLock();
log.error("Request to delete %s failed %s", location, t);
Duration errorDuration = elapsedErrorDuration();
if (!(t instanceof PrestoException) && errorDuration.compareTo(minErrorDuration) > 0) {
String message =
format(
"Error closing remote buffer (%s - requests failed for %s)",
location, errorDuration);
t = new PrestoException(REMOTE_BUFFER_CLOSE_FAILED, message, t);
}
handleFailure(t, resultFuture);
}
},
executor);
}
private void scheduleAsyncCleanupRequest(Backoff cleanupBackoff, Request request, String action) {
Futures.addCallback(
httpClient.executeAsync(request, createStatusResponseHandler()),
new FutureCallback<StatusResponse>() {
@Override
public void onSuccess(StatusResponse result) {
// assume any response is good enough
}
@Override
public void onFailure(Throwable t) {
if (t instanceof RejectedExecutionException) {
// client has been shutdown
return;
}
// record failure
if (cleanupBackoff.failure()) {
logError(t, "Unable to %s task at %s", action, request.getUri());
return;
}
// reschedule
long delayNanos = cleanupBackoff.getBackoffDelayNanos();
if (delayNanos == 0) {
scheduleAsyncCleanupRequest(cleanupBackoff, request, action);
} else {
errorScheduledExecutor.schedule(
() -> scheduleAsyncCleanupRequest(cleanupBackoff, request, action),
delayNanos,
NANOSECONDS);
}
}
},
executor);
}
private synchronized void sendGetResults() {
URI uri = HttpUriBuilder.uriBuilderFrom(location).appendPath(String.valueOf(token)).build();
HttpResponseFuture<PagesResponse> resultFuture =
httpClient.executeAsync(
prepareGet().setHeader(PRESTO_MAX_SIZE, maxResponseSize.toString()).setUri(uri).build(),
new PageResponseHandler(blockEncodingSerde));
future = resultFuture;
Futures.addCallback(
resultFuture,
new FutureCallback<PagesResponse>() {
@Override
public void onSuccess(PagesResponse result) {
checkNotHoldsLock();
resetErrors();
List<Page> pages;
try {
synchronized (HttpPageBufferClient.this) {
if (taskInstanceId == null) {
taskInstanceId = result.getTaskInstanceId();
}
if (!isNullOrEmpty(taskInstanceId)
&& !result.getTaskInstanceId().equals(taskInstanceId)) {
// TODO: update error message
throw new PrestoException(REMOTE_TASK_MISMATCH, REMOTE_TASK_MISMATCH_ERROR);
}
if (result.getToken() == token) {
pages = result.getPages();
token = result.getNextToken();
} else {
pages = ImmutableList.of();
}
}
} catch (PrestoException e) {
handleFailure(e, resultFuture);
return;
}
// add pages
if (clientCallback.addPages(HttpPageBufferClient.this, pages)) {
pagesReceived.addAndGet(pages.size());
rowsReceived.addAndGet(pages.stream().mapToLong(Page::getPositionCount).sum());
} else {
pagesRejected.addAndGet(pages.size());
rowsRejected.addAndGet(pages.stream().mapToLong(Page::getPositionCount).sum());
}
synchronized (HttpPageBufferClient.this) {
// client is complete, acknowledge it by sending it a delete in the next request
if (result.isClientComplete()) {
completed = true;
}
if (future == resultFuture) {
future = null;
errorDelayMillis = 0;
}
lastUpdate = DateTime.now();
}
requestsCompleted.incrementAndGet();
clientCallback.requestComplete(HttpPageBufferClient.this);
}
@Override
public void onFailure(Throwable t) {
log.debug("Request to %s failed %s", uri, t);
checkNotHoldsLock();
Duration errorDuration = elapsedErrorDuration();
t = rewriteException(t);
if (!(t instanceof PrestoException) && errorDuration.compareTo(minErrorDuration) > 0) {
String message =
format("%s (%s - requests failed for %s)", WORKER_NODE_ERROR, uri, errorDuration);
t = new PageTransportTimeoutException(message, t);
}
handleFailure(t, resultFuture);
}
},
executor);
}