You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by za...@apache.org on 2026/10/22 04:19:01 UTC
(celeborn) branch main updated: [CELEBORN-2495] Fix memory leak of committedFileInfos when graceful shutdown is disabled
This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 0efa3db8e4 [CELEBORN-2495] Fix memory leak of committedFileInfos when graceful shutdown is disabled
0efa3db8e4 is described below
commit 0efa3db8e41633eb66316560ee1ad4476f02b262
Author: Ubuntu <gu...@gmail.com>
AuthorDate: Thu Oct 22 12:15:07 2026 +0800
[CELEBORN-2495] Fix memory leak of committedFileInfos when graceful shutdown is disabled
JIRA: https://issues.apache.org/jira/browse/CELEBORN-2495
### What changes were proposed in this pull request?
In `StorageManager`:
1. `cleanupExpiredShuffleKey` removes expired entries from `committedFileInfos` after destroying the corresponding writers. Keeping this ordering ensures that an in-flight writer close finishes registering committed metadata before cleanup removes it. Recovery DB deletion remains conditional on graceful shutdown and `cleanDB`.
2. `notifyFileInfoCommitted` only records the committed file into `committedFileInfos` when `workerGracefulShutdown` is enabled, restoring the 0.5.x behavior. This map is only used to persist committed file infos to the recovery DB, so it is useless when graceful shutdown is disabled.
### Why are the changes needed?
With the default configuration (`celeborn.worker.graceful.shutdown.enabled=false`), every committed `DiskFileInfo` is kept forever in `committedFileInfos`:
- `TierWriter.notifyFileCommitted()` (called on every writer `close()`, for both local disk and DFS writers) calls `StorageManager.notifyFileInfoCommitted`, which adds the file info unconditionally.
- `cleanupExpiredShuffleKey` only removes the entry under `if (workerGracefulShutdown)`, so the entries are never released.
This is a regression introduced in 0.6: in 0.5.x, `PartitionDataWriter` only called `notifyFileInfoCommitted` under `if (workerGracefulShutdown)`. That guard was lost when the write path was refactored into `TierWriter`. The issue affects 0.6.1, 0.6.2 and `main`.
Impact: the worker Old Gen grows without bound until `-Xmx`, leading to long Full GCs and eventually OOM. On a production 0.6.1 worker we observed ~14.25 million `DiskFileInfo` instances (and as many `ReduceFileMeta`) retained in the heap, while the disks only held ~6,900 shuffle files and the worker had only 176 active shuffles.
Suggest backporting to `branch-0.6`.
### Does this PR resolve a correctness bug?
- [x] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
Expanded `StorageManagerSuite` to cover:
- Graceful shutdown disabled: committed metadata is not registered, even before shuffle expiry.
- Graceful shutdown enabled: committed metadata is registered and removed on sequential expiry.
- A deterministic concurrent close/cleanup interleaving using the actual `PartitionDataWriter` close/destroy monitor.
Local validation on JDK 17:
```shell
build/mvn -B --no-transfer-progress -pl worker -am test -Dtest=none -DwildcardSuites=org.apache.celeborn.service.deploy.worker.storage.StorageManagerSuite -DtestJarPhase=process-test-classes
```
All 6 tests passed. The concurrent cleanup regression fails against the previous PR head (`2d2aab03b24e045be22d63ffe06ffc77f65427ce`) and passes with the follow-up fix.
`dev/reformat` also completed successfully for all 16 formatter invocations.
Closes #3860 from SoulKyu/fix-committed-fileinfos-leak.
Lead-authored-by: Ubuntu <gu...@gmail.com>
Co-authored-by: zhengtao <sh...@alibaba-inc.com>
Signed-off-by: zhengtao <sh...@alibaba-inc.com>
Co-Authored-By: Codex <no...@openai.com>
AI-Model: gpt-6
AI-Contributed/Feature: 8/8
AI-Contributed/UT: 136/136
---
.../deploy/worker/storage/StorageManager.scala | 8 +-
.../worker/storage/StorageManagerSuite.scala | 136 ++++++++++++++++++++-
2 files changed, 138 insertions(+), 6 deletions(-)
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 4ef0e6f366..418084ab46 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -291,7 +291,8 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
conf.workerGracefulShutdownSaveCommittedFileInfoInterval
private val dbDeleteFailurePolicy =
conf.workerGracefulShutdownDbDeleteFailurePolicy
- private val committedFileInfos =
+ @VisibleForTesting
+ private[storage] val committedFileInfos =
JavaUtils.newConcurrentHashMap[String, ConcurrentHashMap[String, DiskFileInfo]]()
// ShuffleClient can fetch data from a restarted worker only
// when the worker's fetching port is stable.
@@ -741,6 +742,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
}
})
if (workerGracefulShutdown) {
+ // Writer cleanup waits for in-flight close calls to finish registering metadata.
committedFileInfos.remove(shuffleKey)
if (cleanDB) {
try {
@@ -1071,7 +1073,9 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
shuffleKey: String,
fileName: String,
fileInfo: DiskFileInfo): Unit = {
- committedFileInfos.computeIfAbsent(shuffleKey, diskFileInfoMapFunc).put(fileName, fileInfo)
+ if (workerGracefulShutdown) {
+ committedFileInfos.computeIfAbsent(shuffleKey, diskFileInfoMapFunc).put(fileName, fileInfo)
+ }
}
def getActiveShuffleSize: Long = {
diff --git a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
index 0e2a8b4200..57505e115c 100644
--- a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
+++ b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
@@ -20,6 +20,8 @@ package org.apache.celeborn.service.deploy.worker.storage
import java.{lang, util}
import java.io.IOException
import java.nio.file.Files
+import java.util.concurrent.{Callable, ConcurrentHashMap, CountDownLatch, Executors, TimeUnit}
+import java.util.concurrent.atomic.AtomicInteger
import scala.collection.JavaConverters._
@@ -28,15 +30,15 @@ import org.apache.hadoop.fs.FileSystem
import org.mockito.{Mockito, MockitoSugar}
import org.mockito.ArgumentMatchersSugar.any
import org.mockito.invocation.InvocationOnMock
-import org.mockito.stubbing.Stubber
+import org.mockito.stubbing.{Answer, Stubber}
import org.apache.celeborn.CelebornFunSuite
import org.apache.celeborn.common.CelebornConf
-import org.apache.celeborn.common.CelebornConf.{WORKER_DISK_RESERVE_SIZE, WORKER_GRACEFUL_SHUTDOWN_ENABLED, WORKER_GRACEFUL_SHUTDOWN_RECOVER_PATH, WORKER_STORAGE_DIRS}
+import org.apache.celeborn.common.CelebornConf.{WORKER_DISK_RESERVE_SIZE, WORKER_GRACEFUL_SHUTDOWN_ENABLED, WORKER_GRACEFUL_SHUTDOWN_RECOVER_PATH, WORKER_GRACEFUL_SHUTDOWN_SAVE_COMMITTED_FILEINFO_INTERVAL, WORKER_STORAGE_DIRS}
import org.apache.celeborn.common.identity.UserIdentifier
import org.apache.celeborn.common.meta.{DiskFileInfo, DiskInfo, DiskStatus}
-import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType, StorageInfo}
-import org.apache.celeborn.common.util.Utils
+import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMode, PartitionType, StorageInfo}
+import org.apache.celeborn.common.util.{CelebornExitKind, JavaUtils, Utils}
import org.apache.celeborn.service.deploy.worker.WorkerSource
trait MockitoHelper extends MockitoSugar {
@@ -55,6 +57,103 @@ class StorageManagerSuite extends CelebornFunSuite with MockitoHelper {
storageManager.saveAllCommittedFileInfosToDB()
}
+ test("[CELEBORN-2495] committedFileInfos should not leak when graceful shutdown is disabled") {
+ withDiskFile(gracefulShutdown = false) { (storageManager, fileInfo, _) =>
+ storageManager.notifyFileInfoCommitted("app-1-0", "0-0-0", fileInfo)
+ assert(storageManager.committedFileInfos.isEmpty)
+ storageManager.cleanupExpiredShuffleKey(storageManager.shuffleKeySet())
+ assert(storageManager.committedFileInfos.isEmpty)
+ }
+ }
+
+ test("[CELEBORN-2495] committedFileInfos should be removed after expiry with graceful shutdown enabled") {
+ withDiskFile(gracefulShutdown = true) { (storageManager, fileInfo, _) =>
+ storageManager.notifyFileInfoCommitted("app-1-0", "0-0-0", fileInfo)
+ assert(storageManager.committedFileInfos.get("app-1-0").get("0-0-0") eq fileInfo)
+ storageManager.cleanupExpiredShuffleKey(storageManager.shuffleKeySet())
+ assert(storageManager.committedFileInfos.isEmpty)
+ assert(storageManager.shuffleKeySet().isEmpty)
+ }
+ }
+
+ test(
+ "[CELEBORN-2495] cleanup should remove metadata committed while waiting for a writer to close") {
+ withDiskFile(gracefulShutdown = true) { (storageManager, fileInfo, conf) =>
+ val manager = spy(storageManager)
+ val closing = new CountDownLatch(1)
+ val cleanupStarted = new CountDownLatch(1)
+ val finishClose = new CountDownLatch(1)
+ val tier = mock[TierWriterBase]
+ val policy = mock[StoragePolicy]
+ val constructorManager = mock[StorageManager]
+ when(constructorManager.storagePolicy).thenReturn(policy)
+ when(policy.createFileWriter(
+ any[PartitionDataWriterContext],
+ any[AtomicInteger],
+ any[FlushNotifier])).thenReturn(tier)
+ Mockito.doAnswer(new Answer[java.lang.Long] {
+ override def answer(invocation: InvocationOnMock): java.lang.Long = {
+ closing.countDown()
+ assert(finishClose.await(30, TimeUnit.SECONDS))
+ manager.notifyFileInfoCommitted("app-1-0", "0-0-0", fileInfo)
+ 0L
+ }
+ }).when(tier).close()
+
+ // Exercise the real close/destroy monitor while controlling when commit completes.
+ val context = new PartitionDataWriterContext(
+ Long.MaxValue,
+ PartitionSplitMode.HARD,
+ false,
+ genPartitionLocation(0, Array(0L)),
+ "app-1",
+ 0,
+ new UserIdentifier("t1", "u1"),
+ PartitionType.REDUCE,
+ false,
+ false)
+ val writer = new PartitionDataWriter(
+ constructorManager,
+ mock[WorkerSource],
+ conf,
+ DeviceMonitor.EmptyMonitor(),
+ context)
+ val writers = new ConcurrentHashMap[String, PartitionDataWriter]()
+ writers.put(fileInfo.getFilePath, writer)
+ manager.workingDirWriters.put(
+ fileInfo.getFile.getParentFile.getParentFile.getParentFile,
+ writers)
+ Mockito.doAnswer(new Answer[java.lang.Boolean] {
+ override def answer(invocation: InvocationOnMock): java.lang.Boolean = {
+ cleanupStarted.countDown()
+ invocation.callRealMethod().asInstanceOf[java.lang.Boolean]
+ }
+ }).when(manager).cleanFileInternal("app-1-0", fileInfo, failOnHdfsDeleteError = false)
+
+ val expiredShuffleKeys = manager.shuffleKeySet()
+ val pool = Executors.newFixedThreadPool(2)
+ try {
+ val closeResult = pool.submit(new Callable[java.lang.Long] {
+ override def call(): java.lang.Long = writer.close()
+ })
+ assert(closing.await(30, TimeUnit.SECONDS))
+ val cleanupResult = pool.submit(new Runnable {
+ override def run(): Unit = manager.cleanupExpiredShuffleKey(expiredShuffleKeys)
+ })
+ assert(cleanupStarted.await(30, TimeUnit.SECONDS))
+ finishClose.countDown()
+ closeResult.get(30, TimeUnit.SECONDS)
+ cleanupResult.get(30, TimeUnit.SECONDS)
+ assert(manager.shuffleKeySet().isEmpty)
+ assert(manager.committedFileInfos.isEmpty)
+ } finally {
+ finishClose.countDown()
+ pool.shutdownNow()
+ pool.awaitTermination(30, TimeUnit.SECONDS)
+ }
+ }
+ }
+
test("updateDiskInfosWithDiskReserveSize") {
// reserve size set to 5g
val conf = new CelebornConf().set(WORKER_DISK_RESERVE_SIZE, Utils.byteStringAsBytes("5g"))
@@ -217,6 +316,35 @@ class StorageManagerSuite extends CelebornFunSuite with MockitoHelper {
}
}
+ private def withDiskFile(gracefulShutdown: Boolean)(
+ run: (StorageManager, DiskFileInfo, CelebornConf) => Unit): Unit = {
+ val root = Files.createTempDirectory("committed-file-info-")
+ val recoveryDir = Files.createDirectory(root.resolve("recovery"))
+ val conf = new CelebornConf()
+ .set(WORKER_STORAGE_DIRS, Seq(root.resolve("data").toString))
+ .set(WORKER_GRACEFUL_SHUTDOWN_ENABLED, gracefulShutdown)
+ .set(WORKER_GRACEFUL_SHUTDOWN_RECOVER_PATH, recoveryDir.toString)
+ .set(WORKER_GRACEFUL_SHUTDOWN_SAVE_COMMITTED_FILEINFO_INTERVAL, TimeUnit.HOURS.toMillis(1))
+ val storageManager = new StorageManager(conf, new WorkerSource(conf))
+ try {
+ storageManager.localDisksSnapshot().foreach(_.setUsableSpace(Long.MaxValue))
+ val (_, fileInfo, _) = storageManager.createDiskFile(
+ genPartitionLocation(0, Array(0L)),
+ "app-1",
+ 0,
+ "0-0-0",
+ new UserIdentifier("t1", "u1"),
+ PartitionType.REDUCE,
+ partitionSplitEnabled = false)
+ run(storageManager, fileInfo, conf)
+ } finally {
+ storageManager.cleanupExpiredShuffleKey(storageManager.shuffleKeySet(), cleanDB = false)
+ storageManager.committedFileInfos.clear()
+ storageManager.close(CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN)
+ JavaUtils.deleteRecursively(root.toFile)
+ }
+ }
+
private def genPartitionLocation(epoch: Int, offsets: Array[Long]): PartitionLocation = {
val location: PartitionLocation =
new PartitionLocation(0, epoch, "localhost", 0, 0, 0, 0, PartitionLocation.Mode.PRIMARY)