From e772f51c7780fed3717be3111e6a1f1858aad31a Mon Sep 17 00:00:00 2001
From: Tomas Langer <tomas.langer@oracle.com>
Date: Mon, 4 Nov 2024 18:34:18 +0100
Subject: [PATCH] Attempt to fix infinite await in Jersey output writer.
 (#9460)

The count down latch is now counted down in close method of the output stream, which should cover all the possible cases (unless close is not called, which would cause other major issues as well, so the case is not handled).

Signed-off-by: Tomas Langer <tomas.langer@oracle.com>
---
 .../microprofile/server/JaxRsService.java     | 48 ++++++++++++++++++-
 .../webserver/http1/Http1ServerResponse.java  |  5 +-
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/microprofile/server/src/main/java/io/helidon/microprofile/server/JaxRsService.java b/microprofile/server/src/main/java/io/helidon/microprofile/server/JaxRsService.java
index 114619a70eb..5acf12c4a75 100644
--- a/microprofile/server/src/main/java/io/helidon/microprofile/server/JaxRsService.java
+++ b/microprofile/server/src/main/java/io/helidon/microprofile/server/JaxRsService.java
@@ -335,7 +335,8 @@ public OutputStream writeResponseStatusAndHeaders(long contentLengthParam,
             if (contentLength > 0) {
                 res.header(HeaderValues.create(HeaderNames.CONTENT_LENGTH, String.valueOf(contentLength)));
             }
-            this.outputStream = res.outputStream();
+            // in case there is an exception during close operation, we would lose the information and wait indefinitely
+            this.outputStream = new ReleaseLatchStream(cdl, res.outputStream());
             return outputStream;
         }
 
@@ -364,6 +365,10 @@ public void commit() {
             } catch (IOException e) {
                 cdl.countDown();
                 throw new UncheckedIOException(e);
+            } catch (Throwable e) {
+                // always release on commit, regardless of what happened
+                cdl.countDown();
+                throw e;
             }
         }
 
@@ -382,7 +387,7 @@ public boolean enableResponseBuffering() {
             return true;        // enable buffering in Jersey
         }
 
-        public void await() {
+        void await() {
             try {
                 cdl.await();
             } catch (InterruptedException e) {
@@ -391,6 +396,45 @@ public void await() {
         }
     }
 
+    private static class ReleaseLatchStream extends OutputStream {
+        private final CountDownLatch cdl;
+        private final OutputStream delegate;
+
+        private ReleaseLatchStream(CountDownLatch cdl, OutputStream delegate) {
+            this.cdl = cdl;
+            this.delegate = delegate;
+        }
+
+        @Override
+        public void write(int b) throws IOException {
+            delegate.write(b);
+        }
+
+        @Override
+        public void write(byte[] b) throws IOException {
+            delegate.write(b);
+        }
+
+        @Override
+        public void write(byte[] b, int off, int len) throws IOException {
+            delegate.write(b, off, len);
+        }
+
+        @Override
+        public void flush() throws IOException {
+            delegate.flush();
+        }
+
+        @Override
+        public void close() throws IOException {
+            try {
+                delegate.close();
+            } finally {
+                cdl.countDown();
+            }
+        }
+    }
+
     private static class BaseUriRequestUri {
         private final URI baseUri;
         private final URI requestUri;
diff --git a/webserver/webserver/src/main/java/io/helidon/webserver/http1/Http1ServerResponse.java b/webserver/webserver/src/main/java/io/helidon/webserver/http1/Http1ServerResponse.java
index ea41127d14c..03abb2ce312 100644
--- a/webserver/webserver/src/main/java/io/helidon/webserver/http1/Http1ServerResponse.java
+++ b/webserver/webserver/src/main/java/io/helidon/webserver/http1/Http1ServerResponse.java
@@ -19,6 +19,7 @@
 import java.io.BufferedOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.io.UncheckedIOException;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Objects;
@@ -843,7 +844,7 @@ public void close() {
             closingDelegate.closing();     // inform of imminent call to close for last flush
             try {
                 delegate.close();
-            } catch (IOException e) {
+            } catch (IOException | UncheckedIOException e) {
                 throw new ServerConnectionException("Failed to close server output stream", e);
             }
         }
@@ -856,7 +857,7 @@ void commit() {
             try {
                 flush();
                 closingDelegate.commit();
-            } catch (IOException e) {
+            } catch (IOException | UncheckedIOException e) {
                 throw new ServerConnectionException("Failed to flush server output stream", e);
             }
         }