[SPARK-57652][SS][TESTS] Deflake snapshotStartBatchId-with-transformWithState StateDataSource test

[HyukjinKwon]

What changes were proposed in this pull request?

Replace the fixed Thread.sleep(5000) in the "snapshotStartBatchId with transformWithState" test of StateDataSourceTransformWithStateSuite with a deterministic eventually(...) wait that polls until the RocksDB snapshot files the snapshotStartBatchId reader needs (snapshot version 2 for the partitions read) have actually been uploaded by the asynchronous maintenance thread.

The root cause is known, but to guard against regressions the timeout assertion now prints the actual state-directory contents, so a recurrence in scheduled jobs is immediately diagnosable (snapshot still pending vs. cleaned up vs. wrong path) rather than a bare failure.

Why are the changes needed?

The snapshot upload is asynchronous (background maintenance thread), so the fixed sleep is racy under CI load. When it is too short, the snapshot .zip is not yet uploaded and the reader fails on the scheduled Maven (Scala 2.13, JDK 21) and JDK 25 builds:

[CANNOT_LOAD_STATE_STORE.UNCATEGORIZED] An error occurred during loading state.
Caused by: java.io.FileNotFoundException: .../state/0/1/2.zip does not exist

Does this PR introduce any user-facing change?

No. Test-only.

How was this patch tested?

• Before (failing job): Build / Maven (Scala 2.13, JDK 21) → sql#core - slow tests — snapshotStartBatchId with transformWithState ... *** FAILED *** (FileNotFoundException ... 2.zip).
• After (passing, run 10x to confirm the flake is gone): ✅ 10/10 passed — the test was executed 10 consecutive times in one sbt session, all green.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Opus 4.8

[dongjoon-hyun]

+1, LGTM.

[dongjoon-hyun]

So, technically, we increase the maximum timeout from 5s to 60s?

[HyukjinKwon]

Self-review (code-review skill, high effort) — posted as a comment, not an approval.

Scope: 1 file, +21/-2 (test only). Reviewed for correctness, removed-behavior, cross-file, simplification, altitude, conventions.

Correctness ✅ (no bugs found)

• Version is right: the wait polls for snapshot version 2 (2(_.*)?\.zip), which is exactly what the reader needs — snapshotStartBatchId = 1 ⇒ snapshot version 1 + 1 = 2, and matches the original failure path (.../state/0/1/2.zip).
• Partitions are right: it waits on partitions 1 and 4, the same partitions the test reads below (snapshotPartitionId 1 and 4); operatorId = 0 ⇒ state/0. SHUFFLE_PARTITIONS = 5 so those partitions exist.
• Null-safe: Option(partitionDir.listFiles()) handles the dir-not-yet-created case (returns Seq.empty, keeps polling) instead of NPE.
• Regex is anchored: matches is whole-string, so 2(_.*)?\.zip matches 2.zip and 2_<uid>.zip (v1 and checkpoint-v2 naming) but not 12.zip/20.zip — no false positives across the suite subclasses.
• No lost coverage vs the removed Thread.sleep: changelogs are uploaded synchronously at commit, so only the snapshot .zip needed the wait; the new wait covers exactly that, deterministically. Waiting immediately before StopStream also guarantees the file is present when state freezes (strictly better than the sleep).
• Validated by running the test 10× in one sbt session (all green, linked in the description).

Minor (maintainability)

• The snapshot version 2 is hardcoded to track the hardcoded snapshotStartBatchId = 1. If a future edit changes snapshotStartBatchId, this wait silently goes stale (times out / waits on the wrong file). Low risk since both are test-local constants a few lines apart, but a one-line comment tying them together wouldn't hurt.

Optional follow-up (altitude, out of scope here)

• A deeper alternative would be to make the snapshotStartBatchId reader tolerant of a not-yet-uploaded snapshot at the product level, rather than each test waiting. That's a behavior change beyond a test deflake; flagging only as a possible future direction.

No blocking concerns.

[HyukjinKwon]

⚠️ Hold merge — this fix is not robust; I'm pushing a correction.

A full Maven (Scala 2.13, JDK 21) integration run surfaced that the (encoding = avro) variant of this test still fails with the new wait:

snapshotStartBatchId with transformWithState (with changelog checkpointing) (encoding = avro) *** FAILED ***
The code passed to eventually never returned normally. Attempted 608 times over 1.0 minutes.
Last failure message: snapshotUploaded was false Snapshot (version 2) for partition 1 was not uploaded in time.

Root cause is deeper than my first patch assumed: the background maintenance thread snapshots the current version, so a version-2 .zip only ever gets created while version 2 is the current version (between the 2nd and 3rd batches). Waiting for it at the end (when version 5 is current) can never make it appear — so when maintenance didn't happen to snapshot v2 during processing, the wait just times out (a different flake, not a fix).

Correct fix: wait for the version-2 snapshot right after version 2 is committed (while it is still current), which deterministically forces maintenance to create it before more batches advance the version. I'll push that and re-validate all variants (unsaferow + avro) with multiple runs before this should be merged.

(My earlier validation only multi-ran the unsaferow variant — that was the gap.)

[HyukjinKwon]

✅ Updated with a robust fix and re-validated; back to ready for review.

The previous patch was insufficient (the avro variant could still time out, as the integration run showed). The fix now waits for the version-2 snapshot while version 2 is the current version (right after the 2nd batch), via a waitForStateSnapshot(version, partitions) helper — this deterministically forces the maintenance thread to create the snapshot, since it only ever snapshots the current version.

Re-validated across both encodings, 8× (the gap last time was only multi-running unsaferow):
run — 8 consecutive sbt runs, each Tests: succeeded 2, failed 0 (both unsaferow and avro), 16/16 green, including the previously-flaky avro variant.

The on-timeout diagnostic (prints the state-dir contents) is retained for future debuggability.

[HyukjinKwon]

Merged to master and branch-4.x.