acrites opened a new pull request, #40739:
URL: https://github.com/apache/beam/pull/40739
# 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:
```java
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`:
```java
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]