Bug2.groovy

@Grab('io.undertow:undertow-core:2.2.25.Final')
@Grab('io.undertow:undertow-servlet:2.2.25.Final')
@Grab('com.squareup.okhttp3:okhttp:4.9.1')

import io.undertow.Undertow
import io.undertow.channels.DetachableStreamSinkChannel
import io.undertow.server.HttpServerExchange
import io.undertow.servlet.Servlets
import io.undertow.servlet.api.DeploymentInfo
import io.undertow.servlet.api.ServletContainer
import okhttp3.Call
import okhttp3.Callback
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response

import javax.servlet.AsyncContext
import javax.servlet.WriteListener
import javax.servlet.http.HttpServlet
import javax.servlet.http.HttpServletRequest
import javax.servlet.http.HttpServletResponse
import java.util.concurrent.CountDownLatch

DeploymentInfo deploymentInfo = Servlets.deployment()
  .setClassLoader(this.class.classLoader)
  .setContextPath('/')
  .setDeploymentName('test.war')
  .addServlets(
    Servlets.servlet("TestServlet2", TestServlet2)
      .addMapping('/test').setAsyncSupported(true)
  )

ServletContainer container = ServletContainer.Factory.newInstance()
def manager = container.addDeployment(deploymentInfo)
manager.deploy()

Undertow undertow = Undertow.builder().addHttpListener(8080, '::1')
  .setHandler(manager.start()).build()
undertow.start()

class TestServlet2 extends HttpServlet {
  @Override
  void doGet(HttpServletRequest req, HttpServletResponse res) {
    if (!req.getAttribute('dispatched')) {
      AsyncContext asyncContext = req.startAsync()
      asyncContext.start {
        res.outputStream.writeListener = new WriteListener() {
          volatile dispatch = true

          @Override
          void onWritePossible() throws IOException {
            if (!dispatch) {
              return
            }
            dispatch = false
            req.setAttribute('dispatched', 1)
            asyncContext.dispatch()
          }

          @Override
          void onError(Throwable t) {
            t.printStackTrace(System.err)
          }
        }
      }
    } else {
      res.outputStream << 'foo'
      res.flushBuffer()

      // workaround: not always feasible (e.g. sendError() calls already responseDone())
//      HttpServerExchange exchange = req.exchange
//      DetachableStreamSinkChannel channel = exchange.@responseChannel
//      channel?.resumeWrites()
    }
  }
}

println 'Server started on http://[::1]:8080/'

def client = new OkHttpClient()

int requestCount = 100
def latch = new CountDownLatch(requestCount)

requestCount.times {
  Request request = new Request.Builder()
    .url('http://[::1]:8080/test')
    .get()
    .build()

  client.newCall(request).enqueue(new Callback() {
    @Override
    void onResponse(Call call, Response response) {
      latch.countDown()
      assert response.body().string() == 'foo'
      response.close()
    }

    @Override
    void onFailure(Call call, IOException e) {
      e.printStackTrace(System.err)
      System.exit 1
    }
  })
}

println "Waiting for all the requests to finish"
latch.await()
println "Finished"
undertow.stop()