Add discard() and optimal writeTo

Discard is used for ensuring a datastream is drained if needed, but a
no-op when no drain is needed. Added an optimal writeTo for
MultiByteBufferDataStream.

Author: mtdowling

io/src/main/java/software/amazon/smithy/java/io/datastream/DataStream.java

@@ -101,6 +101,20 @@ default void writeTo(OutputStream out) throws IOException {
         }
     }
 
+    /**
+     * Drain and release the underlying source without producing the bytes.
+     *
+     * <p>For in-memory backings this is a no-op. The data already arrived and lives in heap memory, so there's
+     * nothing to release. For sources backed by an {@link InputStream} or reactive
+     * {@link java.util.concurrent.Flow.Publisher}, this method consumes (or cancels and closes) the source so any
+     * underlying transport (socket, file descriptor, subscriber registration) can be released.
+     *
+     * @throws IOException if releasing the underlying source fails.
+     */
+    default void discard() throws IOException {
+        // no-op for in-memory backings
+    }
+
     /**
      * Read the contents of the stream into a ByteBuffer by reading all bytes from {@link #asInputStream()}.
      *

io/src/main/java/software/amazon/smithy/java/io/datastream/InputStreamDataStream.java

@@ -38,6 +38,21 @@ public void writeTo(OutputStream out) throws IOException {
         asInputStream().transferTo(out);
     }
 
+    @Override
+    public void discard() throws IOException {
+        if (consumed || closed) {
+            return;
+        }
+
+        consumed = true;
+        closed = true;
+        try (var is = inputStream) {
+            // skipNBytes is allowed to refuse and return early on some streams; transferTo to a
+            // null sink is the contract-correct way to drain to EOF without holding the bytes.
+            is.transferTo(OutputStream.nullOutputStream());
+        }
+    }
+
     @Override
     public boolean isReplayable() {
         return false;

io/src/main/java/software/amazon/smithy/java/io/datastream/MultiByteBufferDataStream.java

@@ -35,6 +35,8 @@
  */
 final class MultiByteBufferDataStream implements DataStream {
 
+    private static final int COPY_BUF_SIZE = 16384;
+
     private final List<ByteBuffer> buffers;
     private final long contentLength;
     private final String contentType;
@@ -124,14 +126,20 @@ public InputStream nextElement() {
 
     @Override
     public void writeTo(OutputStream out) throws IOException {
+        byte[] scratch = null;
         for (var buf : buffers) {
-            var dup = buf.duplicate();
-            if (dup.hasArray()) {
-                out.write(dup.array(), dup.arrayOffset() + dup.position(), dup.remaining());
+            if (buf.hasArray()) {
+                out.write(buf.array(), buf.arrayOffset() + buf.position(), buf.remaining());
             } else {
-                var tmp = new byte[dup.remaining()];
-                dup.get(tmp);
-                out.write(tmp);
+                if (scratch == null) {
+                    scratch = new byte[COPY_BUF_SIZE];
+                }
+                var dup = buf.duplicate();
+                while (dup.hasRemaining()) {
+                    int n = Math.min(scratch.length, dup.remaining());
+                    dup.get(scratch, 0, n);
+                    out.write(scratch, 0, n);
+                }
             }
         }
     }

io/src/main/java/software/amazon/smithy/java/io/datastream/PublisherDataStream.java

@@ -143,6 +143,54 @@ public void onComplete() {
         }
     }
 
+    @Override
+    public void discard() throws IOException {
+        if (consumed) {
+            return;
+        }
+
+        consumed = true;
+        var error = new AtomicReference<Throwable>();
+        var done = new CountDownLatch(1);
+
+        publisher.subscribe(new Flow.Subscriber<>() {
+            @Override
+            public void onSubscribe(Flow.Subscription s) {
+                s.request(Long.MAX_VALUE);
+            }
+
+            @Override
+            public void onNext(ByteBuffer buf) {
+                // drop on the floor
+            }
+
+            @Override
+            public void onError(Throwable t) {
+                error.set(t);
+                done.countDown();
+            }
+
+            @Override
+            public void onComplete() {
+                done.countDown();
+            }
+        });
+
+        try {
+            done.await();
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new IOException("Interrupted while discarding publisher data", e);
+        }
+
+        Throwable t = error.get();
+        if (t instanceof IOException ioe) {
+            throw ioe;
+        } else if (t != null) {
+            throw new IOException("Publisher error", t);
+        }
+    }
+
     private void innerSubscribe(Flow.Subscriber<? super ByteBuffer> subscriber) {
         consumed = true;
         publisher.subscribe(subscriber);

io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedDataStream.java

@@ -45,6 +45,11 @@ public void writeTo(OutputStream out) throws IOException {
         delegate.writeTo(out);
     }
 
+    @Override
+    public void discard() throws IOException {
+        delegate.discard();
+    }
+
     @Override
     public long contentLength() {
         return contentLength;

io/src/test/java/software/amazon/smithy/java/io/datastream/InputStreamDataStreamTest.java

@@ -111,4 +111,42 @@ public void writeToNotReplayable() throws IOException {
 
         assertThrows(IllegalStateException.class, () -> ds.writeTo(new ByteArrayOutputStream()));
     }
+
+    @Test
+    public void discardDrainsAndClosesUnderlyingStream() throws IOException {
+        // Track that the underlying stream got closed and drained without copying out the bytes.
+        var data = new byte[] {1, 2, 3, 4, 5};
+        var closed = new boolean[] {false};
+        var src = new ByteArrayInputStream(data) {
+            @Override
+            public void close() throws IOException {
+                closed[0] = true;
+                super.close();
+            }
+        };
+        var ds = DataStream.ofInputStream(src);
+
+        ds.discard();
+
+        assertThat(closed[0], is(true));
+        assertThat(ds.isAvailable(), is(false));
+        // Drained to EOF.
+        assertThat(src.available(), equalTo(0));
+    }
+
+    @Test
+    public void discardIsIdempotent() throws IOException {
+        var ds = DataStream.ofInputStream(new ByteArrayInputStream(new byte[] {1, 2, 3}));
+
+        ds.discard();
+        ds.discard(); // Second discard must not throw
+    }
+
+    @Test
+    public void discardAfterConsumptionIsNoOp() throws IOException {
+        var ds = DataStream.ofInputStream(new ByteArrayInputStream(new byte[] {1, 2, 3}));
+        ds.asInputStream();
+
+        ds.discard(); // Don't reopen the stream just to discard it
+    }
 }

io/src/test/java/software/amazon/smithy/java/io/datastream/MultiByteBufferDataStreamTest.java

@@ -0,0 +1,109 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.io.datastream;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+public class MultiByteBufferDataStreamTest {
+
+    @Test
+    public void writeToHeapBuffers() throws IOException {
+        // hasArray() == true for heap buffers, exercising the array-fast path.
+        var b1 = ByteBuffer.wrap("hello ".getBytes(StandardCharsets.UTF_8));
+        var b2 = ByteBuffer.wrap("world".getBytes(StandardCharsets.UTF_8));
+        var ds = DataStream.ofByteBuffers(List.of(b1, b2), 11, "text/plain");
+        var out = new ByteArrayOutputStream();
+
+        ds.writeTo(out);
+
+        assertArrayEquals("hello world".getBytes(StandardCharsets.UTF_8), out.toByteArray());
+    }
+
+    @Test
+    public void writeToReadOnlyBuffersUsesScratch() throws IOException {
+        // asReadOnlyBuffer().hasArray() returns false (HeapByteBufferR) — exercises the scratch
+        // copy path. This is the case that motivated the optimization: the JDK HttpClient produces
+        // read-only heap buffers via ZeroCopyBodySubscriber, and writeTo previously allocated a new
+        // byte[] per chunk. Verify the bytes still come out correctly.
+        var b1 = ByteBuffer.wrap("read-only ".getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer();
+        var b2 = ByteBuffer.wrap("path".getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer();
+        assertThat(b1.hasArray(), is(false));
+        var ds = DataStream.ofByteBuffers(List.of(b1, b2), 14, "text/plain");
+        var out = new ByteArrayOutputStream();
+
+        ds.writeTo(out);
+
+        assertArrayEquals("read-only path".getBytes(StandardCharsets.UTF_8), out.toByteArray());
+    }
+
+    @Test
+    public void writeToReadOnlyBuffersDoesNotMutatePosition() throws IOException {
+        // Defends the contract: replayable means writeTo is safe to call again. The previous
+        // implementation called dup.get(tmp) which advanced the duplicate's position, but the
+        // current implementation must use a duplicate so the original buffer's position survives.
+        var b1 = ByteBuffer.wrap("a".getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer();
+        var b2 = ByteBuffer.wrap("b".getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer();
+        var ds = DataStream.ofByteBuffers(List.of(b1, b2), 2, "text/plain");
+
+        var out1 = new ByteArrayOutputStream();
+        ds.writeTo(out1);
+        var out2 = new ByteArrayOutputStream();
+        ds.writeTo(out2);
+
+        assertArrayEquals("ab".getBytes(StandardCharsets.UTF_8), out1.toByteArray());
+        assertArrayEquals("ab".getBytes(StandardCharsets.UTF_8), out2.toByteArray());
+    }
+
+    @Test
+    public void writeToBufferLargerThanScratch() throws IOException {
+        // Drive the inner copy loop past one scratch fill (scratch is 16 KiB internally).
+        byte[] payload = new byte[40_000];
+        for (int i = 0; i < payload.length; i++) {
+            payload[i] = (byte) (i & 0xff);
+        }
+        var b1 = ByteBuffer.wrap(payload, 0, 25_000).asReadOnlyBuffer();
+        var b2 = ByteBuffer.wrap(payload, 25_000, 15_000).asReadOnlyBuffer();
+        var ds = DataStream.ofByteBuffers(List.of(b1, b2), payload.length, "application/octet-stream");
+        var out = new ByteArrayOutputStream();
+
+        ds.writeTo(out);
+
+        assertArrayEquals(payload, out.toByteArray());
+    }
+
+    @Test
+    public void writeToMixedBuffers() throws IOException {
+        // First buffer hits the heap-array fast path, second buffer hits the scratch copy path.
+        // Validates the scratch is allocated lazily on the first non-array buffer.
+        var b1 = ByteBuffer.wrap("heap ".getBytes(StandardCharsets.UTF_8));
+        var b2 = ByteBuffer.wrap("readonly".getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer();
+        var ds = DataStream.ofByteBuffers(List.of(b1, b2), 13, "text/plain");
+        var out = new ByteArrayOutputStream();
+
+        ds.writeTo(out);
+
+        assertArrayEquals("heap readonly".getBytes(StandardCharsets.UTF_8), out.toByteArray());
+    }
+
+    @Test
+    public void asByteBufferStitchesAcrossChunks() {
+        var b1 = ByteBuffer.wrap("aa".getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer();
+        var b2 = ByteBuffer.wrap("bb".getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer();
+        var ds = DataStream.ofByteBuffers(List.of(b1, b2), 4, "text/plain");
+
+        assertThat(ds.asByteBuffer(), equalTo(ByteBuffer.wrap("aabb".getBytes(StandardCharsets.UTF_8))));
+    }
+}

io/src/test/java/software/amazon/smithy/java/io/datastream/PublisherDataStreamTest.java

@@ -78,4 +78,59 @@ void writeToEmpty() throws IOException {
 
         assertEquals(0, out.size());
     }
+
+    @Test
+    void discardDrainsPublisherToCompletion() throws IOException {
+        var publisher = new SubmissionPublisher<ByteBuffer>();
+        var ds = DataStream.ofPublisher(publisher, null, -1);
+
+        var discardThread = Thread.startVirtualThread(() -> {
+            try {
+                ds.discard();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+
+        // Wait for discard's subscriber to register before submitting items.
+        while (publisher.getNumberOfSubscribers() < 1) {
+            Thread.onSpinWait();
+        }
+        publisher.submit(ByteBuffer.wrap("dropped".getBytes(StandardCharsets.UTF_8)));
+        publisher.close();
+
+        try {
+            discardThread.join();
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+
+        // After discard, the stream is consumed and no longer available for replay.
+        assertEquals(false, ds.isAvailable());
+    }
+
+    @Test
+    void discardIsIdempotent() throws IOException {
+        var publisher = new SubmissionPublisher<ByteBuffer>();
+        var ds = DataStream.ofPublisher(publisher, null, -1);
+
+        Thread.startVirtualThread(publisher::close);
+        ds.discard();
+
+        // Second call exits immediately without re-subscribing.
+        ds.discard();
+    }
+
+    @Test
+    void discardPropagatesPublisherError() {
+        var publisher = new SubmissionPublisher<ByteBuffer>();
+        var ds = DataStream.ofPublisher(publisher, null, -1);
+
+        var ex = new RuntimeException("boom");
+        Thread.startVirtualThread(() -> publisher.closeExceptionally(ex));
+
+        var thrown = assertThrows(IOException.class, ds::discard);
+        assertEquals("Publisher error", thrown.getMessage());
+        assertEquals(ex, thrown.getCause());
+    }
 }

io/src/test/java/software/amazon/smithy/java/io/datastream/WrappedDataStreamTest.java

@@ -48,4 +48,25 @@ public void writeToDelegates() throws IOException {
 
         assertArrayEquals(data, out.toByteArray());
     }
+
+    @Test
+    public void discardDelegates() throws IOException {
+        // The wrapped stream's discard must reach through to the underlying source so InputStream
+        // / Publisher resources actually get released. Verify via observable side effect on
+        // InputStreamDataStream (drain-to-EOF + close).
+        var closed = new boolean[] {false};
+        var src = new ByteArrayInputStream("payload".getBytes(StandardCharsets.UTF_8)) {
+            @Override
+            public void close() throws IOException {
+                closed[0] = true;
+                super.close();
+            }
+        };
+        var inner = DataStream.ofInputStream(src);
+        var ds = DataStream.withMetadata(inner, "text/plain", 7L, false);
+
+        ds.discard();
+
+        assertThat(closed[0], is(true));
+    }
 }