This is an automated email from the ASF dual-hosted git repository.
danny0405 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new b6fe7c3fd8b8 fix(metrics): handle stale DistributedRegistry across
SparkContext restarts (#19790)
b6fe7c3fd8b8 is described below
commit b6fe7c3fd8b86d3d62b8bd74fdfcc1dfa183f316
Author: Prashant Wason <[email protected]>
AuthorDate: Tue Oct 6 19:15:30 2026 -0700
fix(metrics): handle stale DistributedRegistry across SparkContext restarts
(#19790)
DistributedRegistry is a JVM-wide singleton cached in Registry.REGISTRY_MAP
and a
Spark AccumulatorV2. After a SparkContext restart in the same JVM,
register()
threw "Cannot register an Accumulator twice". Detect the stale state
explicitly,
register a fresh accumulator that shares the stale instance's counters, and
swap
it into REGISTRY_MAP and HoodieSparkEngineContext's
DISTRIBUTED_REGISTRY_MAP.
---
.../client/common/HoodieSparkEngineContext.java | 30 ++++--
.../SparkHoodieBackedTableMetadataWriter.java | 4 +-
...ieBackedTableMetadataWriterTableVersionSix.java | 4 +-
.../apache/hudi/metrics/DistributedRegistry.java | 36 +++++++-
.../hudi/metrics/TestDistributedRegistry.java | 102 +++++++++++++++++++++
5 files changed, 161 insertions(+), 15 deletions(-)
diff --git
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/client/common/HoodieSparkEngineContext.java
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/client/common/HoodieSparkEngineContext.java
index 97317b61e490..77eff393353c 100644
---
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/client/common/HoodieSparkEngineContext.java
+++
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/client/common/HoodieSparkEngineContext.java
@@ -137,7 +137,7 @@ public class HoodieSparkEngineContext extends
HoodieEngineContext {
@Override
public <I, O> List<O> map(List<I> data, SerializableFunction<I, O> func, int
parallelism) {
- final Map<String, Registry> registries = DISTRIBUTED_REGISTRY_MAP;
+ final Map<String, Registry> registries = getDistributedRegistries();
return javaSparkContext.parallelize(data, parallelism).map(i -> {
setRegistries(registries);
return func.apply(i);
@@ -146,7 +146,7 @@ public class HoodieSparkEngineContext extends
HoodieEngineContext {
@Override
public <I, K, V> List<V> mapToPairAndReduceByKey(List<I> data,
SerializablePairFunction<I, K, V> mapToPairFunc, SerializableBiFunction<V, V,
V> reduceFunc, int parallelism) {
- final Map<String, Registry> registries = DISTRIBUTED_REGISTRY_MAP;
+ final Map<String, Registry> registries = getDistributedRegistries();
return javaSparkContext.parallelize(data, parallelism).mapToPair(input -> {
setRegistries(registries);
Pair<K, V> pair = mapToPairFunc.call(input);
@@ -158,7 +158,7 @@ public class HoodieSparkEngineContext extends
HoodieEngineContext {
public <I, K, V> Stream<ImmutablePair<K, V>>
mapPartitionsToPairAndReduceByKey(
Stream<I> data, SerializablePairFlatMapFunction<Iterator<I>, K, V>
flatMapToPairFunc,
SerializableBiFunction<V, V, V> reduceFunc, int parallelism) {
- final Map<String, Registry> registries = DISTRIBUTED_REGISTRY_MAP;
+ final Map<String, Registry> registries = getDistributedRegistries();
return javaSparkContext.parallelize(data.collect(Collectors.toList()),
parallelism)
.mapPartitionsToPair((PairFlatMapFunction<Iterator<I>, K, V>) iterator
-> {
setRegistries(registries);
@@ -179,7 +179,7 @@ public class HoodieSparkEngineContext extends
HoodieEngineContext {
@Override
public <I, O> List<O> flatMap(List<I> data, SerializableFunction<I,
Stream<O>> func, int parallelism) {
- final Map<String, Registry> registries = DISTRIBUTED_REGISTRY_MAP;
+ final Map<String, Registry> registries = getDistributedRegistries();
return javaSparkContext.parallelize(data, parallelism).flatMap(x -> {
setRegistries(registries);
return func.apply(x).iterator();
@@ -193,7 +193,7 @@ public class HoodieSparkEngineContext extends
HoodieEngineContext {
@Override
public <I, K, V> Map<K, V> mapToPair(List<I> data,
SerializablePairFunction<I, K, V> func, Integer parallelism) {
- final Map<String, Registry> registries = DISTRIBUTED_REGISTRY_MAP;
+ final Map<String, Registry> registries = getDistributedRegistries();
if (Objects.nonNull(parallelism)) {
return javaSparkContext.parallelize(data, parallelism).mapToPair(input
-> {
setRegistries(registries);
@@ -281,13 +281,25 @@ public class HoodieSparkEngineContext extends
HoodieEngineContext {
@Override
public Registry getMetricRegistry(String tableName, String registryName) {
final String prefixedName = tableName.isEmpty() ? registryName : tableName
+ "." + registryName;
- return DISTRIBUTED_REGISTRY_MAP.computeIfAbsent(prefixedName, key -> {
- Registry registry = Registry.getRegistryOfClass(tableName, registryName,
DistributedRegistry.class.getName());
- ((DistributedRegistry) registry).register(javaSparkContext);
- return registry;
+ // Registers on every lookup, not only on first insert, so that an entry
cached under a previous
+ // SparkContext is recovered too.
+ return DISTRIBUTED_REGISTRY_MAP.compute(prefixedName, (key, cached) -> {
+ Registry registry = cached != null
+ ? cached : Registry.getRegistryOfClass(tableName, registryName,
DistributedRegistry.class.getName());
+ return ((DistributedRegistry) registry).register(javaSparkContext);
});
}
+ /**
+ * Returns the distributed registries to ship to executors, first recovering
any entry registered to a
+ * previous SparkContext: the map is static and outlives SparkContext
restarts, and an unregistered
+ * accumulator fails task serialization.
+ */
+ private Map<String, Registry> getDistributedRegistries() {
+ DISTRIBUTED_REGISTRY_MAP.replaceAll((key, registry) ->
((DistributedRegistry) registry).register(javaSparkContext));
+ return DISTRIBUTED_REGISTRY_MAP;
+ }
+
/**
* Accumulator-backed registries owned by this context rather than by the
process, keyed by table base
* path. Unlike {@link #getMetricRegistry}, these are never published into
{@code Registry.REGISTRY_MAP}:
diff --git
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriter.java
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriter.java
index ec21dac28ea7..27a442f50fcb 100644
---
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriter.java
+++
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriter.java
@@ -119,9 +119,9 @@ public class SparkHoodieBackedTableMetadataWriter extends
HoodieBackedTableMetad
if (metadataWriteConfig.isMetricsOn()) {
Registry registry;
if (metadataWriteConfig.isExecutorMetricsEnabled() &&
metadataWriteConfig.getMetricsReporterType() != MetricsReporterType.INMEMORY) {
- registry = Registry.getRegistry("HoodieMetadata",
DistributedRegistry.class.getName());
+ DistributedRegistry distRegistry = (DistributedRegistry)
Registry.getRegistry("HoodieMetadata", DistributedRegistry.class.getName());
HoodieSparkEngineContext sparkEngineContext =
(HoodieSparkEngineContext) engineContext;
- ((DistributedRegistry)
registry).register(sparkEngineContext.getJavaSparkContext());
+ registry =
distRegistry.register(sparkEngineContext.getJavaSparkContext());
} else {
registry = Registry.getRegistry("HoodieMetadata");
}
diff --git
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriterTableVersionSix.java
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriterTableVersionSix.java
index 8190a5112d59..337115c4b897 100644
---
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriterTableVersionSix.java
+++
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkHoodieBackedTableMetadataWriterTableVersionSix.java
@@ -97,9 +97,9 @@ public class
SparkHoodieBackedTableMetadataWriterTableVersionSix extends HoodieB
if (metadataWriteConfig.isMetricsOn()) {
Registry registry;
if (metadataWriteConfig.isExecutorMetricsEnabled() &&
metadataWriteConfig.getMetricsReporterType() != MetricsReporterType.INMEMORY) {
- registry = Registry.getRegistry("HoodieMetadata",
DistributedRegistry.class.getName());
+ DistributedRegistry distRegistry = (DistributedRegistry)
Registry.getRegistry("HoodieMetadata", DistributedRegistry.class.getName());
HoodieSparkEngineContext sparkEngineContext =
(HoodieSparkEngineContext) engineContext;
- ((DistributedRegistry)
registry).register(sparkEngineContext.getJavaSparkContext());
+ registry =
distRegistry.register(sparkEngineContext.getJavaSparkContext());
} else {
registry = Registry.getRegistry("HoodieMetadata");
}
diff --git
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java
index c8b9fdde6458..da7d42ea6aee 100644
---
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java
+++
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java
@@ -41,6 +41,8 @@ public class DistributedRegistry extends
AccumulatorV2<Map<String, Long>, Map<St
private final String name;
ConcurrentHashMap<String, Long> counters = new ConcurrentHashMap<>();
+ // The registered accumulator that replaced this one if it went stale, see
register().
+ private transient volatile DistributedRegistry replacement;
/** Driver-only, to detect a SparkContext restart in the same JVM (shells,
notebooks, Spark Connect). */
private transient String registeredAppId;
@@ -53,13 +55,43 @@ public class DistributedRegistry extends
AccumulatorV2<Map<String, Long>, Map<St
return name;
}
- public void register(JavaSparkContext jsc) {
- if (!isRegistered()) {
+ /**
+ * Registers this accumulator with the given SparkContext if it is not
registered yet.
+ *
+ * <p>A cached instance that was registered to a previous SparkContext can
no longer be registered
+ * (AccumulatorV2 rejects a second registration). In that case a fresh
accumulator sharing the same
+ * counters is registered instead and swapped into {@link
Registry#REGISTRY_MAP}. Callers should use the
+ * returned instance; updates made through the stale instance are still
reported via the shared counters.
+ *
+ * @return this instance, or its registered replacement.
+ */
+ public synchronized DistributedRegistry register(JavaSparkContext jsc) {
+ if (isRegistered()) {
+ return this;
+ }
+ if (registeredAppId == null) {
jsc.sc().register(this);
// Only when this call actually registers: stamping unconditionally
would re-brand an accumulator
// bound to a dead context and mask the staleness this field exists to
detect.
this.registeredAppId = jsc.sc().applicationId();
+ return this;
+ }
+ // Stale: registered to a previous SparkContext that no longer exists, so
the id is gone from the
+ // current AccumulatorContext but AccumulatorV2 refuses to register this
instance again.
+ if (replacement != null) {
+ // Another cache still holds this stale instance after it was already
replaced. Delegate, so a
+ // replacement that went stale in turn (another SparkContext restart) is
replaced the same way.
+ replacement = replacement.register(jsc);
+ return replacement;
}
+ DistributedRegistry fresh = new DistributedRegistry(this.name);
+ // Share the counters rather than copying them, so updates made through a
caller that still holds
+ // this stale instance are reported by the replacement too.
+ fresh.counters = this.counters;
+ fresh.register(jsc);
+ Registry.REGISTRY_MAP.replaceAll((key, registry) -> registry == this ?
fresh : registry);
+ replacement = fresh;
+ return fresh;
}
/** False when bound to a different (typically stopped) context, meaning it
must be recreated. */
diff --git
a/hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/metrics/TestDistributedRegistry.java
b/hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/metrics/TestDistributedRegistry.java
index 6058eeef3522..928ca6de052b 100644
---
a/hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/metrics/TestDistributedRegistry.java
+++
b/hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/metrics/TestDistributedRegistry.java
@@ -23,12 +23,14 @@ import org.apache.hudi.common.metrics.Registry;
import org.apache.hudi.testutils.HoodieClientTestUtils;
import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.util.AccumulatorContext;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@@ -330,4 +332,104 @@ public class TestDistributedRegistry {
Assertions.assertTrue(countsWithoutPrefix.containsKey(METRIC_1));
Assertions.assertTrue(countsWithoutPrefix.containsKey(METRIC_2));
}
+
+ @Test
+ public void testRegisterIdempotent() {
+ // Given: a DistributedRegistry registered to the SparkContext
+ String registryName = REGISTRY_NAME + "_testIdempotent";
+ DistributedRegistry registry = new DistributedRegistry(registryName);
+ registry.add(METRIC_1, 42);
+ DistributedRegistry result = registry.register(jsc);
+
+ // Then: first registration returns the same instance
+ Assertions.assertSame(registry, result);
+ Assertions.assertTrue(registry.isRegistered());
+
+ // When: register() is called again on the same SparkContext
+ DistributedRegistry result2 = registry.register(jsc);
+
+ // Then: it is a no-op, returns same instance
+ Assertions.assertSame(registry, result2);
+ }
+
+ @Test
+ public void testRegisterHandlesStaleAccumulator() {
+ // Given: a registered registry cached in the global REGISTRY_MAP
+ String registryName = REGISTRY_NAME + "_testStale";
+ DistributedRegistry staleRegistry = new DistributedRegistry(registryName);
+ staleRegistry.add(METRIC_1, 100);
+ staleRegistry.add(METRIC_2, 200);
+ String cacheKey = Registry.makeKey("", registryName);
+ Registry.REGISTRY_MAP.put(cacheKey, staleRegistry);
+ try {
+ Assertions.assertSame(staleRegistry, staleRegistry.register(jsc));
+
+ // Given: the accumulator went stale, as after a SparkContext restart
once the previous context
+ // dropped it: it is no longer in the AccumulatorContext but can't be
registered again
+ makeStale(staleRegistry);
+
+ // When: register() is called
+ DistributedRegistry result = staleRegistry.register(jsc);
+
+ // Then: a registered replacement carrying the counters is returned and
cached
+ Assertions.assertNotSame(staleRegistry, result);
+ Assertions.assertTrue(result.isRegistered());
+ Assertions.assertEquals(registryName, result.getName());
+ Map<String, Long> counts = result.getAllCounts();
+ Assertions.assertEquals(100, counts.get(METRIC_1));
+ Assertions.assertEquals(200, counts.get(METRIC_2));
+ Assertions.assertSame(result, Registry.REGISTRY_MAP.get(cacheKey));
+
+ // Then: another holder of the stale instance gets the same replacement
+ Assertions.assertSame(result, staleRegistry.register(jsc));
+
+ // Then: updates through the stale instance are reported by the
replacement
+ staleRegistry.add(METRIC_1, 5);
+ Assertions.assertEquals(105, result.getAllCounts().get(METRIC_1));
+
+ // Given: the replacement went stale too, as after another SparkContext
restart
+ makeStale(result);
+
+ // When: a holder of the original stale instance calls register()
+ DistributedRegistry result2 = staleRegistry.register(jsc);
+
+ // Then: a new registered replacement sharing the counters is returned
and cached
+ Assertions.assertNotSame(result, result2);
+ Assertions.assertTrue(result2.isRegistered());
+ Assertions.assertEquals(105, result2.getAllCounts().get(METRIC_1));
+ Assertions.assertSame(result2, Registry.REGISTRY_MAP.get(cacheKey));
+ Assertions.assertSame(result2, result.register(jsc));
+ } finally {
+ Registry.REGISTRY_MAP.remove(cacheKey);
+ }
+ }
+
+ @Test
+ public void testEngineContextRecoversStaleRegistry() {
+ // Given: a registry obtained through the engine context, which caches it
in a static map
+ String registryName = REGISTRY_NAME + "_testEngineStale";
+ DistributedRegistry staleRegistry = (DistributedRegistry)
engineContext.getMetricRegistry("", registryName);
+ staleRegistry.add(METRIC_1, 7);
+ try {
+ // Given: the cached accumulator went stale
+ makeStale(staleRegistry);
+
+ // Then: engine operations ship only registered accumulators to executors
+ Assertions.assertEquals(Arrays.asList(2, 4),
engineContext.map(Arrays.asList(1, 2), i -> i * 2, 2));
+
+ // Then: lookups return the registered replacement with the counters
+ DistributedRegistry result = (DistributedRegistry)
engineContext.getMetricRegistry("", registryName);
+ Assertions.assertNotSame(staleRegistry, result);
+ Assertions.assertTrue(result.isRegistered());
+ Assertions.assertEquals(7, result.getAllCounts().get(METRIC_1));
+ Assertions.assertSame(result,
Registry.REGISTRY_MAP.get(Registry.makeKey("", registryName)));
+ } finally {
+ Registry.REGISTRY_MAP.remove(Registry.makeKey("", registryName));
+ }
+ }
+
+ private static void makeStale(DistributedRegistry registry) {
+ AccumulatorContext.remove(registry.id());
+ Assertions.assertFalse(registry.isRegistered());
+ }
}
