Adds non-blocking poll() for BytesSupplier (#2478)

deepjavalibrary · Mar 24, 2023 · 6f5bb95 · 6f5bb95
1 parent fb304a3
commit 6f5bb95
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/api/src/main/java/ai/djl/modality/ChunkedBytesSupplier.java b/api/src/main/java/ai/djl/modality/ChunkedBytesSupplier.java
@@ -71,6 +71,16 @@ public byte[] nextChunk(long timeout, TimeUnit unit) throws InterruptedException
         return data.getAsBytes();
     }
 
+    /**
+     * Retrieves and removes the head of chunk or returns {@code null} if data is not available.
+     *
+     * @return the head of chunk or returns {@code null} if data is not available
+     */
+    public byte[] poll() {
+        BytesSupplier data = queue.poll();
+        return data == null ? null : data.getAsBytes();
+    }
+
     /** {@inheritDoc} */
     @Override
     public byte[] getAsBytes() {

diff --git a/api/src/test/java/ai/djl/modality/ChunkedBytesSupplierTest.java b/api/src/test/java/ai/djl/modality/ChunkedBytesSupplierTest.java
@@ -24,6 +24,7 @@ public class ChunkedBytesSupplierTest {
     @Test
     public void test() throws InterruptedException, IOException {
         ChunkedBytesSupplier supplier = new ChunkedBytesSupplier();
+        Assert.assertNull(supplier.poll());
         Assert.assertThrows(() -> supplier.nextChunk(1, TimeUnit.MICROSECONDS));
 
         supplier.appendContent(new byte[] {1, 2}, false);
@@ -37,7 +38,7 @@ public void test() throws InterruptedException, IOException {
         data.appendContent(new byte[] {1, 2}, true);
 
         Assert.assertTrue(data.hasNext());
-        Assert.assertEquals(data.nextChunk(1, TimeUnit.MILLISECONDS).length, 0);
+        Assert.assertEquals(data.poll().length, 0);
         Assert.assertEquals(data.nextChunk(1, TimeUnit.MILLISECONDS), new byte[] {1, 2});
 
         Assert.assertFalse(data.hasNext());