micheal-o commented on PR #59144:
URL: https://github.com/apache/spark/pull/59144#issuecomment-6030266239
**This is an AI assisted review. This is an initial review and doesn't mean
I'm fine with this change or approving it.**
## Review Summary
This PR does two things: (1) `StateStore.get`/`getReadOnly` retry once on a
CLOSED->ACQUIRED load failure, using an unpublished replacement provider, and
(2) Cartesian locality prefers the stateful left side. Retrying once is enough,
since nobody else can see or close the unpublished replacement while it loads,
and first-writer-wins holds. I agree with HyukjinKwon that two providers for
the same id coexisting is already part of today's lifecycle. I don't agree that
`close()` blocks until the store is released, though: it waits at most
`lockAcquireTimeoutMs` (120s by default), so the losing replacement can leak.
The new locality walk also re-implements
`DAGScheduler.getPreferredLocsInternal` without its safeguards.
## Required Changes
1. **StateStore.scala:1671-1681 - the losing replacement is queued for close
while its store is still acquired.** Walk-through:
- T1 retries and loads its store from replacement R.
- T2 publishes Q in the meantime.
- T1 sees `contains`, enqueues R, and returns R's store. That store stays
acquired for the rest of the task, because `UnsafeCartesianRDD.compute` buffers
rdd2 and then reads rdd1 lazily.
- Maintenance runs snapshot, then cleanup, then `closeProvider`.
`RocksDBStateMachine.close()` blocks a maintenance thread for up to
`lockAcquireTimeoutMs` and then throws `unreleasedThreadError`.
- The `finally` in `closeProvider` (1805) calls `setUnloaded()` anyway.
The NonFatal handler's second `closeProvider` then sees the provider as
unloaded and returns, so `rocksDB.close()` never runs. R's native memory and
local dir leak for the life of the executor.
The ownership-move path can hit this by chance today; this branch hits it
deterministically for any task that runs longer than the timeout.
Fix options:
- Abort R's store, enqueue R (it is now RELEASED and closes cleanly), and
load through the published Q with `getProvider(false)`.
- Or enqueue R from a task-completion listener instead of right away.
The test at RocksDBStateStoreLockHardeningSuite.scala:430-475 asserts the
current behavior and will need updating.
2. **StateStore.scala:1662 / 1717-1745 - the unpublished provider leaks if
`reportActiveStoreInstance` throws.**
- With `createUnpublishedProvider = true`, `createAndInit` opens RocksDB
eagerly (a native handle plus a local dir).
- `getStateStoreProvider` then makes the `askSync` RPC at 1745, which can
time out.
- If it throws, nothing references the provider: it isn't in
`loadedProviders`, and the handler at 1664-1669 only covers `loadStore`.
Fix: create the unpublished provider in a separate helper (see Code
Design), and close it before rethrowing on any failure after creation.
3. **CartesianRDD.scala:76-107 - drop
`preferredLocationsThroughNarrowDependencies`.**
- It duplicates `DAGScheduler.getPreferredLocsInternal`, but without the
`visited` set (SPARK-695), so a diamond-shaped narrow lineage is explored
exponentially, and it ignores `getCacheLocs`.
- It runs n1*n2 times per stage, and each run makes a coordinator
`askSync` from the DAGScheduler event loop. That is about 40k RPCs per batch at
200x200.
- `CartesianRDD.getDependencies` already lists rdd1 first. So when the
flag is set, return just `rdd1.preferredLocations(s1)`, even if it is `Nil`,
and the scheduler walks the left chain and then falls back to rdd2 by itself.
- Tests can assert through `sc.dagScheduler.getPreferredLocs`.
4. **RocksDBStateStoreLockHardeningSuite.scala:352/410/445/493 - the retry
tests don't reproduce the race as it happens in production.**
- Every test calls `close()` on a provider that is still cached, which no
production path does. In production the provider is first removed by an
ownership move, then closed by maintenance while the task waits in
`acquireStamp`; that is the `else None` branch at 1652.
- Pause instance 1 before `acquireStamp`, then remove and close it via
`StateStore.closeProvider` from the test thread.
- Also add tests for:
- the losing replacement held past `lockAcquireTimeoutMs` (the suite
already uses 2000ms);
- `getProvider(true)` throwing;
- a replacement that fails inside `rocksDB.load`. The current
failing-replacement test throws before `super.getStore`.
## Suggestions
- **StateStore.scala:1650-1662 - the retry always builds a fresh provider,
even when `loadedProviders` already holds a different, healthy one.**
- Interleaving: P is closed, T3 publishes P2, then T1 creates P3. P3 does
a full DFS load and is guaranteed to lose (item 1).
- If the cached instance isn't `firstProvider`, retry through the cache
with `getProvider(false)`, and create an unpublished provider only when the
cache is empty or still holds the stale instance.
- **StateStore.scala:1651-1656 - can the identity-removal branch run in
production?** I couldn't find a production path that leaves a CLOSED provider
in `loadedProviders`: every close path removes it under the lock first. When
the branch does fire, enqueueing an already-closed provider makes
`verifyForMaintenance` throw and logs a spurious "Closing provider due to
error" warning. Please confirm it is reachable, or skip the enqueue when the
provider is already closed.
- **StateStore.scala:1664-1669 - a replacement that failed to load has no
stamp and was never published.** Close it directly on the task thread instead
of queuing it for `MaintenanceOpRequest.All`, which runs snapshot plus DFS
cleanup on an instance that never loaded.
- **StateStore.scala:1658/1667/1679 - none of these enqueues call
`maintenanceTask.triggerNow()`.** The ownership-move path does (1767), so these
providers wait for the next periodic tick.
- **StateStore.scala:1686-1692 - `isClosedProviderLoadTransition` matches on
message-parameter strings.** Renaming a state or an operation would silently
disable the retry. The `newState == "ACQUIRED"` check is also always true for
LOAD. Prefer a typed signal, e.g. a dedicated subclass raised from
`RocksDBStateMachine.acquireStamp`.
- **StateStore.scala:1666 - use `NonFatal` instead of `case error:
Throwable`.** It also shadows the outer `error` from 1649.
- **CartesianProductExec.scala:97 -
`left.exists(_.isInstanceOf[StateStoreWriter])` also matches a writer below an
Exchange.** In that case there is no state store in this stage, yet right-side
locality is still dropped. Stop the traversal at exchanges.
- **CartesianProductExec.scala:97 - only the left side is handled.** For
`static.join(statefulStream)` the stateful input is on the right, its siblings
still scatter, and the query relies on the retry alone. Either handle that side
too or state the limitation.
- **StateStore.scala:1662 - `getProvider(true)` still calls
`reportActiveStoreInstance` for a provider that isn't published yet.** If
another task's report removes this id in the meantime, the provider is
published anyway and stays a non-owner until maintenance evicts it.
- nit: `preferFirstParentLocations` (CartesianRDD.scala:54) and
`preferLeftLocations` (CartesianProductExec.scala:42) name the same thing
differently. Pick one.
- nit: CartesianRDD.scala:91 - this comment in core RDD code refers to SQL's
StateStoreRDD and codegen wrappers. It goes away with item 3.
- nit: CartesianProductExec.scala:96 - the comment just restates the code.
Please remove it.
## Performance
- With the flag on, all `numRight` siblings of a left partition prefer the
same executor and serialize in `acquireStamp` on one provider.
- Each sibling recomputes the stateful partition, and siblings can occupy
executor slots just to wait.
- Any sibling that holds the store longer than `lockAcquireTimeoutMs`
makes the next one fail with `unreleasedThreadError`. Spreading siblings across
executors hides this today.
- Please benchmark with `numRight` at or above the executor's core count,
and report task wait times.
- Each retry does a full RocksDB load from DFS into a new local dir while P
may still hold memory. That is fine for a rare path, but the description should
say so.
- The walk's n1*n2 coordinator RPCs and its exponential behavior on diamond
lineages are covered by item 3.
## Code Design
- The root cause is that the Cartesian join recomputes the stateful left
subtree once per right partition, and each copy commits the same state version.
Locality hints and the retry only reduce the symptoms. A planner-level fix
would remove the cause: broadcast or materialize the non-stateful side, or
don't plan `CartesianProductExec` over a `StateStoreWriter`. Worth discussing
in the JIRA.
- Have you checked this with checkpoint v2
(`enableStateStoreCheckpointIds`)? There, n2 sibling commits of one partition
each report a different unique id, and the first-writer-wins argument depends
on how that resolves.
- StateStore.scala:1710 - the `createUnpublishedProvider` flag turns
`getStateStoreProvider` into "create, don't cache", and the `Boolean =>
(StateStoreProvider, Boolean)` callback in `withClosedProviderRetry` is hard to
follow. Split out a dedicated helper that creates an unpublished provider, and
keep `getStateStoreProvider` for the cached path. That also makes item 2 easy
to fix.
- This retry only applies to RocksDB, but it lives in the provider-agnostic
`StateStore` object.
## Testing
- StreamingCartesianStateStorePlacementSuite.scala:98-113
- `stateStoreLocations` copies the production walk, so the assertion is
circular. Compare against the location the coordinator reports for the provider.
- It checks a hand-built `cartesian.copy(...)`, not the executed plan.
- `local[2]` has a single host, so real placement is never exercised, and
the test name ("independently scheduled tasks") claims more than the test
checks.
- RocksDBStateStoreLockHardeningSuite.scala:465/504 - these assert on
`unloadedProvidersToClose` while `SharedSparkSession`'s maintenance thread can
drain it, which is timing-dependent. Stop maintenance for these tests or assert
through a hook.
- RocksDBStateStoreLockHardeningSuite.scala:42-46, 70
- `instance == 2` depends on a global creation counter, so any extra
`createAndInit` shifts the numbering.
- `loadStarted`/`continueLoad` are plain `var`s read from other threads
without `@volatile`.
- RDDSuite.scala:921-961 - no cases for:
- a `ShuffleDependency` boundary;
- several narrow deps where only the second has locations;
- a diamond lineage.
- No test that a non-stateful `CartesianProductExec` keeps the existing
union behavior.
## PR Hygiene
- The description says the retry "handles the remaining race safely". Please
qualify that and state:
- after a lost publish race, two acquired stores for one partition can
coexist on one executor;
- the retry only applies to RocksDB;
- the right-side stateful case isn't covered;
- the locality flag only changes behavior when the right side has
locations. When it has none (e.g. `spark.range`), the DAGScheduler walk already
lands on the left state-store location.
- These are two independent fixes: a core scheduling hint, and an SS
provider retry. Consider splitting them into separate PRs so core reviewers can
own the CartesianRDD part and each fix can be reverted on its own.
- There is no internal conf kill switch for either behavior change.
- nit: `UnsafeCartesianRDD`'s new constructor parameter has no default. A
`false` default keeps source compatibility for callers.
--
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]
