Fix seeking of empty chunkers. (#17830)

Empty chunkers are strange in that they emit one empty chunk and then end. This change makes them properly resetable. Previously, reset() on a empty chunker could result in a NullPointerException. Also, it's important to call close() on the underlying data stream even if it's empty. Closes #17797. PiperOrigin-RevId: 517119206 Change-Id: Iff7908d6cd0633aa2a355ea89f8e647a9fefffcd Co-authored-by: Benjamin Peterson <[email protected]>
bazelbuild · Mar 21, 2023 · 44aa3b8 · 44aa3b8
1 parent ee32eff
commit 44aa3b8
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 5 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/remote/Chunker.java b/src/main/java/com/google/devtools/build/lib/remote/Chunker.java
@@ -159,14 +159,14 @@ public void reset() throws IOException {
   public void seek(long toOffset) throws IOException {
     // For compressed stream, we need to reinitialize the stream since the offset refers to the
     // uncompressed form.
-    if (initialized && toOffset >= offset && !compressed) {
+    if (initialized && size > 0 && toOffset >= offset && !compressed) {
       ByteStreams.skipFully(data, toOffset - offset);
       offset = toOffset;
     } else {
       reset();
       initialize(toOffset);
     }
-    if (data.finished()) {
+    if (size > 0 && data.finished()) {
       close();
     }
   }
@@ -216,7 +216,7 @@ public Chunk next() throws IOException {
     maybeInitialize();
 
     if (size == 0) {
-      data = null;
+      close();
       return emptyChunk;
     }
 

diff --git a/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java b/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java
@@ -84,8 +84,17 @@ public void nextShouldThrowIfNoMoreData() throws IOException {
 
   @Test
   public void emptyData() throws Exception {
-    byte[] data = new byte[0];
-    Chunker chunker = Chunker.builder().setInput(data).build();
+    var inp =
+        new ByteArrayInputStream(new byte[0]) {
+          private boolean closed;
+
+          @Override
+          public void close() throws IOException {
+            closed = true;
+            super.close();
+          }
+        };
+    Chunker chunker = Chunker.builder().setInput(0, inp).build();
 
     assertThat(chunker.hasNext()).isTrue();
 
@@ -96,6 +105,7 @@ public void emptyData() throws Exception {
     assertThat(next.getOffset()).isEqualTo(0);
 
     assertThat(chunker.hasNext()).isFalse();
+    assertThat(inp.closed).isTrue();
 
     assertThrows(NoSuchElementException.class, () -> chunker.next());
   }
@@ -193,6 +203,21 @@ public void seekForwards() throws IOException {
     assertThat(chunker.hasNext()).isFalse();
   }
 
+  @Test
+  public void seekEmptyData() throws IOException {
+    var chunker = Chunker.builder().setInput(new byte[0]).build();
+    for (var i = 0; i < 2; i++) {
+      chunker.seek(0);
+      var next = chunker.next();
+      assertThat(next).isNotNull();
+      assertThat(next.getData()).isEmpty();
+      assertThat(next.getOffset()).isEqualTo(0);
+
+      assertThat(chunker.hasNext()).isFalse();
+      assertThrows(NoSuchElementException.class, chunker::next);
+    }
+  }
+
   @Test
   public void testSingleChunkCompressed() throws IOException {
     byte[] data = {72, 101, 108, 108, 111, 32, 87, 111, 114, 108, 100, 33};