[acrites]

Summary of Changes and Bug Fixes

Ran an audit on WindmillMap.java, WindmillSet.java, and WindmillMapViaMultimap.java, which identified a few logical errors. None of the files had unit tests, so I added some basic tests, as well as tests identifying all of the specific bugs mentioned below.

---

1. WindmillSet.java

Bug: addIfAbsent() Always Returned false

• Issue: Apache Beam's SetState.addIfAbsent(value) contract specifies:

  Returns a ReadableState<Boolean> that reads true if the element was successfully added (i.e. was not already in the set), or false if the element was already present.

• Root Cause: The underlying Windmill primitive TagBag.putIfAbsent(key, false) returns null if the item was absent and added, or the existing boolean (false) if already present. WindmillSet used:

  return Optional.ofNullable(putState.read()).orElse(false);

  When the item was absent and inserted, putState.read() returned null, which defaulted via orElse(false) to false. When already present, it returned false. Consequently, addIfAbsent().read() always evaluated to false.
• Fix: Updated read() to check whether the previous state was null:

  return putState.read() == null;

---

2. WindmillMapViaMultimap.java

Bug: computeIfAbsent() Returned Inserted Value Instead of null

• Issue: Beam's MapState.computeIfAbsent(key, mappingFunction) contract specifies:

  If a value is computed and inserted, returns null or a state reading null (matching putIfAbsent). Otherwise returns the existing value.

• Root Cause: On cache misses, WindmillMapViaMultimap computed the value, inserted it into the multimap, and returned ReadableStates.immediate(inserted) instead of null.
• Fix: Corrected the insertion branch to return ReadableStates.immediate(null). When the key is already present, it continues to return the existing value via Iterators.getOnlyElement(existingValues).

---

3. WindmillMap.java

Bug 1: Key Identity vs. Coder Structural Equality (e.g. byte[] Keys)

• Issue: State internal collections (cachedValues, negativeCache, localAdditions, localRemovals) used raw user keys (Map<K, V>, Set<K>).
• Root Cause: Keys whose equality is defined structurally by their coder (such as byte[], array types, or objects relying on structural serialization) do not implement equals() and hashCode() based on their content. For example, two distinct byte[] arrays with identical contents were treated as distinct keys, resulting in missed cache hits, duplicate entries, and incorrect state lookups.
• Fix: Introduced StructuralKey<K> which wraps originalKey and delegates equals() and hashCode() to keyCoder.structuralValue(originalKey). All internal tracking maps and sets now operate on StructuralKey<K>.

Bug 2: put() Premature Early Return on null Values and Overwritten Removals

• Issue: put(key, value) had two flaws:
  1. For NullableCoder, calling put(key, null) on an empty map caused oldValue to be null. Since Objects.equals(null, null) == true, put() early-returned without recording the key insertion into cachedValues or localAdditions.
  2. If a key was pending removal in localRemovals or marked in negativeCache, re-putting an equal value was skipped, leaving the key marked for deletion.
• Fix: put() now explicitly requires wasPresent = cachedValues.containsKey(structKey) and ensures !localRemovals.contains(structKey) && !negativeCache.contains(structKey) before performing an early return. It also removes the key from localRemovals and negativeCache upon mutation.

Bug 3: computeIfAbsent() Unnecessary State Dirtying on Storage Reads

• Issue: When a key was absent from cachedValues but present in Windmill persistent storage (persistedValue != null), computeIfAbsent() previously invoked this.put(key, persistedValue).
• Root Cause: Calling put() added the key to localAdditions. During commit (persistDirectly()), unmodified data fetched from storage was unnecessarily re-serialized and written back to Windmill.
• Fix: Unified absent-key computation in computeAndPut(). When an existing key is read from persistent storage, it is placed directly into cachedValues.put(structKey, persistedValue) without dirtying localAdditions.

Bug 4: getOrDefault() Ignored negativeCache, localRemovals, and complete

• Issue:
  1. getOrDefault(key, defaultValue).read() dispatched an asynchronous fetch (getFutureForKey(key)) even when the key was known to be absent (negativeCache, localRemovals, or complete == true).
  2. If the key was found in localRemovals or negativeCache, it returned null instead of defaultValue.
• Fix: Added fast-paths in getOrDefault().read() that return defaultValue immediately when present in localRemovals, negativeCache, or when complete without the key in cachedValues. Updated readLater() to avoid prefetching known absent/cached keys.

