HyukjinKwon opened a new pull request, #57021:
URL: https://github.com/apache/spark/pull/57021
### 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
### Was this patch authored or co-authored using generative AI tooling?
Yes.
Co-authored-by: Isaac
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
