[api] Refactor PublisherBytesSupplier.java

api/src/main/java/ai/djl/inference/streaming/PublisherBytesSupplier.java

@@ -14,13 +14,10 @@
 
 import ai.djl.ndarray.BytesSupplier;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 /**
@@ -29,16 +26,14 @@
  */
 public class PublisherBytesSupplier implements BytesSupplier {
 
-    private final List<byte[]> allData;
-    private final AtomicBoolean completed;
     private Consumer<byte[]> subscriber;
-    private final AtomicInteger dataPushed;
+    private CountDownLatch latch;
+    private CompletableFuture<Void> future;
 
     /** Constructs a {@link PublisherBytesSupplier}. */
     public PublisherBytesSupplier() {
-        allData = new ArrayList<>();
-        completed = new AtomicBoolean();
-        dataPushed = new AtomicInteger();
+        latch = new CountDownLatch(1);
+        future = new CompletableFuture<>();
     }
 
     /**
@@ -48,83 +43,46 @@ public PublisherBytesSupplier() {
      * @param lastChunk true if this is the last chunk
      */
     public void appendContent(byte[] data, boolean lastChunk) {
-        synchronized (allData) {
-            allData.add(data);
+        if (subscriber == null) {
+            try {
+                if (!latch.await(2, TimeUnit.MINUTES)) {
+                    throw new IllegalStateException("Wait for subscriber timeout.");
+                }
+                if (subscriber == null) {
+                    // workaround Spotbugs
+                    throw new IllegalStateException("subscriber is not set.");
+                }
+            } catch (InterruptedException e) {
+                throw new IllegalStateException("Append content interrupted.", e);
+            }
         }
+        subscriber.accept(data);
         if (lastChunk) {
-            completed.set(true);
+            subscriber.accept(null);
+            future.complete(null);
         }
-        pushData();
     }
 
     /**
      * Adds the subscriber to the {@link BytesSupplier} to get notified about additional data.
      *
      * @param subscriber a consumer function that will receive bytes when new daata is added and
      *     null when completed
+     * @return a {@code CompletableFuture} object
      */
-    public void subscribe(Consumer<byte[]> subscriber) {
+    public CompletableFuture<Void> subscribe(Consumer<byte[]> subscriber) {
         if (this.subscriber != null) {
             throw new IllegalStateException(
                     "The PublisherBytesSupplier only allows a single Subscriber");
         }
         this.subscriber = subscriber;
-        pushData();
-    }
-
-    private void pushData() {
-        if (subscriber == null) {
-            return;
-        }
-
-        int dataAvailable;
-        synchronized (allData) {
-            dataAvailable = allData.size();
-        }
-
-        int sent = dataPushed.getAndSet(dataAvailable);
-        if (sent < dataAvailable) {
-            synchronized (this) {
-                for (; sent < dataAvailable; sent++) {
-                    subscriber.accept(allData.get(sent));
-                }
-                if (completed.get()) {
-                    subscriber.accept(null);
-                }
-            }
-        }
-    }
-
-    /** Waits until completed before passing thread (BLOCKS THREAD!). */
-    @SuppressWarnings("PMD.EmptyControlStatement")
-    public void waitToRead() {
-        // Block until complete!!!
-        while (!completed.get()) {
-            // Do nothing
-        }
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public byte[] getAsBytes() {
-        if (!completed.get()) {
-            throw new IllegalStateException(
-                    "PublisherByteSupplier must be completely filled before reading.");
-        }
-
-        try (ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
-            for (byte[] data : allData) {
-                bos.write(data);
-            }
-            return bos.toByteArray();
-        } catch (IOException e) {
-            throw new AssertionError("Failed to read BytesSupplier", e);
-        }
+        latch.countDown();
+        return future;
     }
 
     /** {@inheritDoc} */
     @Override
     public ByteBuffer toByteBuffer() {
-        return ByteBuffer.wrap(getAsBytes());
+        throw new UnsupportedOperationException("Not supported.");
     }
 }

api/src/test/java/ai/djl/inference/streaming/PublisherBytesSupplierTest.java

@@ -15,32 +15,38 @@
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class PublisherBytesSupplierTest {
 
     @Test
-    public void test() {
+    public void test() throws ExecutionException, InterruptedException {
         AtomicInteger contentCount = new AtomicInteger();
         PublisherBytesSupplier supplier = new PublisherBytesSupplier();
 
-        // Add to supplier without subscriber
-        supplier.appendContent(new byte[] {1}, false);
-        Assert.assertEquals(contentCount.get(), 0);
+        new Thread(
+                        () -> {
+                            // Add to supplier without subscriber
+                            supplier.appendContent(new byte[] {1}, false);
+                            // Add to supplier with subscriber
+                            supplier.appendContent(new byte[] {1}, true);
+                        })
+                .start();
 
         // Subscribing with data should trigger subscriptions
-        supplier.subscribe(
-                d -> {
-                    if (d == null) {
-                        // Do nothing on completion
-                        return;
-                    }
-                    contentCount.getAndIncrement();
-                });
-        Assert.assertEquals(contentCount.get(), 1);
+        CompletableFuture<Void> future =
+                supplier.subscribe(
+                        d -> {
+                            if (d == null) {
+                                // Do nothing on completion
+                                return;
+                            }
+                            contentCount.getAndIncrement();
+                        });
 
-        // Add to supplier with subscriber
-        supplier.appendContent(new byte[] {1}, true);
+        future.get();
         Assert.assertEquals(contentCount.get(), 2);
     }
 }