Bug 5: Key Resurrection / Filter Leak in entries().read() with Weighted

• Issue: In entries().read() under the Weighted branch, after processing persistent entries, remaining cached keys were iterated from cachedValues.keySet(). If a key was in cachedValues but marked in localRemovals, it was mistakenly emitted.
• Fix: Added !localRemovals.contains(sk) checks to prevent deleted keys from resurfacing. In the non-weighted branch, reused userEntries() to avoid duplicate conversion logic.

Bug 6: Dead / Broken Key Serialization in persistDirectly()

• Issue: In persistDirectly(), the loop over localRemovals attempted to serialize keys to keyStream, but the line keyCoder.encode(...) was commented out. It generated empty byte strings (ByteString.EMPTY) and never appended them to the commit bundle.
• Fix: Removed the broken and dead serialization block, keeping only the correct commitBuilder.addTagValueDeleteRequests() calls.

Bug 7: isEmpty() Optimization

• Issue: isEmpty() was performing unnecessary checks when cachedValues already contained active, non-removed entries.
• Fix: Added a fast-path with documentation: if !cachedValues.isEmpty(), the map is guaranteed to be non-empty (as any removal removes the key from cachedValues).

---

4. Test Suites Added

Three dedicated test suites were created with comprehensive coverage for all identified edge cases and contract requirements:

1. WindmillSetTest.java:

   • testAddIfAbsentReturnsTrueWhenAbsent: Verifies addIfAbsent() returns true on initial addition.
   • testAddIfAbsentReturnsFalseWhenPresent: Verifies addIfAbsent() returns false when the element already exists.
   • testAddAndRead: Verifies standard add() and read() behavior.
   • testRemove: Verifies element removal.
   • testContains: Verifies contains() state checks.
   • testClear: Verifies set clearing.
   • testPersistDirectly: Verifies mutation serialization to Windmill.WorkItemCommitRequest.
   • testAddIfAbsentWithByteArray: Tests structural equality for byte array elements.
   • testReadEmpty: Verifies reading an empty set.

2. WindmillMapTest.java:

   • testByteArrayKeyStructuralEquality: Verifies distinct byte[] instances with identical contents correctly hit the cache, removals, and negative cache.
   • testPutNullValueWithNullableCoderDoesNotEarlyReturn: Verifies put(key, null) on an empty map is not dropped as a no-op.
   • testPutSameValueAfterRemovalReAdds: Verifies putting an identical value after remove() unmarks removal and adds the key.
   • testComputeIfAbsentReturnsNullOnAbsent: Verifies computeIfAbsent() returns null when a key is newly computed and inserted.
   • testComputeIfAbsentReturnsExistingValueWhenPresent: Verifies computeIfAbsent() returns the existing value on key collision.
   • testComputeIfAbsentDoesNotDirtyLocalAdditionsWhenPresentInStorage: Verifies fetching existing keys from storage does not mark them as local additions during commit.
   • testGetOrDefaultReturnsDefaultWhenRemoved: Verifies getOrDefault() returns default rather than null for removed keys.
   • testGetOrDefaultReturnsDefaultWhenInNegativeCache: Verifies getOrDefault() returns default for negative-cached keys.
   • testGetOrDefaultReturnsDefaultWhenComplete: Verifies getOrDefault() returns default when map is fully read.
   • testEntriesWeightedDoesNotResurrectRemovedKeys: Verifies weighted entries().read() does not emit keys pending in localRemovals.
   • testBasicPutAndGet: Basic map read/write.
   • testRemoveAndContainsKey: Verifies removal and containsKey() checks.
   • testIsEmpty: Verifies empty and non-empty assertions.
   • testClear: Verifies clearing all map entries.
   • testValuesAndKeysIterables: Verifies values().read() and keys().read().
   • testPersistDirectly: Verifies commit building for inserts, updates, and deletes.

3. WindmillMapViaMultimapTest.java:

   • testComputeIfAbsentReturnsNullOnAbsent: Verifies contract compliance (null returned when inserted).
   • testComputeIfAbsentReturnsExistingValueWhenPresent: Verifies existing value returned on collision.
   • testSingleValueIterableAdaptorGetOnlyElement: Confirms Iterators.getOnlyElement succeeds on single-element values.
   • testSingleValueIterableAdaptorThrowsOnMultiValuedKey: Confirms IllegalArgumentException is raised when multiple values exist for a single key.
   • testBasicPutAndGet: Verifies basic put/get through the multimap backing.
   • testRemoveAndContainsKey: Verifies key removal.
   • testClearAndIsEmpty: Verifies clear and empty checks.

