Revert "Register Servlet WriteListener during request dispatch"

This reverts commit e958aee5adca061d2650f9f5378ce93a9e92ae7d.

3 files changed, 18 insertions, 58 deletions

diff --git a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HttpRequestDispatch.java b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HttpRequestDispatch.java
index 57e5ba0cf03..01a54665f49 100644
--- a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HttpRequestDispatch.java
+++ b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HttpRequestDispatch.java
@@ -143,7 +143,6 @@ class HttpRequestDispatch {
 
     @SuppressWarnings("try")
     private ServletRequestReader handleRequest() throws IOException {
-        servletResponseController.registerWriteListener();
         HttpRequest jdiscRequest = HttpRequestFactory.newJDiscRequest(jDiscContext.container, servletRequest);
         ContentChannel requestContentChannel;
 
@@ -175,7 +174,6 @@ class HttpRequestDispatch {
 
     ContentChannel handleRequestFilterResponse(Response response) {
         try {
-            servletResponseController.registerWriteListener();
             servletRequest.getInputStream().close();
             ContentChannel responseContentChannel = servletResponseController.responseHandler.handleResponse(response);
             servletResponseController.finishedFuture().whenComplete(completeRequestCallback);
diff --git a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletOutputStreamWriter.java b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletOutputStreamWriter.java
index 4ad181d2bf6..5abdb98edbd 100644
--- a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletOutputStreamWriter.java
+++ b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletOutputStreamWriter.java
@@ -36,8 +36,7 @@ public class ServletOutputStreamWriter {
     private enum State {
         NOT_STARTED,
         WAITING_FOR_WRITE_POSSIBLE_CALLBACK,
-        WAITING_FOR_FIRST_BUFFER,
-        WAITING_FOR_SUBSEQUENT_BUFFER,
+        WAITING_FOR_BUFFER,
         WRITING_BUFFERS,
         FINISHED_OR_ERROR
     }
@@ -78,32 +77,13 @@ public class ServletOutputStreamWriter {
         this.metricReporter = metricReporter;
     }
 
-    public void registerWriteListener() {
-        outputStream.setWriteListener(writeListener);
-    }
-
     public void sendErrorContentAndCloseAsync(ByteBuffer errorContent) {
-        boolean thisThreadShouldWrite;
-
         synchronized (monitor) {
             // Assert that no content has been written as it is too late to write error response if the response is committed.
-            switch (state) {
-                case NOT_STARTED:
-                    queueErrorContent_holdingLock(errorContent);
-                    state = State.WAITING_FOR_WRITE_POSSIBLE_CALLBACK;
-                    thisThreadShouldWrite = false;
-                    break;
-                case WAITING_FOR_FIRST_BUFFER:
-                    queueErrorContent_holdingLock(errorContent);
-                    state = State.WRITING_BUFFERS;
-                    thisThreadShouldWrite = true;
-                    break;
-                default:
-                    throw createAndLogAssertionError("Invalid state: " + state);
-            }
-        }
-        if (thisThreadShouldWrite) {
-            writeBuffersInQueueToOutputStream();
+            assertStateIs(state, State.NOT_STARTED);
+            queueErrorContent_holdingLock(errorContent);
+            state = State.WAITING_FOR_WRITE_POSSIBLE_CALLBACK;
+            outputStream.setWriteListener(writeListener);
         }
     }
 
@@ -124,12 +104,12 @@ public class ServletOutputStreamWriter {
             switch (state) {
                 case NOT_STARTED:
                     state = State.WAITING_FOR_WRITE_POSSIBLE_CALLBACK;
+                    outputStream.setWriteListener(writeListener);
                     break;
                 case WAITING_FOR_WRITE_POSSIBLE_CALLBACK:
                 case WRITING_BUFFERS:
                     break;
-                case WAITING_FOR_FIRST_BUFFER:
-                case WAITING_FOR_SUBSEQUENT_BUFFER:
+                case WAITING_FOR_BUFFER:
                     thisThreadShouldWrite = true;
                     state = State.WRITING_BUFFERS;
                     break;
@@ -171,7 +151,7 @@ public class ServletOutputStreamWriter {
                 contentPart = responseContentQueue.pollFirst();
 
                 if (contentPart == null && lastOperationWasFlush) {
-                    state = State.WAITING_FOR_SUBSEQUENT_BUFFER;
+                    state = State.WAITING_FOR_BUFFER;
                     return;
                 }
             }
@@ -272,16 +252,12 @@ public class ServletOutputStreamWriter {
 
     private static void assertStateIs(State currentState, State expectedState) {
         if (currentState != expectedState) {
-            throw createAndLogAssertionError("Expected state " + expectedState + ", got state " + currentState);
+            AssertionError error = new AssertionError("Expected state " + expectedState + ", got state " + currentState);
+            log.log(Level.WARNING, "Assertion failed.", error);
+            throw error;
         }
     }
 
-    private static AssertionError createAndLogAssertionError(String detailedMessage) {
-        AssertionError error = new AssertionError(detailedMessage);
-        log.log(Level.WARNING, "Assertion failed.", error);
-        return error;
-    }
-
     public void fail(Throwable t) {
         setFinished(Optional.of(t));
     }
@@ -289,27 +265,16 @@ public class ServletOutputStreamWriter {
     private final WriteListener writeListener = new WriteListener() {
         @Override
         public void onWritePossible() throws IOException {
-            boolean shouldWriteBuffers = false;
             synchronized (monitor) {
-                switch (state) {
-                    case NOT_STARTED:
-                        state = State.WAITING_FOR_FIRST_BUFFER;
-                        break;
-                    case WAITING_FOR_WRITE_POSSIBLE_CALLBACK:
-                        state = State.WRITING_BUFFERS;
-                        shouldWriteBuffers = true;
-                        break;
-                    case FINISHED_OR_ERROR:
-                        return;
-                    case WAITING_FOR_FIRST_BUFFER:
-                    case WAITING_FOR_SUBSEQUENT_BUFFER:
-                    case WRITING_BUFFERS:
-                        throw createAndLogAssertionError("Invalid state: " + state);
+                if (state == State.FINISHED_OR_ERROR) {
+                    return;
                 }
+
+                assertStateIs(state, State.WAITING_FOR_WRITE_POSSIBLE_CALLBACK);
+                state = State.WRITING_BUFFERS;
             }
-            if (shouldWriteBuffers) {
-                writeBuffersInQueueToOutputStream();
-            }
+
+            writeBuffersInQueueToOutputStream();
         }
 
         @Override
diff --git a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletResponseController.java b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletResponseController.java
index a06d39ed016..0d8cf936d1d 100644
--- a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletResponseController.java
+++ b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletResponseController.java
@@ -69,9 +69,6 @@ public class ServletResponseController {
                 new ServletOutputStreamWriter(servletResponse.getOutputStream(), executor, metricReporter);
     }
 
-    public void registerWriteListener() {
-        servletOutputStreamWriter.registerWriteListener();
-    }
 
     private static int getStatusCode(Throwable t) {
         if (t instanceof BindingNotFoundException) {