fix(metrics): handle stale DistributedRegistry across SparkContext restarts. - #19790
fix(metrics): handle stale DistributedRegistry across SparkContext restarts.#19790prashantwason wants to merge 1 commit into
Conversation
…starts
Summary:
Intent:
- Fix IllegalStateException("Cannot register an Accumulator twice") when executor metrics are enabled and multiple Hudi write clients are created within the same JVM (e.g. Marmaray processing multiple tables sequentially)
Changes:
- DistributedRegistry.register() now handles the stale-singleton case: when a cached DistributedRegistry was registered to a previous SparkContext, it creates a fresh AccumulatorV2, copies counters, registers it, and swaps it into Registry.REGISTRY_MAP
- Changed register() return type from void to DistributedRegistry so callers use the potentially-replaced instance
- Updated callers in SparkHoodieBackedTableMetadataWriter, SparkHoodieBackedTableMetadataWriterTableVersionSix, and HoodieSparkEngineContext
Test Plan:
Added TestDistributedRegistry.testRegisterIdempotent and testRegisterHandlesStaleAccumulator
Jira Issues: T3-HUDI-9618
---
<sub>Generated by the 🪄 [pr-create](https://sg.uberinternal.com/code.uber.internal/uber-code/devexp-agent-marketplace/-/blob/claude-code/plugins/dev/uber-dev/skills/pr-create/SKILL.md) skill in devexp-agent-marketplace</sub>
|
Could the description reference #19063? This is gap 2 there. Worth a line saying it is the same stale singleton seen from the other side: gap 2's text describes
|
|
@danny0405 @rahil-c can you help review this PR ? |
danny0405
left a comment
danny0405
left a comment
There was a problem hiding this comment.
Two inline findings below. I validated the affected PR classes against Spark 3.5.5: direct stale-accumulator recovery preserves counters and accepts executor updates, but the engine cache can bypass recovery. Both added tests also pass with the recovery block removed. These were focused checks, not a full repository test run.
| return DISTRIBUTED_REGISTRY_MAP.computeIfAbsent(prefixedName, key -> { | ||
| Registry registry = Registry.getRegistryOfClass(tableName, registryName, DistributedRegistry.class.getName()); | ||
| ((DistributedRegistry) registry).register(javaSparkContext); | ||
| return registry; | ||
| return ((DistributedRegistry) registry).register(javaSparkContext); |
There was a problem hiding this comment.
[P1] Recover existing entries in the engine registry cache
DISTRIBUTED_REGISTRY_MAP is static and survives SparkContext restarts, so putting the recovery call inside computeIfAbsent skips it for registries already cached by the previous context. Under the stale-accumulator condition this PR handles, a subsequent lookup therefore still returns the unregistered instance. I reproduced this with the PR classes on Spark 3.5.5: obtain a registry through the first engine context, restart Spark, remove the old accumulator ID from AccumulatorContext to induce the stale state, and look up the same registry through the new engine context. The lookup returns the old instance with isRegistered() == false, and engineContext.map(...) fails with Task not serializable, caused by Accumulator must be registered before send to executor. Please scope/reset this cache with the SparkContext or recover its stale entries before reuse; replacing only Registry.REGISTRY_MAP does not update the map captured by engine operations.
| // When: register() is called with the new SparkContext | ||
| // In local mode, isRegistered() may still return true since AccumulatorContext | ||
| // persists across stop/start. Force the stale path by calling register on jsc2. | ||
| // The key invariant: it must not throw IllegalStateException. | ||
| DistributedRegistry result = staleRegistry.register(jsc2); |
There was a problem hiding this comment.
[P2] Force the stale state so this test exercises recovery
Stopping and recreating the SparkContext does not remove this strongly referenced accumulator from the JVM-wide AccumulatorContext on Spark 3.5.5. Consequently, this call returns from the isRegistered() guard and never exercises the catch block or cache replacement. I ran both added tests successfully, then removed the entire recovery block while retaining the new return type, and both tests still passed. Please explicitly remove staleRegistry.id() from AccumulatorContext and assert !staleRegistry.isRegistered() before calling register(). Then assert that the result is a different instance and that Registry.REGISTRY_MAP.get(cacheKey) is that result, alongside the counter checks, so removing or breaking the fix actually fails the regression test.
| DistributedRegistry fresh = new DistributedRegistry(this.name); | ||
| fresh.counters.putAll(this.counters); | ||
| jsc.sc().register(fresh); | ||
| Registry.REGISTRY_MAP.forEach((key, registry) -> { |
There was a problem hiding this comment.
This swap can only reach Registry.REGISTRY_MAP, but HoodieSparkEngineContext keeps a second cache, DISTRIBUTED_REGISTRY_MAP, which is private static and not visible from here. createWrapperFileSystemRegistries resolves the HoodieWrapperFileSystem registries through getMetricRegistry, so after a swap that path would still hand back the dead accumulator and silently report nothing, even once the metadata table path recovers.
Would it work better to do the detect-and-replace in getMetricRegistry, where both caches are in scope? That would also cover the computeIfAbsent short circuit Danny raised above, since the two look like the same root cause.
There was a problem hiding this comment.
🤖 Agreed that getMetricRegistry is the natural spot — it's the only place both caches are visible (DISTRIBUTED_REGISTRY_MAP directly, and Registry.REGISTRY_MAP via getRegistryOfClass), and createWrapperFileSystemRegistries funnels through it too, so a swap confined to DistributedRegistry.register() can't reach DISTRIBUTED_REGISTRY_MAP. One thing to watch though: the detect-and-replace has to run on the value returned by computeIfAbsent, not inside the lambda — otherwise the short-circuit Danny flagged means an already-cached stale entry is never re-evaluated. A staleness check plus an explicit put/compute on the returned registry would cover both caches and that computeIfAbsent path.
hudi-agent
left a comment
hudi-agent
left a comment
There was a problem hiding this comment.
Thanks for working on this! The PR makes DistributedRegistry.register() recover from a stale accumulator left behind by a dead SparkContext by creating a fresh instance, copying counters, and swapping it into the registry cache. The core recovery logic reads correctly; one concurrency edge case around the new shared-state mutation is worth double-checking, in addition to the DISTRIBUTED_REGISTRY_MAP consistency point already raised by other reviewers. Please take a look at any inline comments, and this should be ready for a Hudi committer or PMC member to take it from here. One nit on a self-undermining test comment; production changes look clean.
|
|
||
| public void register(JavaSparkContext jsc) { | ||
| if (!isRegistered()) { | ||
| public DistributedRegistry register(JavaSparkContext jsc) { |
There was a problem hiding this comment.
🤖 Since the whole point here is a JVM-wide singleton shared across multiple write clients, could two threads race in this method? The isRegistered() check and the catch/swap are unsynchronized, so two concurrent callers on the same stale instance could each build a fresh, both call jsc.sc().register(fresh) (registering two accumulators), and then race on the REGISTRY_MAP swap — leaving one registered accumulator orphaned and callers holding different instances. If concurrent write-client init is possible, would it be worth synchronizing register() (or guarding the swap)? @nsivabalan
| jsc.stop(); | ||
| SparkConf conf = HoodieClientTestUtils.getSparkConfForTest( | ||
| TestDistributedRegistry.class.getSimpleName() + "_stale"); | ||
| JavaSparkContext jsc2 = new JavaSparkContext(conf); |
There was a problem hiding this comment.
🤖 nit: the comment here admits the catch path may never be hit in local mode, which undercuts what the test name promises — could you either rename it to something like testRegisterWithNewSparkContext that reflects what it actually exercises, or add a note explaining how the stale-singleton catch branch is separately validated?
| DistributedRegistry fresh = new DistributedRegistry(this.name); | ||
| fresh.counters.putAll(this.counters); | ||
| jsc.sc().register(fresh); | ||
| Registry.REGISTRY_MAP.forEach((key, registry) -> { |
There was a problem hiding this comment.
🤖 Agreed that getMetricRegistry is the natural spot — it's the only place both caches are visible (DISTRIBUTED_REGISTRY_MAP directly, and Registry.REGISTRY_MAP via getRegistryOfClass), and createWrapperFileSystemRegistries funnels through it too, so a swap confined to DistributedRegistry.register() can't reach DISTRIBUTED_REGISTRY_MAP. One thing to watch though: the detect-and-replace has to run on the value returned by computeIfAbsent, not inside the lambda — otherwise the short-circuit Danny flagged means an already-cached stale entry is never re-evaluated. A staleness check plus an explicit put/compute on the returned registry would cover both caches and that computeIfAbsent path.
7 participants
Describe the issue this Pull Request addresses
DistributedRegistry is a JVM-wide singleton (cached in Registry.REGISTRY_MAP) and a Spark AccumulatorV2. When executor metrics are enabled (hoodie.metrics.executor.enable=true) and multiple Hudi write clients are created within the same JVM (e.g. batch frameworks processing multiple tables sequentially), the cached DistributedRegistry instance retains metadata set from a prior SparkContext registration.
On a subsequent initRegistry() call:
This crashes the write with HoodieException: Failed to instantiate Metadata table.
Summary and Changelog
Fix DistributedRegistry.register() to handle the stale-singleton case where the accumulator was registered to a previous SparkContext that no longer exists.
Impact
No public API or user-facing feature change. No performance impact — the fix only alters behavior in the error path (stale accumulator from a dead SparkContext). The normal registration path is unchanged.
Risk Level
Low. The catch block only fires when AccumulatorV2.register() throws IllegalStateException, which is the exact crash this fixes. The fresh instance is functionally identical to the original, with counters preserved.
Documentation Update
None. No new configs or user-facing changes.
Contributor's checklist