[github-actions]

📋 Review Summary

This PR fixes several important contract and state-tracking bugs across WindmillSet, WindmillMapViaMultimap, and WindmillMap (including SetState.addIfAbsent always returning false, MapState.computeIfAbsent returning the inserted value instead of null, StructuralKey equality for array/custom keys, and cache/removal handling in getOrDefault, entries, and put), accompanied by new unit test suites.

🔍 General Feedback

• The fixes in WindmillSet, WindmillMapViaMultimap, and WindmillMap are well-targeted and resolve clear contract violations.
• One subtle edge case remains in WindmillMap.entries().read() on the non-weighted (lazy) path: calling userEntries() directly reads the live cachedValues map instead of the cachedCopy snapshot, which can cause duplicate entries or ConcurrentModificationException if get(k) or put(k, v) is called before/during iteration of the returned Iterable.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

💬 Inline Findings

• runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMap.java (line 418): 🟠 High: In the non-weighted branch of entries().read(), userEntries() builds its list from the live cachedValues at the time entries().read() is called, whereas Iterables.filter checks !cachedCopy.containsKey(structKey) lazily when the returned Iterable is iterated. However, if any changes were meant to be isolated to cachedCopy—or more importantly, userEntries() is called while cachedCopy is already constructed on line 410—Wait: MapState specifies that subsequent mutations should not affect iterables already returned by entries().read(). While userEntries() snapshots cachedValues into a List eagerly, cachedCopy on line 410 is a redundant copy of cachedValues. Even more subtly, in the Weighted branch above (line 404), cachedValues.putIfAbsent(structKey, e.getValue()) is used: if a key was locally mapped to null (via NullableCoder, which this PR explicitly supports in put), putIfAbsent will overwrite the local null value with the stale non-null value from Windmill storage because Map.putIfAbsent treats a null value as absent! Using !cachedValues.containsKey(structKey) before cachedValues.put(structKey, e.getValue()) on line 404 avoids overwriting locally stored null values.

• runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMap.java (line 404): 🟠 High: Map.putIfAbsent overwrites existing entries whose value is null. Since this PR adds support for null values with NullableCoder (tested in testPutNullValueSupported), if a user calls map.put(k, null) to overwrite a persisted non-null value with null and then calls entries().read(), cachedValues.putIfAbsent(structKey, e.getValue()) will overwrite the cached null with the old persisted value from Windmill. Check !cachedValues.containsKey(structKey) instead.

                if (!localRemovals.contains(structKey) && !cachedValues.containsKey(structKey)) {
                  cachedValues.put(structKey, e.getValue());
                }

• runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMap.java (line 284): 🟡 Medium: In WindmillMap.remove(K key) (lines 293-296), if (localRemovals.add(structKey)) only removes structKey from cachedValues and localAdditions if structKey was not already in localRemovals. However, in WindmillMap.entries().read() or get().read() prior to this PR (or if localRemovals already contained structKey), if localRemovals is not cleared on put... put does remove from localRemovals, so localRemovals and cachedValues are disjoint unless computeIfAbsent or another path populates cachedValues without clearing localRemovals. Still, removing from cachedValues and localAdditions unconditionally (or ensuring negativeCache.remove(structKey) is also called on remove) is safer and cleaner.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 75.97%. Comparing base (76f6e6d) to head (01f523a).
⚠️ Report is 298 commits behind head on master.

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #40739       +/-   ##
=============================================
+ Coverage     56.16%   75.97%   +19.80%     
+ Complexity     2288     2273       -15     
=============================================
  Files          1124      259      -865     
  Lines        178336    14377   -163959     
  Branches       1489     1479       -10     
=============================================
- Hits         100170    10923    -89247     
+ Misses        75645     2763    -72882     
+ Partials       2521      691     -1830     

Flag	Coverage Δ
go	?
java	75.97% <ø> (+0.18%)	⬆️
python	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

Assigning reviewers:

R: @Abacn added as fallback since no labels match configuration
Backup reviewer: kennknowles (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).