From 3cb803b88f6fd279bff2ac33ffcf51421bc3128e Mon Sep 17 00:00:00 2001 From: Abhishek Agarwal <1477457+abhishekagarwal87@users.noreply.github.com> Date: Sat, 3 Dec 2022 21:30:37 +0530 Subject: [PATCH 1/2] Fix the truncated query response --- .../main/java/org/apache/druid/server/QueryResource.java | 3 +-- .../java/org/apache/druid/sql/http/ResultFormat.java | 5 ++++- .../main/java/org/apache/druid/sql/http/SqlResource.java | 9 ++++++--- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/server/src/main/java/org/apache/druid/server/QueryResource.java b/server/src/main/java/org/apache/druid/server/QueryResource.java index 743ca9e60ba7..f4d1f7794487 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResource.java +++ b/server/src/main/java/org/apache/druid/server/QueryResource.java @@ -230,8 +230,7 @@ public void write(OutputStream outputStream) throws WebApplicationException // json serializer will always close the yielder jsonWriter.writeValue(os, yielder); - os.flush(); // Some types of OutputStream suppress flush errors in the .close() method. - os.close(); + os.flush(); } catch (Exception ex) { e = ex; diff --git a/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java b/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java index 0c450b55b5f3..96eb49e6e60c 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java +++ b/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java @@ -146,7 +146,10 @@ public interface Writer extends Closeable void writeRowEnd() throws IOException; /** - * End of the response. Must allow the user to know that they have read all data successfully. + * End of the response. Any data buffered by the writer must be flushed as well in this method. Some implementations + * might be writing a trailer at the end for clients to confirm that response is not truncated. However, that is + * no longer needed now since http client can rely on standard chunked encoding mechanics to detect a truncated + * response. */ void writeResponseEnd() throws IOException; } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 968a722caa5d..c63818c48eeb 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -143,9 +143,12 @@ public void write(OutputStream output) throws IOException, WebApplicationExcepti Exception e = null; CountingOutputStream os = new CountingOutputStream(output); Yielder yielder = finalYielder; - - try (final ResultFormat.Writer writer = sqlQuery.getResultFormat() - .createFormatter(os, jsonMapper)) { + // The writer below is not closed by us, since it would also close output that was passed as an + // argument. Instead, we would just call the ResultFormat.Writer#writeResponseEnd to flush + // any buffered data. + final ResultFormat.Writer writer = sqlQuery.getResultFormat() + .createFormatter(os, jsonMapper); + try { writer.writeResponseStart(); if (sqlQuery.includeHeader()) { From 7f2d5c3aaeb9c24c50aff83c37b40195a4bb6112 Mon Sep 17 00:00:00 2001 From: Abhishek Agarwal <1477457+abhishekagarwal87@users.noreply.github.com> Date: Tue, 6 Dec 2022 17:06:21 +0530 Subject: [PATCH 2/2] Use a different approach --- .../org/apache/druid/server/QueryResource.java | 2 ++ .../apache/druid/sql/http/ArrayLinesWriter.java | 2 ++ .../org/apache/druid/sql/http/ArrayWriter.java | 2 ++ .../java/org/apache/druid/sql/http/CsvWriter.java | 3 ++- .../apache/druid/sql/http/ObjectLinesWriter.java | 2 ++ .../org/apache/druid/sql/http/ObjectWriter.java | 2 ++ .../org/apache/druid/sql/http/ResultFormat.java | 14 ++++++++++---- .../org/apache/druid/sql/http/SqlResource.java | 9 +++------ 8 files changed, 25 insertions(+), 11 deletions(-) diff --git a/server/src/main/java/org/apache/druid/server/QueryResource.java b/server/src/main/java/org/apache/druid/server/QueryResource.java index f4d1f7794487..b821830a15ca 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResource.java +++ b/server/src/main/java/org/apache/druid/server/QueryResource.java @@ -231,6 +231,8 @@ public void write(OutputStream outputStream) throws WebApplicationException jsonWriter.writeValue(os, yielder); os.flush(); + + // Do not close the output stream since that is not managed by QueryResource. } catch (Exception ex) { e = ex; diff --git a/sql/src/main/java/org/apache/druid/sql/http/ArrayLinesWriter.java b/sql/src/main/java/org/apache/druid/sql/http/ArrayLinesWriter.java index beda6deceaba..0d0d18b43e48 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/ArrayLinesWriter.java +++ b/sql/src/main/java/org/apache/druid/sql/http/ArrayLinesWriter.java @@ -41,6 +41,8 @@ public ArrayLinesWriter(final OutputStream outputStream, final ObjectMapper json this.serializers = jsonMapper.getSerializerProviderInstance(); this.outputStream = outputStream; this.jsonGenerator = jsonMapper.writer().getFactory().createGenerator(outputStream); + // Disable auto closing of target stream since that is not managed by this writer. + jsonGenerator.configure(JsonGenerator.Feature.AUTO_CLOSE_TARGET, false); jsonGenerator.setRootValueSeparator(new SerializedString("\n")); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/ArrayWriter.java b/sql/src/main/java/org/apache/druid/sql/http/ArrayWriter.java index cd863d5bf175..eeb1907892a3 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/ArrayWriter.java +++ b/sql/src/main/java/org/apache/druid/sql/http/ArrayWriter.java @@ -44,6 +44,8 @@ public ArrayWriter(final OutputStream outputStream, final ObjectMapper jsonMappe this.jsonGenerator = jsonMapper.getFactory().createGenerator(outputStream); this.outputStream = outputStream; + // Disable auto closing of target stream since that is not managed by this writer. + jsonGenerator.configure(JsonGenerator.Feature.AUTO_CLOSE_TARGET, false); // Disable automatic JSON termination, so clients can detect truncated responses. jsonGenerator.configure(JsonGenerator.Feature.AUTO_CLOSE_JSON_CONTENT, false); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/CsvWriter.java b/sql/src/main/java/org/apache/druid/sql/http/CsvWriter.java index e5e997306845..ba4f49f4a7d7 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/CsvWriter.java +++ b/sql/src/main/java/org/apache/druid/sql/http/CsvWriter.java @@ -126,6 +126,7 @@ public void writeRowEnd() @Override public void close() throws IOException { - writer.close(); + // Just flush the writer but do not close it, since we do not want to close the target output stream + writer.flush(); } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/ObjectLinesWriter.java b/sql/src/main/java/org/apache/druid/sql/http/ObjectLinesWriter.java index a593b9b21b2a..02d4db5454f9 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/ObjectLinesWriter.java +++ b/sql/src/main/java/org/apache/druid/sql/http/ObjectLinesWriter.java @@ -41,6 +41,8 @@ public ObjectLinesWriter(final OutputStream outputStream, final ObjectMapper jso this.serializers = jsonMapper.getSerializerProviderInstance(); this.outputStream = outputStream; this.jsonGenerator = jsonMapper.writer().getFactory().createGenerator(outputStream); + // Disable auto closing of target stream since that is not managed by this writer. + jsonGenerator.configure(JsonGenerator.Feature.AUTO_CLOSE_TARGET, false); jsonGenerator.setRootValueSeparator(new SerializedString("\n")); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/ObjectWriter.java b/sql/src/main/java/org/apache/druid/sql/http/ObjectWriter.java index bdab65a1f7e6..12d75b336694 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/ObjectWriter.java +++ b/sql/src/main/java/org/apache/druid/sql/http/ObjectWriter.java @@ -47,6 +47,8 @@ public ObjectWriter(final OutputStream outputStream, final ObjectMapper jsonMapp this.jsonGenerator = jsonMapper.getFactory().createGenerator(outputStream); this.outputStream = outputStream; + // Disable auto closing of target stream since that is not managed by this writer. + jsonGenerator.configure(JsonGenerator.Feature.AUTO_CLOSE_TARGET, false); // Disable automatic JSON termination, so clients can detect truncated responses. jsonGenerator.configure(JsonGenerator.Feature.AUTO_CLOSE_JSON_CONTENT, false); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java b/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java index 96eb49e6e60c..9d72d7bc7b27 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java +++ b/sql/src/main/java/org/apache/druid/sql/http/ResultFormat.java @@ -146,12 +146,18 @@ public interface Writer extends Closeable void writeRowEnd() throws IOException; /** - * End of the response. Any data buffered by the writer must be flushed as well in this method. Some implementations - * might be writing a trailer at the end for clients to confirm that response is not truncated. However, that is - * no longer needed now since http client can rely on standard chunked encoding mechanics to detect a truncated - * response. + * End of the response. Some implementations might be writing a trailer at the end for clients to confirm + * that response is not truncated. However, that is no longer needed now since http client can rely on + * standard chunked encoding mechanics to detect a truncated response. */ void writeResponseEnd() throws IOException; + + /** + * Implementations must not close the target output stream that they are writing to. + * @throws IOException + */ + @Override + void close() throws IOException; } @JsonCreator diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index c63818c48eeb..6a6f93ff194c 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -143,12 +143,9 @@ public void write(OutputStream output) throws IOException, WebApplicationExcepti Exception e = null; CountingOutputStream os = new CountingOutputStream(output); Yielder yielder = finalYielder; - // The writer below is not closed by us, since it would also close output that was passed as an - // argument. Instead, we would just call the ResultFormat.Writer#writeResponseEnd to flush - // any buffered data. - final ResultFormat.Writer writer = sqlQuery.getResultFormat() - .createFormatter(os, jsonMapper); - try { + // The ResultWriter gets closed automatically however that doesn't close os. + try (final ResultFormat.Writer writer = sqlQuery.getResultFormat() + .createFormatter(os, jsonMapper)) { writer.writeResponseStart(); if (sqlQuery.includeHeader()) {