shlomitubul opened a new pull request, #3721: URL: https://github.com/apache/celeborn/pull/3721
### What changes were proposed in this pull request? This is the `main`-targeted version of #3706 (originally opened against `branch-0.6`), with the data-loss fix that @SteNicholas's review identified. When the worker-side `celeborn.worker.commitFiles.timeout` fires and `future.cancel(true)` / `task.cancel(true)` interrupt the per-partition commit tasks, `Controller`'s `handleAsync` error branch had three issues: 1. The response was built with empty `committedPrimaryIds` / `committedReplicaIds`, discarding every partition that committed before the timer fired. 2. `context.reply()` was never called, so the driver waited out `celeborn.client.rpc.commitFiles.askTimeout` (e.g. `Cannot receive any reply ... in 300000 milliseconds`) instead of receiving the worker's verdict. 3. The `COMMIT_FILES_TIME` timer was never stopped on this path (latent leak), unlike the success path. This PR: - Extracts the cancellation-path response into a pure, unit-tested `Controller.buildCommitFilesResponseOnCancel`. - Returns `StatusCode.PARTIAL_SUCCESS` with the actually-committed partitions (and their storage infos / bitmaps / sizes) when anything committed; preserves `COMMIT_FILE_EXCEPTION` (all requested failed) when nothing committed. - Calls `context.reply(response)` and `workerSource.stopTimer(...)` on the error path. **The key correctness point** (raised in review of #3706): tasks still **queued**, or interrupted before reaching a terminal state, when cancellation fires land in **none** of the `committed` / `empty` / `failed` sets. So the failed lists are computed as `requested − committed − empty` rather than only the explicitly-failed ids. Otherwise the driver's `CommitHandler.checkDataLost` — which keys only off the failed sets and never diffs against the requested set — cannot distinguish an in-flight (has data, uncommitted) partition from an empty (no data) one, since both are absent from `committed` and `failed`. It would silently treat the in-flight partition as empty-and-valid → wrong reducer results with **no** `FetchFailedException`. Reporting `requested − committed − empty` as failed makes the driver recompute exactly the partitions that didn't make it. ### Why are the changes needed? On heavy shuffles where the per-worker partition count makes `close()` exceed the timeout, the driver today receives no reply (times out at `commitFiles.askTimeout`) and marks the **entire** shuffle data-lost, even when most partitions committed — every reducer then hits `SHUFFLE_DATA_LOST` → `FetchFailedException` → whole-map-stage recompute. With this change the driver gets a definitive `PARTIAL_SUCCESS` with the real committed/failed split, so only the partitions that genuinely didn't commit are recomputed. ### Does this PR introduce any user-facing change? No client API or wire-format change. `StatusCode.PARTIAL_SUCCESS` is already part of the `CommitFilesResponse` protocol and already handled as a terminal, non-retry status by `CommitHandler`. User-visible effect: reduced blast radius when a worker's commit times out under load. ### How was this patch tested? New `ControllerSuite` unit tests for `buildCommitFilesResponseOnCancel`: - timeout with some partitions committed, one empty, and two queued/in-flight → `PARTIAL_SUCCESS`, committed preserved, the two in-flight ids reported as failed, the empty id **not** failed; - nothing committed → `COMMIT_FILE_EXCEPTION` with all requested ids failed. Built and tested with Java 17 / scalafmt (`spotless:check`) clean on `worker`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- 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]
