Fix silent data truncation when mutating partially-cached StateIterables

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java

@@ -436,6 +436,7 @@ public void remove(Set<Object> toRemoveStructuralValues) {
       // returning invalid results as part of a future request.
       if (existing.getBlocks().get(existing.getBlocks().size() - 1).getNextToken() != null) {
         cache.remove(IterableCacheKey.INSTANCE);
+        return;
       }
 
       // Combine all the individual blocks into one block containing all the values since
@@ -550,6 +551,7 @@ private void appendHelper(List<T> newValues, long newWeight) {
       // returning invalid results as part of a future request.
       if (existing.getBlocks().get(existing.getBlocks().size() - 1).getNextToken() != null) {
         cache.remove(IterableCacheKey.INSTANCE);
+        return;
       }
 
       // Combine all the individual blocks into one block containing all the values since

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java

@@ -196,6 +196,45 @@ public void testAppendCached() throws Exception {
     assertEquals(encode("A1", "A2", "A3", "A4", "A5"), fakeClient.getData().get(key("A")));
   }
 
+  @Test
+  public void testPartialIterationThenAppendDoesNotTruncateCache() throws Exception {
+    // Initial data contains multiple elements across multiple chunks.
+    FakeBeamFnStateClient fakeClient =
+        new FakeBeamFnStateClient(
+            StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList("A1", "A2", "A3", "A4")), 2);
+
+    Cache<?, ?> cache = Caches.eternal();
+    {
+      // First user state: simulates BagState.isEmpty().read() or partial iteration, followed by
+      // append.
+      BagUserState<String> userState =
+          new BagUserState<>(cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of());
+
+      // Calling iterator().hasNext() only loads the first chunk into the cache with nextToken !=
+      // null.
+      assertTrue(userState.get().iterator().hasNext());
+
+      // User appends a new element.
+      userState.append("A5");
+
+      // asyncClose persists the append to runner and updates the cache.
+      userState.asyncClose();
+    }
+
+    {
+      // The next user state accesses the state for the same key.
+      BagUserState<String> userState =
+          new BagUserState<>(cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of());
+
+      // All elements (A1, A2, A3, A4, A5) must be present; unread pages (A2, A3, A4) must not be
+      // lost.
+      assertArrayEquals(
+          new String[] {"A1", "A2", "A3", "A4", "A5"},
+          Iterables.toArray(userState.get(), String.class));
+      userState.asyncClose();
+    }
+  }
+
   @Test
   public void testClear() throws Exception {
     FakeBeamFnStateClient fakeClient =

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java

@@ -145,8 +145,12 @@ public void testAppend() throws Exception {
 
       // Test on a partially accessed iterable
       iterable = create(5, expected);
-      iterable.append(Ints.asList(42, 43));
+      // Access the iterator to load the first block into the cache so that the cache is partially
+      // populated with nextToken != null before calling append. Without this call, the cache is
+      // empty and append would return early at the existing == null check instead of exercising
+      // the partial cache eviction logic.
       Boolean ignored = iterable.iterator().hasNext();
+      iterable.append(Ints.asList(42, 43));
       // The append is dropped since the entire iterable isn't in memory so next time we iterate we
       // expect to be given what the runner tells us is there.
       verifyFetch(iterable.iterator(), expected);
@@ -176,8 +180,12 @@ public void testRemove() throws Exception {
 
       // Test on a partially accessed iterable
       iterable = create(5, expected);
-      iterable.remove(toRemove);
+      // Access the iterator to load the first block into the cache so that the cache is partially
+      // populated with nextToken != null before calling remove. Without this call, the cache is
+      // empty and remove would return early at the existing == null check instead of exercising
+      // the partial cache eviction logic.
       Boolean ignored = iterable.iterator().hasNext();
+      iterable.remove(toRemove);
       // The remove is dropped since the entire iterable isn't in memory so next time we iterate we
       // expect to be given what the runner tells us is there.
       verifyFetch(iterable.iterator(), expected);