This is an automated email from the ASF dual-hosted git repository.
cryptoe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new fbc45bebeb2 refactor: Add shared AsyncRequest helper class for async
HTTP requests using AsyncContext (#20784)
fbc45bebeb2 is described below
commit fbc45bebeb279d90a0e2e1097c3fd53186cec9aa
Author: Kashif Faraz <[email protected]>
AuthorDate: Sat Oct 31 03:05:54 2026 +0530
refactor: Add shared AsyncRequest helper class for async HTTP requests
using AsyncContext (#20784)
* refactor: extract shared AsyncRequest helper for async HTTP responses
Moves OverlordResource's private nested AsyncRequest class into
org.apache.druid.server.http (druid-server) as a standalone, reusable
class, and migrates other *Resource classes with similar
AsyncContext-based response handling (SegmentListerResource,
TaskManagementResource, MessageRelayResource) to use it. This removes
duplicated boilerplate for guarding against double-completion of async
requests, writing JSON/Smile responses, and handling async timeouts.
QueryResource and WorkerResource keep their existing AsyncContext usage
since their response-writing logic (raw exception fallback and chunked
octet-stream transfer, respectively) does not fit this abstraction.
* refactor: simplify AsyncRequest timeout handling across callers
All 5 call sites of AsyncRequest.addTimeoutListener() did the same two
things on timeout: cancel a future and complete the request with a
fixed status code (and no body). Replace the generic listener with
completeWithStatusOnTimeout(future, statusCode), which does both,
removing the repeated boilerplate from each caller.
While doing this, fix a latent ordering bug: the request must be
completed before the future is cancelled, not after, since cancelling
a future wired through MoreExecutors.directExecutor() can synchronously
invoke the future's failure callback on the same thread, racing to
complete the response with the wrong status before the timeout status
is written. (Caught by OverlordResourceTest's gateway-timeout test.)
* refactor: fluent builder API for AsyncRequest
Replace AsyncRequest's public complete()/completeWithError()/
completeWithStatus()/completeWithStatusOnTimeout() methods with a
single fluent Builder (AsyncRequest.create(...)) that specifies the
ObjectMapper, content type, the future to wait on, its timeout, and
the status codes/responses to send on timeout, success and failure.
All response-writing logic is now private, reachable only through the
builder's start() (for future-backed responses) and completeNow() (for
immediate ones).
Responses are described with a new AsyncResponse value type
(status-only, result with optional TypeReference/JavaType, or error
with message), keeping ObjectMapper/sendError details out of callers.
Update OverlordResource, SegmentListerResource, TaskManagementResource,
and MessageRelayResource to the new builder API.
Add AsyncRequestTest covering completeNow/start success, failure,
timeout (including the future-completes-after-timeout race), builder
validation, and a regression test for preserving polymorphic type info
on generic/type-erased results. Teach MockHttpServletResponse.sendError
to record the status and message instead of throwing, to support this.
* refactor: collapse AsyncRequest.Builder's future() and timeoutMillis()
into waitOn()
Every caller always set both together, so there was no real use case
for specifying a future without a timeout. Replace the two separate
builder methods with a single waitOn(future, timeoutMillis).
* refactor: remove misleading @Nullable from Builder's
future/onSuccess/onFailure fields
These fields are null only transiently before the corresponding
builder method is called, not valid null values callers should expect
to encounter, so @Nullable was misleading here.
* feat: validate timeoutMillis is positive in AsyncRequest.Builder.waitOn
All existing callers already validate their user-supplied timeout
before calling waitOn(), but guard the method itself against a zero
or negative timeout to make the contract explicit.
* refactor: make AsyncRequest.Builder generic to remove unchecked casts
Split Builder into two stages: Builder (ObjectMapper/content type,
plus completeNow() for immediate responses) and FutureBuilder<T>
(obtained from Builder.waitOn(ListenableFuture<T>, timeoutMillis)),
which carries the future's result type T through onTimeout/onSuccess/
onFailure/start(). This removes the unchecked casts that were
previously needed to store the future and onSuccess/onFailure
functions as a common erased type, and makes completeNow() and
waitOn() mutually exclusive by construction rather than by a runtime
check.
Also reformat the AsyncRequest.onSuccess(...) call sites so the
lambda argument starts on its own line, for consistency with
onFailure.
* refactor: make AsyncResponse generic and drop redundant ObjectMapper param
AsyncResponse is now AsyncResponse<T>, so TypeReference<T> instantiations
at onSuccess call sites can use the diamond operator and have T correctly
inferred, instead of silently erasing to Object via a wildcard parameter.
Also drops the ObjectMapper argument from AsyncResponse.of(...) overloads
since AsyncRequest already carries one from create(), making the
duplicate parameter redundant.
* refactor: collapse AsyncRequest's two builders into one, add completed()
Replaces the separate immediate-completion entry point (Builder.completeNow)
with a standalone AsyncRequest.completed(...) method, and folds the
non-generic Builder + generic FutureBuilder<T> into a single generic
Builder<T> obtained via create(request, mapper, future, timeoutMillis).
This removes a builder stage that no longer served a purpose once
completeNow moved out, and lets create() infer T directly from the future
argument instead of needing a separate waitOn() step.
* refactor: move contentType/ObjectMapper off create() into
Builder.contentType
create() no longer takes an ObjectMapper or content type, since those are
only needed when the request is actually completed, not when it starts
waiting on a future. Builder.contentType(objectMapper, contentType) now
carries both together as a required pair, validated in start().
Also fixes a response content-type regression introduced by extracting
AsyncRequest: SegmentListerResource/TaskManagementResource were passing
the raw Accept header as content type (null when absent), whereas the
original inline AsyncContext code never set a content type at all. Both
resources now resolve a concrete content type alongside their ObjectMapper
in createContext(), matching the original behavior.
* Apply suggestion from @kfaraz
* Defer AsyncResponse construction until request completion is confirmed
AsyncRequest.complete now takes a Supplier<AsyncResponse<T>> instead of an
already-built AsyncResponse, so onSuccess/onFailure/timeout response
construction only happens if the request hasn't already been completed.
---
.../indexing/overlord/http/OverlordResource.java | 103 ++-----
.../worker/http/TaskManagementResource.java | 90 ++----
.../messages/server/MessageRelayResource.java | 93 ++-----
.../org/apache/druid/server/http/AsyncRequest.java | 304 +++++++++++++++++++++
.../druid/server/http/SegmentListerResource.java | 154 +++--------
.../apache/druid/server/http/AsyncRequestTest.java | 286 +++++++++++++++++++
.../server/mocks/MockHttpServletResponse.java | 17 +-
7 files changed, 718 insertions(+), 329 deletions(-)
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java
b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java
index 6c2efddd746..f018d91cdb2 100644
---
a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java
@@ -27,10 +27,7 @@ import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
-import com.google.common.util.concurrent.FutureCallback;
-import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.Inject;
import com.sun.jersey.spi.container.ResourceFilters;
import io.netty.handler.codec.http.HttpResponseStatus;
@@ -66,6 +63,8 @@ import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.UOE;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.metadata.LockFilterPolicy;
+import org.apache.druid.server.http.AsyncRequest;
+import org.apache.druid.server.http.AsyncRequest.AsyncResponse;
import org.apache.druid.server.http.HttpMediaType;
import org.apache.druid.server.http.ServletResourceUtils;
import org.apache.druid.server.http.security.ConfigResourceFilter;
@@ -85,9 +84,7 @@ import org.apache.druid.utils.CollectionUtils;
import org.joda.time.Interval;
import javax.annotation.Nullable;
-import javax.servlet.AsyncContext;
import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
import javax.ws.rs.Consumes;
import javax.ws.rs.DELETE;
import javax.ws.rs.DefaultValue;
@@ -109,7 +106,6 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
/**
*
@@ -523,54 +519,37 @@ public class OverlordResource
@Context HttpServletRequest request
)
{
- final AsyncRequest asyncRequest = new AsyncRequest(request);
final Optional<TaskActionClient> taskActionClient =
taskMaster.getTaskActionClient(holder.getTask());
if (!taskActionClient.isPresent()) {
// Encourage client to try again soon, when we'll likely have a redirect
set up
- asyncRequest.complete(Status.SERVICE_UNAVAILABLE.getStatusCode(), null);
+ AsyncRequest.completed(request, jsonMapper,
AsyncResponse.status(Status.SERVICE_UNAVAILABLE.getStatusCode()));
return;
}
final ListenableFuture<?> future =
taskActionClient.get().submitAsync(holder.getAction());
- asyncRequest.context.addListener(
- ServletResourceUtils.createAsyncTimeoutListener(_ -> {
- asyncRequest.complete(HttpResponseStatus.GATEWAY_TIMEOUT.code(),
null);
- future.cancel(true);
- })
- );
-
// Use a default timeout of 15 minutes
- asyncRequest.context.setTimeout(15 * 60_000);
-
- Futures.addCallback(
- future,
- new FutureCallback<Object>()
- {
- @Override
- public void onSuccess(Object result)
- {
- // Use null-safe map since some task actions may return null result
- final Map<String, Object> payload = new HashMap<>();
- payload.put("result", result);
-
- asyncRequest.complete(Status.OK.getStatusCode(), payload);
- }
-
- @Override
- public void onFailure(Throwable t)
- {
- log.warn(t, "Failed to perform task action");
- if (t instanceof DruidException druidException) {
- asyncRequest.complete(druidException.getStatusCode(),
druidException.toErrorResponse());
- } else {
- final Map<String, String> payload = new HashMap<>();
- payload.put("error", t.getMessage());
-
asyncRequest.complete(Status.INTERNAL_SERVER_ERROR.getStatusCode(), payload);
- }
- }
- },
- MoreExecutors.directExecutor()
- );
+ AsyncRequest.create(request, future, 15 * 60_000)
+ .contentType(jsonMapper, MediaType.APPLICATION_JSON)
+ .onTimeout(HttpResponseStatus.GATEWAY_TIMEOUT.code())
+ .onSuccess(
+ result -> {
+ // Use null-safe map since some task actions may return
null result
+ final Map<String, Object> payload = new HashMap<>();
+ payload.put("result", result);
+ return AsyncResponse.of(Status.OK.getStatusCode(),
payload);
+ }
+ )
+ .onFailure(t -> {
+ log.warn(t, "Failed to perform task action");
+ if (t instanceof DruidException druidException) {
+ return AsyncResponse.of(druidException.getStatusCode(),
druidException.toErrorResponse());
+ } else {
+ final Map<String, String> payload = new HashMap<>();
+ payload.put("error", t.getMessage());
+ return
AsyncResponse.of(Status.INTERNAL_SERVER_ERROR.getStatusCode(), payload);
+ }
+ })
+ .start();
}
@GET
@@ -920,36 +899,4 @@ public class OverlordResource
)
);
}
-
- private class AsyncRequest
- {
- final AtomicBoolean completed = new AtomicBoolean(false);
- final AsyncContext context;
-
- AsyncRequest(HttpServletRequest request)
- {
- this.context = request.startAsync();
- }
-
- void complete(int statusCode, @Nullable Object result)
- {
- if (!completed.compareAndSet(false, true)) {
- log.warn("Request has already been completed");
- return;
- }
-
- try {
- final HttpServletResponse response = (HttpServletResponse)
context.getResponse();
- response.setStatus(statusCode);
- if (result != null) {
- response.setContentType(MediaType.APPLICATION_JSON);
- jsonMapper.writeValue(response.getOutputStream(), result);
- }
- context.complete();
- }
- catch (Exception e) {
- log.noStackTrace().warn(e, "Request timed out or is closed already");
- }
- }
- }
}
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/worker/http/TaskManagementResource.java
b/indexing-service/src/main/java/org/apache/druid/indexing/worker/http/TaskManagementResource.java
index 63d30f06d0a..6d25809a738 100644
---
a/indexing-service/src/main/java/org/apache/druid/indexing/worker/http/TaskManagementResource.java
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/worker/http/TaskManagementResource.java
@@ -21,10 +21,7 @@ package org.apache.druid.indexing.worker.http;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes;
-import com.google.common.util.concurrent.FutureCallback;
-import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.Inject;
import com.sun.jersey.spi.container.ResourceFilters;
import org.apache.druid.guice.annotations.Json;
@@ -36,10 +33,10 @@ import org.apache.druid.indexing.worker.WorkerTaskManager;
import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.server.coordination.ChangeRequestHistory;
import org.apache.druid.server.coordination.ChangeRequestsSnapshot;
-import org.apache.druid.server.http.ServletResourceUtils;
+import org.apache.druid.server.http.AsyncRequest;
+import org.apache.druid.server.http.AsyncRequest.AsyncResponse;
import org.apache.druid.server.http.security.StateResourceFilter;
-import javax.servlet.AsyncContext;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.ws.rs.Consumes;
@@ -51,7 +48,6 @@ import javax.ws.rs.QueryParam;
import javax.ws.rs.core.Context;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
-import java.io.IOException;
/**
* Endpoints used by Overlord to Manage tasks on this Middle Manager.
@@ -100,7 +96,6 @@ public class TaskManagementResource
* @param timeout after which response is sent even if there are no new
segment updates.
* @param req
* @return null to avoid "MUST return a non-void type" warning.
- * @throws IOException
*/
@GET
@Produces({MediaType.APPLICATION_JSON,
SmileMediaTypes.APPLICATION_JACKSON_SMILE})
@@ -109,7 +104,7 @@ public class TaskManagementResource
@QueryParam("hash") long hash,
@QueryParam("timeout") long timeout,
@Context final HttpServletRequest req
- ) throws IOException
+ )
{
if (timeout <= 0) {
sendErrorResponse(req, HttpServletResponse.SC_BAD_REQUEST, "timeout must
be positive.");
@@ -125,56 +120,20 @@ public class TaskManagementResource
)
);
- final AsyncContext asyncContext = req.startAsync();
-
- asyncContext.addListener(
- ServletResourceUtils.createAsyncTimeoutListener(event -> {
- // HTTP 204 NO_CONTENT is sent to the client.
- future.cancel(true);
- event.getAsyncContext().complete();
- })
- );
-
- Futures.addCallback(
- future,
- new FutureCallback<ChangeRequestsSnapshot>()
- {
- @Override
- public void onSuccess(ChangeRequestsSnapshot result)
- {
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.setStatus(HttpServletResponse.SC_OK);
-
context.inputMapper.writerFor(WorkerHolder.WORKER_SYNC_RESP_TYPE_REF)
-
.writeValue(asyncContext.getResponse().getOutputStream(), result);
- asyncContext.complete();
- }
- catch (Exception ex) {
- log.debug(ex, "Request timed out or closed already.");
- }
- }
-
- @Override
- public void onFailure(Throwable th)
- {
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- if (th instanceof IllegalArgumentException) {
- response.sendError(HttpServletResponse.SC_BAD_REQUEST,
th.getMessage());
- } else {
-
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
th.getMessage());
- }
- asyncContext.complete();
- }
- catch (Exception ex) {
- log.debug(ex, "Request timed out or closed already.");
- }
- }
- },
- MoreExecutors.directExecutor()
- );
-
- asyncContext.setTimeout(timeout);
+ AsyncRequest.create(req, future, timeout)
+ .contentType(context.inputMapper, context.contentType)
+ .onTimeout(HttpServletResponse.SC_NO_CONTENT)
+ .onSuccess(
+ result -> AsyncResponse.of(
+ HttpServletResponse.SC_OK,
+ result,
+ WorkerHolder.WORKER_SYNC_RESP_TYPE_REF
+ )
+ )
+ .onFailure(th -> th instanceof IllegalArgumentException
+ ?
AsyncResponse.error(HttpServletResponse.SC_BAD_REQUEST, th.getMessage())
+ :
AsyncResponse.error(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
th.getMessage()))
+ .start();
return null;
}
@@ -198,27 +157,28 @@ public class TaskManagementResource
}
}
- private void sendErrorResponse(HttpServletRequest req, int code, String
error) throws IOException
+ private void sendErrorResponse(HttpServletRequest req, int code, String
error)
{
- AsyncContext asyncContext = req.startAsync();
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.sendError(code, error);
- asyncContext.complete();
+ AsyncRequest.completed(req, jsonMapper, AsyncResponse.error(code, error));
}
private ResponseContext createContext(String requestType)
{
boolean isSmile =
SmileMediaTypes.APPLICATION_JACKSON_SMILE.equals(requestType);
- return new ResponseContext(isSmile ? smileMapper : jsonMapper);
+ return isSmile
+ ? new ResponseContext(smileMapper,
SmileMediaTypes.APPLICATION_JACKSON_SMILE)
+ : new ResponseContext(jsonMapper, MediaType.APPLICATION_JSON);
}
private static class ResponseContext
{
private final ObjectMapper inputMapper;
+ private final String contentType;
- ResponseContext(ObjectMapper inputMapper)
+ ResponseContext(ObjectMapper inputMapper, String contentType)
{
this.inputMapper = inputMapper;
+ this.contentType = contentType;
}
}
}
diff --git
a/server/src/main/java/org/apache/druid/messages/server/MessageRelayResource.java
b/server/src/main/java/org/apache/druid/messages/server/MessageRelayResource.java
index c5ec4f1b333..4c8a397d438 100644
---
a/server/src/main/java/org/apache/druid/messages/server/MessageRelayResource.java
+++
b/server/src/main/java/org/apache/druid/messages/server/MessageRelayResource.java
@@ -22,17 +22,14 @@ package org.apache.druid.messages.server;
import com.fasterxml.jackson.databind.JavaType;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes;
-import com.google.common.util.concurrent.FutureCallback;
-import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.messages.MessageBatch;
import org.apache.druid.messages.client.MessageListener;
import org.apache.druid.messages.client.MessageRelayClient;
-import org.apache.druid.server.http.ServletResourceUtils;
+import org.apache.druid.server.http.AsyncRequest;
+import org.apache.druid.server.http.AsyncRequest.AsyncResponse;
-import javax.servlet.AsyncContext;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.ws.rs.GET;
@@ -40,8 +37,6 @@ import javax.ws.rs.Path;
import javax.ws.rs.PathParam;
import javax.ws.rs.QueryParam;
import javax.ws.rs.core.Context;
-import java.io.IOException;
-import java.util.concurrent.atomic.AtomicBoolean;
/**
* Server-side resource for message relaying. Wraps an {@link Outbox} and
{@link MessageListener}.
@@ -82,7 +77,7 @@ public class MessageRelayResource<MessageType>
* Retrieve messages from the outbox for a particular client, as a {@link
MessageBatch} in Smile format.
* The messages are retrieved from {@link Outbox#getMessages(String, long,
long)}.
*
- * This is a long-polling async method, using {@link AsyncContext} to wait
up to {@link #GET_MESSAGES_TIMEOUT} for
+ * This is a long-polling async method, using {@link AsyncRequest} to wait
up to {@link #GET_MESSAGES_TIMEOUT} for
* messages to appear in the outbox.
*
* @return HTTP 200 with Smile response with messages on success; HTTP 204
(No Content) if no messages were put in
@@ -97,78 +92,38 @@ public class MessageRelayResource<MessageType>
@QueryParam("epoch") final Long epoch,
@QueryParam("watermark") final Long watermark,
@Context final HttpServletRequest req
- ) throws IOException
+ )
{
if (epoch == null || watermark == null || clientHost == null ||
clientHost.isEmpty()) {
- AsyncContext asyncContext = req.startAsync();
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.sendError(HttpServletResponse.SC_BAD_REQUEST);
- asyncContext.complete();
+ AsyncRequest.completed(
+ req,
+ smileMapper,
+ SmileMediaTypes.APPLICATION_JACKSON_SMILE,
+ AsyncResponse.error(HttpServletResponse.SC_BAD_REQUEST, null)
+ );
return null;
}
- final AtomicBoolean didRespond = new AtomicBoolean();
final ListenableFuture<MessageBatch<MessageType>> batchFuture =
outbox.getMessages(clientHost, epoch, watermark);
- final AsyncContext asyncContext = req.startAsync();
- asyncContext.setTimeout(GET_MESSAGES_TIMEOUT);
- asyncContext.addListener(
- ServletResourceUtils.createAsyncTimeoutListener(event -> {
- if (didRespond.compareAndSet(false, true)) {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.setStatus(HttpServletResponse.SC_NO_CONTENT);
- event.getAsyncContext().complete();
- batchFuture.cancel(true);
- }
- })
- );
// Save these items, since "req" becomes inaccessible in future exception
handlers.
final String remoteAddr = req.getRemoteAddr();
final String requestURI = req.getRequestURI();
- Futures.addCallback(
- batchFuture,
- new FutureCallback<>()
- {
- @Override
- public void onSuccess(MessageBatch<MessageType> result)
- {
- if (didRespond.compareAndSet(false, true)) {
- log.debug("Sending message batch: %s", result);
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.setStatus(HttpServletResponse.SC_OK);
-
response.setContentType(SmileMediaTypes.APPLICATION_JACKSON_SMILE);
- smileMapper.writerFor(batchType)
-
.writeValue(asyncContext.getResponse().getOutputStream(), result);
- response.getOutputStream().close();
- asyncContext.complete();
- }
- catch (Exception e) {
- log.noStackTrace().warn(e, "Could not respond to request
from[%s] to[%s]", remoteAddr, requestURI);
- }
- }
- }
-
- @Override
- public void onFailure(Throwable e)
- {
- if (didRespond.compareAndSet(false, true)) {
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
-
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
- asyncContext.complete();
- }
- catch (Exception e2) {
- e.addSuppressed(e2);
- }
-
- log.noStackTrace().warn(e, "Request failed from[%s] to[%s]",
remoteAddr, requestURI);
- }
- }
- },
- Execs.directExecutor()
- );
+ AsyncRequest.create(req, batchFuture, GET_MESSAGES_TIMEOUT)
+ .contentType(smileMapper,
SmileMediaTypes.APPLICATION_JACKSON_SMILE)
+ .onTimeout(HttpServletResponse.SC_NO_CONTENT)
+ .onSuccess(
+ result -> {
+ log.debug("Sending message batch: %s", result);
+ return AsyncResponse.of(HttpServletResponse.SC_OK,
result, batchType);
+ }
+ )
+ .onFailure(e -> {
+ log.noStackTrace().warn(e, "Request failed from[%s] to[%s]",
remoteAddr, requestURI);
+ return
AsyncResponse.error(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, null);
+ })
+ .start();
return null;
}
diff --git
a/server/src/main/java/org/apache/druid/server/http/AsyncRequest.java
b/server/src/main/java/org/apache/druid/server/http/AsyncRequest.java
new file mode 100644
index 00000000000..04fd558278a
--- /dev/null
+++ b/server/src/main/java/org/apache/druid/server/http/AsyncRequest.java
@@ -0,0 +1,304 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.server.http;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.JavaType;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.ObjectWriter;
+import com.google.common.base.Preconditions;
+import com.google.common.util.concurrent.FutureCallback;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
+import org.apache.druid.java.util.common.logger.Logger;
+
+import javax.annotation.Nullable;
+import javax.servlet.AsyncContext;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import javax.ws.rs.core.MediaType;
+import java.lang.reflect.Type;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * Wraps the {@link AsyncContext} of a {@link HttpServletRequest} and ensures
that the
+ * response is completed at most once, guarding against writing to a timed-out
or already-completed request.
+ * Use {@link #completed(HttpServletRequest, ObjectMapper, AsyncResponse)} to
complete a request immediately, or
+ * {@link #create(HttpServletRequest, ListenableFuture, long)} to obtain a
{@link Builder} that responds based on
+ * the outcome of the future; call {@link Builder#contentType} to specify how
to write the response entity.
+ */
+public class AsyncRequest
+{
+ private static final Logger log = new Logger(AsyncRequest.class);
+
+ private final AtomicBoolean completed = new AtomicBoolean(false);
+ private final AsyncContext context;
+ private final ObjectMapper objectMapper;
+ private final String contentType;
+
+ /**
+ * Starts async processing of the request and immediately completes it with
the given response, writing any
+ * response entity as JSON using the given {@link ObjectMapper}.
+ */
+ public static <T> AsyncRequest completed(HttpServletRequest request,
ObjectMapper jsonMapper, AsyncResponse<T> response)
+ {
+ return completed(request, jsonMapper, MediaType.APPLICATION_JSON,
response);
+ }
+
+ /**
+ * Starts async processing of the request and immediately completes it with
the given response, writing any
+ * response entity using the given {@link ObjectMapper} and content type.
+ */
+ public static <T> AsyncRequest completed(HttpServletRequest request,
ObjectMapper objectMapper, String contentType, AsyncResponse<T> response)
+ {
+ final AsyncRequest asyncRequest = new AsyncRequest(request, objectMapper,
contentType);
+ asyncRequest.complete(() -> response);
+ return asyncRequest;
+ }
+
+ /**
+ * Creates a {@link Builder} that waits on the given future. Call {@link
Builder#contentType} to specify the
+ * {@link ObjectMapper} and content type to use for writing the response
entity.
+ */
+ public static <T> Builder<T> create(HttpServletRequest request,
ListenableFuture<T> future, long timeoutMillis)
+ {
+ Preconditions.checkArgument(timeoutMillis > 0, "timeoutMillis[%s] must be
greater than zero", timeoutMillis);
+ return new Builder<>(request, future, timeoutMillis);
+ }
+
+ private AsyncRequest(HttpServletRequest request, ObjectMapper objectMapper,
String contentType)
+ {
+ this.context = request.startAsync();
+ this.objectMapper = objectMapper;
+ this.contentType = contentType;
+ }
+
+ /**
+ * Completes the request as described by the {@link AsyncResponse} obtained
from the given supplier, if it
+ * hasn't already been completed. The supplier is invoked only if the
request is not already complete.
+ */
+ private <T> void complete(Supplier<AsyncResponse<T>> responseSupplier)
+ {
+ if (!completed.compareAndSet(false, true)) {
+ // Request has already been completed
+ return;
+ }
+
+ try {
+ final AsyncResponse<T> response = responseSupplier.get();
+ final HttpServletResponse servletResponse = (HttpServletResponse)
context.getResponse();
+ if (response.isError) {
+ if (response.errorMessage == null) {
+ servletResponse.sendError(response.statusCode);
+ } else {
+ servletResponse.sendError(response.statusCode,
response.errorMessage);
+ }
+ } else {
+ servletResponse.setStatus(response.statusCode);
+ if (response.result != null) {
+ servletResponse.setContentType(contentType);
+ final ObjectWriter writer = response.resultType == null
+ ? objectMapper.writer()
+ :
objectMapper.writerFor(objectMapper.getTypeFactory().constructType(response.resultType));
+ writer.writeValue(servletResponse.getOutputStream(),
response.result);
+ }
+ }
+ context.complete();
+ }
+ catch (Exception e) {
+ log.noStackTrace().warn(e, "Request timed out or is closed already");
+ }
+ }
+
+ /**
+ * Describes how an {@link AsyncRequest} should be completed: a status code,
with either a result of type
+ * {@code T} (serialized using the {@link AsyncRequest}'s {@link
ObjectMapper}) or an error message (sent using
+ * {@link HttpServletResponse#sendError}).
+ */
+ public static final class AsyncResponse<T>
+ {
+ private final int statusCode;
+ private final boolean isError;
+ @Nullable
+ private final T result;
+ @Nullable
+ private final Type resultType;
+ @Nullable
+ private final String errorMessage;
+
+ private AsyncResponse(int statusCode, boolean isError, @Nullable T result,
@Nullable Type resultType, @Nullable String errorMessage)
+ {
+ this.statusCode = statusCode;
+ this.isError = isError;
+ this.result = result;
+ this.resultType = resultType;
+ this.errorMessage = errorMessage;
+ }
+
+ /**
+ * Responds with the given status code and no body.
+ */
+ public static AsyncResponse<Void> status(int statusCode)
+ {
+ return new AsyncResponse<>(statusCode, false, null, null, null);
+ }
+
+ /**
+ * Responds with the given status code, writing the given result, if
non-null.
+ */
+ public static <T> AsyncResponse<T> of(int statusCode, @Nullable T result)
+ {
+ return new AsyncResponse<>(statusCode, false, result, null, null);
+ }
+
+ /**
+ * Responds with the given status code, writing the given result using the
given {@link TypeReference} so
+ * that generic type information erased at runtime is preserved.
+ */
+ public static <T> AsyncResponse<T> of(int statusCode, @Nullable T result,
TypeReference<T> resultType)
+ {
+ return new AsyncResponse<>(statusCode, false, result,
resultType.getType(), null);
+ }
+
+ /**
+ * Responds with the given status code, writing the given result using the
given {@link JavaType} so that
+ * generic type information erased at runtime is preserved.
+ */
+ public static <T> AsyncResponse<T> of(int statusCode, @Nullable T result,
@Nullable JavaType resultType)
+ {
+ return new AsyncResponse<>(statusCode, false, result, resultType, null);
+ }
+
+ /**
+ * Responds with the given error status code and message, using {@link
HttpServletResponse#sendError}.
+ */
+ public static AsyncResponse<Void> error(int statusCode, @Nullable String
message)
+ {
+ return new AsyncResponse<>(statusCode, true, null, null, message);
+ }
+ }
+
+ /**
+ * Builds and starts an {@link AsyncRequest} that responds based on the
outcome of a {@link ListenableFuture}
+ * of type {@code T}, obtained from {@link #create(HttpServletRequest,
ListenableFuture, long)}.
+ */
+ public static final class Builder<T>
+ {
+ private final HttpServletRequest request;
+ private final ListenableFuture<T> future;
+ private final long timeoutMillis;
+
+ private ObjectMapper objectMapper;
+ private String contentType;
+ private int onTimeoutStatusCode = HttpServletResponse.SC_NO_CONTENT;
+ private Function<T, AsyncResponse<?>> onSuccess;
+ private Function<Throwable, AsyncResponse<?>> onFailure;
+
+ private Builder(HttpServletRequest request, ListenableFuture<T> future,
long timeoutMillis)
+ {
+ this.request = request;
+ this.future = future;
+ this.timeoutMillis = timeoutMillis;
+ }
+
+ /**
+ * The {@link ObjectMapper} and content type to use for writing the
response entity. Required.
+ */
+ public Builder<T> contentType(ObjectMapper objectMapper, String
contentType)
+ {
+ this.objectMapper = objectMapper;
+ this.contentType = contentType;
+ return this;
+ }
+
+ /**
+ * Status code to respond with, with no body, if the future times out
before completing.
+ */
+ public Builder<T> onTimeout(int statusCode)
+ {
+ this.onTimeoutStatusCode = statusCode;
+ return this;
+ }
+
+ /**
+ * Function invoked with the result of the future on success, to determine
how to respond. Required.
+ */
+ public Builder<T> onSuccess(Function<T, AsyncResponse<?>> onSuccess)
+ {
+ this.onSuccess = onSuccess;
+ return this;
+ }
+
+ /**
+ * Function invoked with the failure cause of the future, to determine how
to respond. Required.
+ */
+ public Builder<T> onFailure(Function<Throwable, AsyncResponse<?>>
onFailure)
+ {
+ this.onFailure = onFailure;
+ return this;
+ }
+
+ /**
+ * Starts async processing of the request using {@link
HttpServletRequest#startAsync()}, wiring up the
+ * configured future, timeout and success/failure handling.
+ */
+ public AsyncRequest start()
+ {
+ Preconditions.checkArgument(objectMapper != null, "contentType must be
specified");
+ Preconditions.checkArgument(contentType != null, "contentType must be
specified");
+ Preconditions.checkArgument(onSuccess != null, "onSuccess must be
specified");
+ Preconditions.checkArgument(onFailure != null, "onFailure must be
specified");
+
+ final AsyncRequest asyncRequest = new AsyncRequest(request,
objectMapper, contentType);
+
+ asyncRequest.context.setTimeout(timeoutMillis);
+ asyncRequest.context.addListener(
+ ServletResourceUtils.createAsyncTimeoutListener(event -> {
+ asyncRequest.complete(() ->
AsyncResponse.status(onTimeoutStatusCode));
+ future.cancel(true);
+ })
+ );
+
+ Futures.addCallback(
+ future,
+ new FutureCallback<T>()
+ {
+ @Override
+ public void onSuccess(T result)
+ {
+ asyncRequest.complete(() -> onSuccess.apply(result));
+ }
+
+ @Override
+ public void onFailure(Throwable t)
+ {
+ asyncRequest.complete(() -> onFailure.apply(t));
+ }
+ },
+ MoreExecutors.directExecutor()
+ );
+
+ return asyncRequest;
+ }
+ }
+}
diff --git
a/server/src/main/java/org/apache/druid/server/http/SegmentListerResource.java
b/server/src/main/java/org/apache/druid/server/http/SegmentListerResource.java
index 20c18942ad4..17cde9059e8 100644
---
a/server/src/main/java/org/apache/druid/server/http/SegmentListerResource.java
+++
b/server/src/main/java/org/apache/druid/server/http/SegmentListerResource.java
@@ -21,10 +21,7 @@ package org.apache.druid.server.http;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes;
-import com.google.common.util.concurrent.FutureCallback;
-import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.Inject;
import com.sun.jersey.spi.container.ResourceFilters;
import org.apache.druid.client.HttpServerInventoryView;
@@ -39,10 +36,10 @@ import
org.apache.druid.server.coordination.DataSegmentChangeRequest;
import org.apache.druid.server.coordination.DataSegmentChangeResponse;
import org.apache.druid.server.coordination.SegmentLoadDropHandler;
import org.apache.druid.server.coordinator.loading.HttpLoadQueuePeon;
+import org.apache.druid.server.http.AsyncRequest.AsyncResponse;
import org.apache.druid.server.http.security.StateResourceFilter;
import javax.annotation.Nullable;
-import javax.servlet.AsyncContext;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.ws.rs.Consumes;
@@ -54,7 +51,6 @@ import javax.ws.rs.QueryParam;
import javax.ws.rs.core.Context;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
-import java.io.IOException;
import java.util.List;
/**
@@ -111,7 +107,6 @@ public class SegmentListerResource
* @param timeout after which response is sent even if there are no new
segment updates.
* @param req
* @return null to avoid "MUST return a non-void type" warning.
- * @throws IOException
*/
@GET
@Produces({MediaType.APPLICATION_JSON,
SmileMediaTypes.APPLICATION_JACKSON_SMILE})
@@ -121,7 +116,7 @@ public class SegmentListerResource
@QueryParam("hash") long hash,
@QueryParam("timeout") long timeout,
@Context final HttpServletRequest req
- ) throws IOException
+ )
{
if (announcer == null) {
sendErrorResponse(req, HttpServletResponse.SC_NOT_FOUND, "announcer is
not available.");
@@ -141,56 +136,20 @@ public class SegmentListerResource
)
);
- final AsyncContext asyncContext = req.startAsync();
-
- asyncContext.addListener(
- ServletResourceUtils.createAsyncTimeoutListener(event -> {
- // HTTP 204 NO_CONTENT is sent to the client.
- future.cancel(true);
- event.getAsyncContext().complete();
- })
- );
-
- Futures.addCallback(
- future,
- new FutureCallback<>()
- {
- @Override
- public void
onSuccess(ChangeRequestsSnapshot<DataSegmentChangeRequest> result)
- {
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.setStatus(HttpServletResponse.SC_OK);
-
context.inputMapper.writerFor(HttpServerInventoryView.SEGMENT_LIST_RESP_TYPE_REF)
-
.writeValue(asyncContext.getResponse().getOutputStream(), result);
- asyncContext.complete();
- }
- catch (Exception ex) {
- log.debug(ex, "Request timed out or closed already.");
- }
- }
-
- @Override
- public void onFailure(Throwable th)
- {
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- if (th instanceof IllegalArgumentException) {
- response.sendError(HttpServletResponse.SC_BAD_REQUEST,
th.getMessage());
- } else {
-
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
th.getMessage());
- }
- asyncContext.complete();
- }
- catch (Exception ex) {
- log.debug(ex, "Request timed out or closed already.");
- }
- }
- },
- MoreExecutors.directExecutor()
- );
-
- asyncContext.setTimeout(timeout);
+ AsyncRequest.create(req, future, timeout)
+ .contentType(context.inputMapper, context.contentType)
+ .onTimeout(HttpServletResponse.SC_NO_CONTENT)
+ .onSuccess(
+ result -> AsyncResponse.of(
+ HttpServletResponse.SC_OK,
+ result,
+ HttpServerInventoryView.SEGMENT_LIST_RESP_TYPE_REF
+ )
+ )
+ .onFailure(th -> th instanceof IllegalArgumentException
+ ?
AsyncResponse.error(HttpServletResponse.SC_BAD_REQUEST, th.getMessage())
+ :
AsyncResponse.error(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
th.getMessage()))
+ .start();
return null;
}
@@ -213,7 +172,7 @@ public class SegmentListerResource
@QueryParam("loadingMode") @Nullable SegmentLoadingMode loadingMode,
List<DataSegmentChangeRequest> changeRequestList,
@Context final HttpServletRequest req
- ) throws IOException
+ )
{
if (loadDropRequestHandler == null) {
sendErrorResponse(req, HttpServletResponse.SC_NOT_FOUND, "load/drop
handler is not available.");
@@ -234,56 +193,20 @@ public class SegmentListerResource
final ListenableFuture<List<DataSegmentChangeResponse>> future =
loadDropRequestHandler.processBatch(changeRequestList, loadingMode ==
null ? SegmentLoadingMode.NORMAL : loadingMode);
- final AsyncContext asyncContext = req.startAsync();
-
- asyncContext.addListener(
- ServletResourceUtils.createAsyncTimeoutListener(event -> {
- // HTTP 204 NO_CONTENT is sent to the client.
- future.cancel(true);
- event.getAsyncContext().complete();
- })
- );
-
- Futures.addCallback(
- future,
- new FutureCallback<>()
- {
- @Override
- public void onSuccess(List<DataSegmentChangeResponse> result)
- {
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.setStatus(HttpServletResponse.SC_OK);
-
context.inputMapper.writerFor(HttpLoadQueuePeon.RESPONSE_ENTITY_TYPE_REF)
-
.writeValue(asyncContext.getResponse().getOutputStream(), result);
- asyncContext.complete();
- }
- catch (Exception ex) {
- log.debug(ex, "Request timed out or closed already.");
- }
- }
-
- @Override
- public void onFailure(Throwable th)
- {
- try {
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- if (th instanceof IllegalArgumentException) {
- response.sendError(HttpServletResponse.SC_BAD_REQUEST,
th.getMessage());
- } else {
-
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
th.getMessage());
- }
- asyncContext.complete();
- }
- catch (Exception ex) {
- log.debug(ex, "Request timed out or closed already.");
- }
- }
- },
- MoreExecutors.directExecutor()
- );
-
- asyncContext.setTimeout(timeout);
+ AsyncRequest.create(req, future, timeout)
+ .contentType(context.inputMapper, context.contentType)
+ .onTimeout(HttpServletResponse.SC_NO_CONTENT)
+ .onSuccess(
+ result -> AsyncResponse.of(
+ HttpServletResponse.SC_OK,
+ result,
+ HttpLoadQueuePeon.RESPONSE_ENTITY_TYPE_REF
+ )
+ )
+ .onFailure(th -> th instanceof IllegalArgumentException
+ ?
AsyncResponse.error(HttpServletResponse.SC_BAD_REQUEST, th.getMessage())
+ :
AsyncResponse.error(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
th.getMessage()))
+ .start();
}
@GET
@@ -304,27 +227,28 @@ public class SegmentListerResource
return
Response.status(Response.Status.OK).entity(capabilitiesResponse).build();
}
- private void sendErrorResponse(HttpServletRequest req, int code, String
error) throws IOException
+ private void sendErrorResponse(HttpServletRequest req, int code, String
error)
{
- AsyncContext asyncContext = req.startAsync();
- HttpServletResponse response = (HttpServletResponse)
asyncContext.getResponse();
- response.sendError(code, error);
- asyncContext.complete();
+ AsyncRequest.completed(req, jsonMapper, AsyncResponse.error(code, error));
}
private ResponseContext createContext(String requestType)
{
boolean isSmile =
SmileMediaTypes.APPLICATION_JACKSON_SMILE.equals(requestType);
- return new ResponseContext(isSmile ? smileMapper : jsonMapper);
+ return isSmile
+ ? new ResponseContext(smileMapper,
SmileMediaTypes.APPLICATION_JACKSON_SMILE)
+ : new ResponseContext(jsonMapper, MediaType.APPLICATION_JSON);
}
private static class ResponseContext
{
private final ObjectMapper inputMapper;
+ private final String contentType;
- ResponseContext(ObjectMapper inputMapper)
+ ResponseContext(ObjectMapper inputMapper, String contentType)
{
this.inputMapper = inputMapper;
+ this.contentType = contentType;
}
}
}
diff --git
a/server/src/test/java/org/apache/druid/server/http/AsyncRequestTest.java
b/server/src/test/java/org/apache/druid/server/http/AsyncRequestTest.java
new file mode 100644
index 00000000000..c5fd3afd317
--- /dev/null
+++ b/server/src/test/java/org/apache/druid/server/http/AsyncRequestTest.java
@@ -0,0 +1,286 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.server.http;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.SettableFuture;
+import org.apache.druid.segment.TestHelper;
+import org.apache.druid.server.coordination.ChangeRequestHistory;
+import org.apache.druid.server.coordination.ChangeRequestsSnapshot;
+import org.apache.druid.server.coordination.DataSegmentChangeRequest;
+import org.apache.druid.server.coordination.SegmentChangeRequestNoop;
+import org.apache.druid.server.http.AsyncRequest.AsyncResponse;
+import org.apache.druid.server.mocks.MockAsyncContext;
+import org.apache.druid.server.mocks.MockHttpServletRequest;
+import org.apache.druid.server.mocks.MockHttpServletResponse;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import javax.servlet.AsyncEvent;
+import javax.servlet.AsyncListener;
+import javax.ws.rs.core.MediaType;
+import java.util.List;
+import java.util.Map;
+
+public class AsyncRequestTest
+{
+ private static final ObjectMapper JSON_MAPPER = TestHelper.JSON_MAPPER;
+
+ private MockHttpServletResponse response;
+ private MockHttpServletRequest request;
+ private CapturingAsyncContext asyncContext;
+
+ @BeforeEach
+ public void setup()
+ {
+ response = new MockHttpServletResponse();
+ request = new MockHttpServletRequest();
+ asyncContext = new CapturingAsyncContext();
+ asyncContext.request = request;
+ asyncContext.response = response;
+ request.asyncContextSupplier = () -> asyncContext;
+ }
+
+ @Test
+ public void test_completed_writesStatusAndBody_whenResponseHasResult()
+ {
+ AsyncRequest.completed(request, JSON_MAPPER, AsyncResponse.of(200,
Map.of("hello", "world")));
+
+ Assertions.assertEquals(200, response.getStatus());
+ Assertions.assertEquals(Map.of("hello", "world"), readJson(new
TypeReference<Map<String, String>>() {}));
+ Assertions.assertTrue(asyncContext.isCompleted());
+ }
+
+ @Test
+ public void test_completed_writesStatusWithoutBody_whenResponseHasNoResult()
+ {
+ AsyncRequest.completed(request, JSON_MAPPER, AsyncResponse.status(204));
+
+ Assertions.assertEquals(204, response.getStatus());
+ Assertions.assertEquals(0, response.baos.size());
+ Assertions.assertTrue(asyncContext.isCompleted());
+ }
+
+ @Test
+ public void test_completed_sendsErrorWithMessage_whenResponseIsError()
+ {
+ AsyncRequest.completed(request, JSON_MAPPER, AsyncResponse.error(400, "bad
request"));
+
+ Assertions.assertEquals(400, response.getStatus());
+ Assertions.assertEquals("bad request", response.getErrorMessage());
+ Assertions.assertTrue(asyncContext.isCompleted());
+ }
+
+ @Test
+ public void test_completed_sendsErrorWithoutMessage_whenErrorMessageIsNull()
+ {
+ AsyncRequest.completed(request, JSON_MAPPER, AsyncResponse.error(500,
null));
+
+ Assertions.assertEquals(500, response.getStatus());
+ Assertions.assertNull(response.getErrorMessage());
+ Assertions.assertTrue(asyncContext.isCompleted());
+ }
+
+ @Test
+ public void test_start_invokesOnSuccess_whenFutureSucceeds()
+ {
+ AsyncRequest.create(request, Futures.immediateFuture("the-result"), 10L)
+ .contentType(JSON_MAPPER, MediaType.APPLICATION_JSON)
+ .onSuccess(result -> AsyncResponse.of(200, result))
+ .onFailure(t -> AsyncResponse.error(500, t.getMessage()))
+ .start();
+
+ Assertions.assertEquals(200, response.getStatus());
+ Assertions.assertEquals("the-result", readJson(new TypeReference<String>()
{}));
+ Assertions.assertTrue(asyncContext.isCompleted());
+ }
+
+ @Test
+ public void test_start_invokesOnFailure_whenFutureFails()
+ {
+ final RuntimeException cause = new RuntimeException("boom");
+
+ AsyncRequest.create(request, Futures.immediateFailedFuture(cause), 10L)
+ .contentType(JSON_MAPPER, MediaType.APPLICATION_JSON)
+ .onSuccess(result -> AsyncResponse.of(200, result))
+ .onFailure(t -> AsyncResponse.error(500, t.getMessage()))
+ .start();
+
+ Assertions.assertEquals(500, response.getStatus());
+ Assertions.assertEquals("boom", response.getErrorMessage());
+ Assertions.assertTrue(asyncContext.isCompleted());
+ }
+
+ @Test
+ public void test_create_throwsException_whenTimeoutMillisIsZero()
+ {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> AsyncRequest.create(request, Futures.immediateVoidFuture(), 0L)
+ );
+ }
+
+ @Test
+ public void test_create_throwsException_whenTimeoutMillisIsNegative()
+ {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> AsyncRequest.create(request, Futures.immediateVoidFuture(), -1L)
+ );
+ }
+
+ @Test
+ public void test_start_throwsException_whenContentTypeIsMissing()
+ {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> AsyncRequest.create(request, Futures.immediateVoidFuture(), 10L)
+ .onSuccess(result -> AsyncResponse.status(200))
+ .onFailure(t -> AsyncResponse.error(500,
t.getMessage()))
+ .start()
+ );
+ }
+
+ @Test
+ public void test_start_throwsException_whenOnSuccessIsMissing()
+ {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> AsyncRequest.create(request, Futures.immediateVoidFuture(), 10L)
+ .contentType(JSON_MAPPER,
MediaType.APPLICATION_JSON)
+ .onFailure(t -> AsyncResponse.error(500,
t.getMessage()))
+ .start()
+ );
+ }
+
+ @Test
+ public void test_start_throwsException_whenOnFailureIsMissing()
+ {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> AsyncRequest.create(request, Futures.immediateVoidFuture(), 10L)
+ .contentType(JSON_MAPPER,
MediaType.APPLICATION_JSON)
+ .onSuccess(result -> AsyncResponse.status(200))
+ .start()
+ );
+ }
+
+ @Test
+ public void
test_start_completesWithTimeoutStatusAndCancelsFuture_whenFutureTimesOut()
throws Exception
+ {
+ final SettableFuture<String> future = SettableFuture.create();
+
+ AsyncRequest.create(request, future, 10L)
+ .contentType(JSON_MAPPER, MediaType.APPLICATION_JSON)
+ .onTimeout(204)
+ .onSuccess(result -> AsyncResponse.of(200, result))
+ .onFailure(t -> AsyncResponse.error(500, t.getMessage()))
+ .start();
+
+ Assertions.assertEquals(10L, asyncContext.getTimeout());
+ Assertions.assertNotNull(asyncContext.capturedListener);
+
+ asyncContext.capturedListener.onTimeout(new AsyncEvent(asyncContext));
+
+ Assertions.assertEquals(204, response.getStatus());
+ Assertions.assertTrue(asyncContext.isCompleted());
+ Assertions.assertTrue(future.isCancelled());
+ }
+
+ @Test
+ public void
test_start_doesNotOverwriteResponse_whenFutureCompletesAfterTimeout() throws
Exception
+ {
+ final SettableFuture<String> future = SettableFuture.create();
+
+ AsyncRequest.create(request, future, 10L)
+ .contentType(JSON_MAPPER, MediaType.APPLICATION_JSON)
+ .onTimeout(204)
+ .onSuccess(result -> AsyncResponse.of(200, result))
+ .onFailure(t -> AsyncResponse.error(500, t.getMessage()))
+ .start();
+
+ asyncContext.capturedListener.onTimeout(new AsyncEvent(asyncContext));
+ Assertions.assertEquals(204, response.getStatus());
+
+ // The future is cancelled by the timeout, but completing it again (as a
real future might race to do)
+ // must not be allowed to overwrite the response that was already sent.
+ future.set("too-late");
+
+ Assertions.assertEquals(204, response.getStatus());
+ Assertions.assertEquals(0, response.baos.size());
+ }
+
+ @Test
+ public void
test_onSuccess_preservesPolymorphicTypeInfo_whenResultTypeIsGenericAndErased()
+ {
+ final ChangeRequestsSnapshot<DataSegmentChangeRequest> snapshot =
ChangeRequestsSnapshot.success(
+ new ChangeRequestHistory.Counter(1, 1),
+ List.of(new SegmentChangeRequestNoop())
+ );
+
+ AsyncRequest.create(
+ request,
+
Futures.<ChangeRequestsSnapshot<DataSegmentChangeRequest>>immediateFuture(snapshot),
+ 10L
+ )
+ .contentType(JSON_MAPPER, MediaType.APPLICATION_JSON)
+ .onSuccess(
+ result -> AsyncResponse.of(200, result, new
TypeReference<>() {})
+ )
+ .onFailure(t -> AsyncResponse.error(500, t.getMessage()))
+ .start();
+
+ Assertions.assertEquals(200, response.getStatus());
+
+ final ChangeRequestsSnapshot<DataSegmentChangeRequest> readBack = readJson(
+ new TypeReference<ChangeRequestsSnapshot<DataSegmentChangeRequest>>()
{}
+ );
+ Assertions.assertEquals(1, readBack.getRequests().size());
+ Assertions.assertTrue(readBack.getRequests().get(0) instanceof
SegmentChangeRequestNoop);
+ }
+
+ private <T> T readJson(TypeReference<T> type)
+ {
+ try {
+ return JSON_MAPPER.readValue(response.baos.toByteArray(), type);
+ }
+ catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ /**
+ * A {@link MockAsyncContext} that captures the single listener added to it,
so tests can manually fire a
+ * timeout event.
+ */
+ private static class CapturingAsyncContext extends MockAsyncContext
+ {
+ private AsyncListener capturedListener;
+
+ @Override
+ public void addListener(AsyncListener listener)
+ {
+ capturedListener = listener;
+ }
+ }
+}
diff --git
a/server/src/test/java/org/apache/druid/server/mocks/MockHttpServletResponse.java
b/server/src/test/java/org/apache/druid/server/mocks/MockHttpServletResponse.java
index 0069dba5e71..3994a00e8e2 100644
---
a/server/src/test/java/org/apache/druid/server/mocks/MockHttpServletResponse.java
+++
b/server/src/test/java/org/apache/druid/server/mocks/MockHttpServletResponse.java
@@ -68,6 +68,8 @@ public class MockHttpServletResponse implements
HttpServletResponse
private int statusCode;
private String contentType;
private String characterEncoding;
+ @Nullable
+ private String errorMessage;
@Override
public void reset()
@@ -122,13 +124,24 @@ public class MockHttpServletResponse implements
HttpServletResponse
@Override
public void sendError(int sc, String msg)
{
- throw new UnsupportedOperationException();
+ statusCode = sc;
+ errorMessage = msg;
}
@Override
public void sendError(int sc)
{
- throw new UnsupportedOperationException();
+ sendError(sc, null);
+ }
+
+ /**
+ * The message passed to the most recent {@link #sendError} call, or null if
{@link #sendError} hasn't been
+ * called or was called without a message.
+ */
+ @Nullable
+ public String getErrorMessage()
+ {
+ return errorMessage;
}
@Override
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
