Use the max response size provided by the HttpPageBufferClient

gerashegalov · Oct 13, 2015 · d3459a1 · d3459a1
1 parent de7a27d
commit d3459a1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/presto-main/src/main/java/com/facebook/presto/server/TaskResource.java b/presto-main/src/main/java/com/facebook/presto/server/TaskResource.java
@@ -56,6 +56,7 @@
 import static com.facebook.presto.PrestoMediaTypes.PRESTO_PAGES;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_BUFFER_COMPLETE;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_CURRENT_STATE;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_MAX_SIZE;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_MAX_WAIT;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_PAGE_NEXT_TOKEN;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_PAGE_TOKEN;
@@ -190,13 +191,14 @@ public Response deleteTask(@PathParam("taskId") TaskId taskId,
  public void getResults(@PathParam("taskId") TaskId taskId,
  @PathParam("outputId") TaskId outputId,
  @PathParam("token") final long token,
+ @HeaderParam(PRESTO_MAX_SIZE) DataSize maxSize,
  @Suspended AsyncResponse asyncResponse)
  throws InterruptedException
  {
  requireNonNull(taskId, "taskId is null");
  requireNonNull(outputId, "outputId is null");
 
- CompletableFuture<BufferResult> bufferResultFuture = taskManager.getTaskResults(taskId, outputId, token, DEFAULT_MAX_SIZE);
+ CompletableFuture<BufferResult> bufferResultFuture = taskManager.getTaskResults(taskId, outputId, token, maxSize);
  bufferResultFuture = addTimeout(
  bufferResultFuture,
  () -> BufferResult.emptyResults(token, false),