Repository: deepjavalibrary/djl
Pull request #2831 (2023-10-31 20:14:07)

[api] Refactor PublisherBytesSupplier.java
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1375 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.08% compared to head [(`e9ec92b`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.18%.
> Report is 903 commits behind head on master.

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2831      +/-   ##
============================================
+ Coverage     72.08%   72.18%   +0.09%     
- Complexity     5126     7136    +2010     
============================================
  Files           473      708     +235     
  Lines         21970    31873    +9903     
  Branches       2351     3313     +962     
============================================
+ Hits          15838    23006    +7168     
- Misses         4925     7289    +2364     
- Partials       1207     1578     +371     
```


| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [...ava/ai/djl/inference/streaming/StreamingBlock.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL2luZmVyZW5jZS9zdHJlYW1pbmcvU3RyZWFtaW5nQmxvY2suamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Dimension.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9EaW1lbnNpb24uamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Unit.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9Vbml0LmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/audio/Audio.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2F1ZGlvL0F1ZGlvLmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...ava/ai/djl/modality/cv/output/DetectedObjects.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9EZXRlY3RlZE9iamVjdHMuamF2YQ==) | `96.29% <100.00%> (+1.29%)` | :arrow_up: |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <100.00%> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...i/djl/modality/cv/transform/RandomResizedCrop.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21SZXNpemVkQ3JvcC5qYXZh) | `94.11% <100.00%> (+5.22%)` | :arrow_up: |
| ... and [225 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [385 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2831?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/inference/streaming/PublisherBytesSupplier.java b/api/src/main/java/ai/djl/inference/streaming/PublisherBytesSupplier.java
index d83c4678f..d5fdfda87 100644
--- a/api/src/main/java/ai/djl/inference/streaming/PublisherBytesSupplier.java
+++ b/api/src/main/java/ai/djl/inference/streaming/PublisherBytesSupplier.java
@@ -14,13 +14,10 @@ package ai.djl.inference.streaming;
 
 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 @@ import java.util.function.Consumer;
  */
 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,13 +43,24 @@ public class PublisherBytesSupplier implements BytesSupplier {
      * @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();
     }
 
     /**
@@ -62,69 +68,21 @@ public class PublisherBytesSupplier implements BytesSupplier {
      *
      * @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.");
     }
 }


Test patch:
diff --git a/api/src/test/java/ai/djl/inference/streaming/PublisherBytesSupplierTest.java b/api/src/test/java/ai/djl/inference/streaming/PublisherBytesSupplierTest.java
index 8c1406881..a8b2bdfab 100644
--- a/api/src/test/java/ai/djl/inference/streaming/PublisherBytesSupplierTest.java
+++ b/api/src/test/java/ai/djl/inference/streaming/PublisherBytesSupplierTest.java
@@ -15,32 +15,38 @@ package ai.djl.inference.streaming;
 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);
     }
 }