Fix seeking of empty chunkers.

benjaminp · copybara-github · commit cfef67da6349 · 2023-03-16T07:46:42.000-07:00
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
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};

Original file line number	Diff line number	Diff line change
`@@ -159,14 +159,14 @@ public void reset() throws IOException {`
`159`	`159`	`public void seek(long toOffset) throws IOException {`
`160`	`160`	`// For compressed stream, we need to reinitialize the stream since the offset refers to the`
`161`	`161`	`// uncompressed form.`
`162`		`- if (initialized && toOffset >= offset && !compressed) {`
	`162`	`+ if (initialized && size > 0 && toOffset >= offset && !compressed) {`
`163`	`163`	`ByteStreams.skipFully(data, toOffset - offset);`
`164`	`164`	`offset = toOffset;`
`165`	`165`	`} else {`
`166`	`166`	`reset();`
`167`	`167`	`initialize(toOffset);`
`168`	`168`	`}`
`169`		`- if (data.finished()) {`
	`169`	`+ if (size > 0 && data.finished()) {`
`170`	`170`	`close();`
`171`	`171`	`}`
`172`	`172`	`}`
`@@ -216,7 +216,7 @@ public Chunk next() throws IOException {`
`216`	`216`	`maybeInitialize();`
`217`	`217`
`218`	`218`	`if (size == 0) {`
`219`		`- data = null;`
	`219`	`+ close();`
`220`	`220`	`return emptyChunk;`
`221`	`221`	`}`
`222`	`222`