From 17f0e708df5eef8e4f147ae9c5bbad7481980d9c Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 26 Oct 2026 13:07:52 -0700 Subject: [PATCH] [Java] Fix silent data truncation when mutating partially-cached State iterables In CachingStateIterable, remove() and appendHelper() detect when the cache is partially loaded (nextToken != null) and call cache.remove(IterableCacheKey.INSTANCE). However, without returning early, execution falls through to the block combination logic which stores a MutatedBlocks with nextToken = null, overwriting the cache with truncated contents. This commit adds the missing return statements after cache.remove(...) so that the cache remains evicted and subsequent iterations fetch the full, correct dataset from the runner. --- .../harness/state/StateFetchingIterators.java | 2 + .../fn/harness/state/BagUserStateTest.java | 39 +++++++++++++++++++ .../state/StateFetchingIteratorsTest.java | 12 +++++- 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 1b4d7ec371ba..174136d779b7 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -436,6 +436,7 @@ public void remove(Set 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 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 diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java index 05a280245d13..952fe0a8cd1f 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java +++ b/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 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 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 = diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java index 43087cabbd37..ee57ea6bbc66 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java +++ b/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);