avoid baos::writeTo on virtual threads

service/src/main/java/org/whispersystems/textsecuregcm/util/BufferingInterceptor.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.util;
+
+import org.glassfish.jersey.message.internal.CommittingOutputStream;
+
+import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.ext.WriterInterceptor;
+import javax.ws.rs.ext.WriterInterceptorContext;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+
+
+/**
+ * This is an elaborate workaround to avoid doing blocking operations under synchronized blocks, which is currently a
+ * suboptimal case for virtual threads.
+ * <p>
+ * Jersey's {@link CommittingOutputStream} has two modes: direct write and buffered writes. In buffered mode, if the
+ * total amount written does not exceed the output stream's buffer size, CommittingOutputStream will compute the
+ * content-length for us. However, when it passes through our write to its own underlying output stream it uses
+ * {@link ByteArrayOutputStream#writeTo(OutputStream)} which performs the write under a synchronized block.
+ * <p>
+ * If we just disable buffering, we lose our content length. However, we can't really set content-length ourselves
+ * without the same access to internal state that CommittingOutputStream has. Fortunately, the underlying OutputStream
+ * wrapped by CommittingOutputStream ALSO has an internal buffer, and can compute the content-length from that if the
+ * content fits. But to make use of that, we need to avoid flushing that output stream until calling close, so that the
+ * underlying output stream can see that it has all the data. Unfortunately the runtime inserts manual flushes after
+ * writes rather than letting the underlying output stream handle it.
+ * <p>
+ * So here we disable buffering on CommittingOutputStream, and buffer ourselves. We don't write anything to the
+ * CommittingOutputStream until we are going to close, and we do nothing on flush.
+ */
+public class BufferingInterceptor implements WriterInterceptor {
+
+  @Override
+  public void aroundWriteTo(final WriterInterceptorContext ctx) throws IOException, WebApplicationException {
+    final OutputStream orig = ctx.getOutputStream();
+    if (Thread.currentThread().isVirtual() && orig instanceof CommittingOutputStream cos) {
+      cos.enableBuffering(0);
+      ctx.setOutputStream(new BufferingOutputStream(cos));
+    }
+    ctx.proceed();
+  }
+
+  private static class BufferingOutputStream extends ByteArrayOutputStream {
+
+    private final CommittingOutputStream original;
+
+    BufferingOutputStream(final CommittingOutputStream original) {
+      this.original = original;
+    }
+
+    @Override
+    public void close() throws IOException {
+      original.write(buf, 0, count);
+      original.close();
+      super.close();
+    }
+  }
+}