[SPARK-53581][CORE] Fix potential thread-safety issue for mapTaskIds.add() #52337

Ngone51 · 2025-09-15T08:08:55Z

What changes were proposed in this pull request?

This a followup for #47037. This PR wraps up the synchronize lock for invocation of OpenHashSet.add() in IndexShuffleBlockResolver.

Why are the changes needed?

OpenHashSet is not thread safe. We should enfoce synchronize lock when invokes the add function to ensure the thread-safety.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

N/A.

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

No.

LuciferYang

@Ngone51 We should create a new Jira ticket because SPARK-46957 has already been released in Spark 4.0.

Ngone51 · 2025-09-15T12:55:51Z

@LuciferYang Thanks for the advice. Created a new ticket SPARK-53581.

core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala

attilapiros · 2025-09-15T20:22:16Z

Now as the PR title changed should not the FOLLOW-UP need to be removed?

Ngone51 · 2025-09-16T01:36:11Z

@attilapiros Thanks. Updated.

attilapiros

LGTM

dongjoon-hyun

+1, LGTM.

…add() ### What changes were proposed in this pull request? This a followup for #47037. This PR wraps up the synchronize lock for invocation of `OpenHashSet.add()` in `IndexShuffleBlockResolver`. ### Why are the changes needed? `OpenHashSet` is not thread safe. We should enfoce synchronize lock when invokes the add function to ensure the thread-safety. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? N/A. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #52337 from Ngone51/fix-thread-safety. Authored-by: Yi Wu <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]> (cherry picked from commit 78aba00) Signed-off-by: Dongjoon Hyun <[email protected]>

dongjoon-hyun · 2025-09-16T22:17:53Z

Thank you, @Ngone51 and all.

Merged to master/4.0/3.5 for

Apache Spark 4.1.0-preview2
Apache Spark 4.0.2
Apache Spark 3.5.7

cc @peter-toth as the release manager of Apache Spark 3.5.7, too.

Ngone51 · 2025-09-17T03:05:35Z

Thanks all!

mridulm · 2025-09-17T03:27:34Z

I am wondering if SortShuffleManager.unregisterShuffle is also impacted by this ...

computeIfAbsent in IndexShuffleBlockResolver.onComplete -> remove in SortShuffleManager.unregisterShuffle -> mapTaskIds.add while iterating over mapTaskIds in unregisterShuffle

I am entirely sure if this combination is possible, but would be a good idea to make SortShuffleManager.unregisterShuffle lock mapTaskIds as well given we have identified the issue in this PR.

Thoughts @Ngone51 ?

Ngone51 · 2025-09-17T03:36:03Z

@mridulm Good catch! I'll make a separate PR to fix it later.

LuciferYang · 2025-09-17T03:47:45Z

I'd like to expand on this issue. Is there a possibility of the reverse timing sequence for the operations mentioned by @mridulm? If SortShuffleManager.unregisterShuffle occurs first, and then the computeIfAbsent in IndexShuffleBlockResolver.onComplete happens afterward, could this lead to a minor memory leak?

Ngone51 · 2025-09-17T04:09:36Z

If SortShuffleManager.unregisterShuffle occurs first, and then the computeIfAbsent in IndexShuffleBlockResolver.onComplete happens afterward, could this lead to a minor memory leak?

@LuciferYang Possible, but I think that's fine. Because even if we have ensured the thread-safety already, this kind of issue could persist due to the race betwen shuffle removal and shuffle adding. The critial issue on this thread-safety problem is that the state of OpenHashSet could be corrupted, leading to incorrect data accessing result.

…egisterShuffle ### What changes were proposed in this pull request? This PR fixes the thread-safetye issue in `SortShuffleManager.unregisterShuffle` by enforcing synchronous lock on `mapTaskIds`'s iteration. Besides, this PR also addresses the [concern](#52337 (comment)) to enfore the type of `taskIdMapsForShuffle` as `ConcurrentHashMap` to ensure its thread-safety. ### Why are the changes needed? Fix the potential thread-safety issue as pointed at #52337 (comment) and also the [concern](#52337 (comment)). ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? N/A ### Was this patch authored or co-authored using generative AI tooling? No. Closes #52386 from Ngone51/fix. Authored-by: Yi Wu <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

