This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 05728c8fa397 [SPARK-57993][SS][TEST] Fix flaky
RocksDBStateStoreIntegrationSuite bounded memory usage test by isolating global
RocksDBMemoryManager
05728c8fa397 is described below
commit 05728c8fa3978e3ba6d57e68c7388d22e02904ea
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Mon Jul 6 12:23:04 2026 +0900
[SPARK-57993][SS][TEST] Fix flaky RocksDBStateStoreIntegrationSuite bounded
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 #57021 from HyukjinKwon/ci-fix/agent2-rocksdb-bounded-mem-isolation.
Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
.../streaming/state/RocksDBStateStoreIntegrationSuite.scala | 11 +++++++++++
1 file changed, 11 insertions(+)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala
index 03e23e6b466c..4060e44ef1ae 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala
@@ -398,6 +398,17 @@ class RocksDBStateStoreIntegrationSuite extends StreamTest
(SQLConf.SHUFFLE_PARTITIONS.key -> "2"), // Use 2 partitions to test
multiple providers
(s"${RocksDBConf.ROCKSDB_SQL_CONF_NAME_PREFIX}.boundedMemoryUsage" ->
"true")) {
+ // Fully unload state store providers left over from previous tests
before touching the
+ // global RocksDBMemoryManager singleton. A preceding test may still
be asynchronously
+ // tearing down its (unbounded) query when this test starts; its
providers only call
+ // RocksDBMemoryManager.unregisterInstance() from
RocksDBStateStoreProvider.close(), which
+ // StreamTest does not force between tests (only afterAll() calls
StateStore.stop()). If a
+ // lingering unbounded provider re-reports its memory usage after the
reset below, it stays
+ // in the singleton and the "0 unbounded instances" assertion flakes
(observed as
+ // "1 did not equal 0", especially under the row-checksum variant
whose extra per-row work
+ // widens the teardown window). Stopping first makes
close()/unregister run synchronously.
+ StateStore.stop()
+
// Clear any existing providers from previous tests
RocksDBMemoryManager.resetWriteBufferManagerAndCache
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
