[SPARK-57993][SS][TEST][branch-4.x] Fix flaky RocksDBStateStoreIntegrationSuite bounded memory usage test by isolating global RocksDBMemoryManager#57025
Closed
HyukjinKwon wants to merge 1 commit into
Conversation
…ounded memory usage test by isolating global RocksDBMemoryManager
### What changes were proposed in this pull request?
This fixes the flaky `bounded memory usage calculation` test in
`RocksDBStateStoreIntegrationSuite` (and, most visibly, its
`RocksDBStateStoreIntegrationSuiteWithRowChecksum` variant) by unloading any state store
providers left over from the previous test before the test touches the process-global
`RocksDBMemoryManager` singleton:
```scala
// Fully unload state store providers left over from previous tests before touching the
// global RocksDBMemoryManager singleton. ...
StateStore.stop()
// Clear any existing providers from previous tests
RocksDBMemoryManager.resetWriteBufferManagerAndCache
```
### Why are the changes needed?
The test asserts, after starting a bounded-memory query, that the global manager holds no
*unbounded* instances:
```scala
eventually(timeout(Span(10, Seconds)), interval(Span(500, Millis))) {
assert(RocksDBMemoryManager.getNumRocksDBInstances(true) == 2)
assert(RocksDBMemoryManager.getNumRocksDBInstances(false) == 0) // <- flakes: "1 did not equal 0"
}
```
The immediately-preceding test, `RocksDB memory tracking integration with UnifiedMemoryManager
with boundedMemory=false`, starts a rate-stream query that registers **unbounded** RocksDB
instances in the process-global `RocksDBMemoryManager` (an `object` singleton). That test calls
`query.stop()` in its `finally`, but `StreamTest` does **not** force state-store teardown between
tests — only `afterAll()` calls `StateStore.stop()`. A lingering unbounded provider is unregistered
from `RocksDBMemoryManager` only when its `RocksDBStateStoreProvider.close()` runs (→ `rocksDB.close()`
→ `RocksDBMemoryManager.unregisterInstance()`), which happens asynchronously after `query.stop()`.
So a leftover unbounded provider can call `updateMemoryUsage(...)` (re-registering itself) **after**
this test's `RocksDBMemoryManager.resetWriteBufferManagerAndCache`, leaving a stale unbounded instance
in the singleton and failing the `getNumRocksDBInstances(false) == 0` assertion with `1 did not equal 0`.
This reproduces far more reliably under the **row-checksum** variant, whose extra per-row work keeps
the prior query's providers busy slightly longer and widens the async-teardown window past the existing
10s `eventually` budget (added by SPARK-55993).
Calling `StateStore.stop()` first forces any residual providers to `close()` (and therefore
`unregisterInstance()`) **synchronously**, so the global singleton is clean before the bounded-memory
query starts. This mirrors what several other state-store suites already do (e.g. `StreamingJoinSuite`,
`MultiStatefulOperatorsSuite` call `StateStore.stop()` in setup) and what `StreamTest.afterAll()` does
between suites.
### Does this PR introduce _any_ user-facing change?
No — test-only change.
### How was this patch tested?
The `Build / Maven (Scala 2.13, JDK 21, ARM)` `sql#core` lane (where this test flakes) was run
repeatedly on a fork via GitHub Actions.
**Before (red — apache/spark `master`):**
- `Build / Maven (Scala 2.13, JDK 21, ARM)` → `sql#core` FAILED with
`RocksDBStateStoreIntegrationSuiteWithRowChecksum ... bounded memory usage calculation ... 1 did not equal 0`:
https://github.com/apache/spark/actions/runs/28745745887
**After (green — with this change, fork verification):**
- `sql#core - other tests` job PASSED (contains `RocksDBStateStoreIntegrationSuite` +
`...WithRowChecksum`): https://github.com/HyukjinKwon/spark/actions/runs/28757237515
- Independent second green run: https://github.com/HyukjinKwon/spark/actions/runs/28757240766
Closes apache#57021 from HyukjinKwon/ci-fix/agent2-rocksdb-bounded-mem-isolation.
Authored-by: Hyukjin Kwon <gurwls223@apache.org>
Signed-off-by: Hyukjin Kwon <hyukjin.kwon@databricks.com>
(cherry picked from commit 05728c8)
Signed-off-by: Hyukjin Kwon <hyukjin.kwon@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Backport of #57021 / SPARK-57993 to
branch-4.x.What changes were proposed in this pull request?
Cherry-pick (
-x) of the merged master commit 05728c8, which fixes the flakybounded memory usage calculationtest inRocksDBStateStoreIntegrationSuite(most visibly itsRocksDBStateStoreIntegrationSuiteWithRowChecksumvariant) by callingStateStore.stop()before touching the process-globalRocksDBMemoryManagersingleton, so leftover unbounded providers from the preceding test are unregistered synchronously.Applies cleanly to
branch-4.x(no conflicts).Why are the changes needed?
Same flake exists on this maintenance branch's CI lanes. See #57021 / SPARK-57993 for the full root-cause analysis.
Does this PR introduce any user-facing change?
No — test-only.
How was this patch tested?
Clean cherry-pick of an already-validated master change (validated on the ARM
sql#corelane in #57021). Existing test.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code