…dd() ### What changes were proposed in this pull request? This a followup for apache#47037. This PR wraps up the synchronize lock for invocation of `OpenHashSet.add()` in `IndexShuffleBlockResolver`. ### Why are the changes needed? `OpenHashSet` is not thread safe. We should enfoce synchronize lock when invokes the add function to ensure the thread-safety. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? N/A. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#52337 from Ngone51/fix-thread-safety. Authored-by: Yi Wu <[email protected]> (cherry picked from commit 78aba00) Signed-off-by: Dongjoon Hyun <[email protected]> Co-authored-by: Yi Wu <[email protected]>

…add() ### What changes were proposed in this pull request? This a followup for apache#47037. This PR wraps up the synchronize lock for invocation of `OpenHashSet.add()` in `IndexShuffleBlockResolver`. ### Why are the changes needed? `OpenHashSet` is not thread safe. We should enfoce synchronize lock when invokes the add function to ensure the thread-safety. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? N/A. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#52337 from Ngone51/fix-thread-safety. Authored-by: Yi Wu <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]> (cherry picked from commit baae014) Signed-off-by: Dongjoon Hyun <[email protected]>

…add() ### What changes were proposed in this pull request? This a followup for apache#47037. This PR wraps up the synchronize lock for invocation of `OpenHashSet.add()` in `IndexShuffleBlockResolver`. ### Why are the changes needed? `OpenHashSet` is not thread safe. We should enfoce synchronize lock when invokes the add function to ensure the thread-safety. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? N/A. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#52337 from Ngone51/fix-thread-safety. Authored-by: Yi Wu <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

…egisterShuffle ### What changes were proposed in this pull request? This PR fixes the thread-safetye issue in `SortShuffleManager.unregisterShuffle` by enforcing synchronous lock on `mapTaskIds`'s iteration. Besides, this PR also addresses the [concern](apache#52337 (comment)) to enfore the type of `taskIdMapsForShuffle` as `ConcurrentHashMap` to ensure its thread-safety. ### Why are the changes needed? Fix the potential thread-safety issue as pointed at apache#52337 (comment) and also the [concern](apache#52337 (comment)). ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? N/A ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#52386 from Ngone51/fix. Authored-by: Yi Wu <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

fix

686d58d

github-actions bot added the CORE label Sep 15, 2025

Ngone51 requested review from LuciferYang, attilapiros and mridulm September 15, 2025 08:09

LuciferYang approved these changes Sep 15, 2025

View reviewed changes

LuciferYang requested changes Sep 15, 2025

View reviewed changes

Ngone51 changed the title ~~[SPARK-46957][FOLLOW-UP][CORE] Fix potential thread-safety issue for mapTaskIds.add()~~ [SPARK-53581][FOLLOW-UP][CORE] Fix potential thread-safety issue for mapTaskIds.add() Sep 15, 2025

attilapiros reviewed Sep 15, 2025

View reviewed changes

core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala Show resolved Hide resolved

Ngone51 changed the title ~~[SPARK-53581][FOLLOW-UP][CORE] Fix potential thread-safety issue for mapTaskIds.add()~~ [SPARK-53581][CORE] Fix potential thread-safety issue for mapTaskIds.add() Sep 16, 2025

cloud-fan approved these changes Sep 16, 2025

View reviewed changes

LuciferYang approved these changes Sep 16, 2025

View reviewed changes

attilapiros approved these changes Sep 16, 2025

View reviewed changes

dongjoon-hyun approved these changes Sep 16, 2025

View reviewed changes

dongjoon-hyun closed this in 78aba00 Sep 16, 2025

Ngone51 mentioned this pull request Sep 18, 2025

[SPARK-53636][CORE] Fix thread-safety issue in SortShuffleManager.unregisterShuffle #52386

Closed

[SPARK-53581][CORE] Fix potential thread-safety issue for mapTaskIds.add() #52337

[SPARK-53581][CORE] Fix potential thread-safety issue for mapTaskIds.add() #52337

Uh oh!

Conversation

Ngone51 commented Sep 15, 2025

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

LuciferYang left a comment

Choose a reason for hiding this comment

Uh oh!

Ngone51 commented Sep 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

attilapiros commented Sep 15, 2025

Uh oh!

Ngone51 commented Sep 16, 2025

Uh oh!

attilapiros left a comment

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun commented Sep 16, 2025

Uh oh!

Ngone51 commented Sep 17, 2025

Uh oh!

mridulm commented Sep 17, 2025

Uh oh!

Ngone51 commented Sep 17, 2025

Uh oh!

LuciferYang commented Sep 17, 2025

Uh oh!

Ngone51 commented Sep 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Ngone51 commented Sep 15, 2025 •

edited

Loading