[PR #1128] [CLOSED] [TE] fix(rdma): pre-compute slice_count to avoid race condition

No labels

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/Mooncake#2028

bsakaguchi commented

2026-10-29 21:32:13 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/Mooncake/pull/1128
Author: @wwq2333
Created: 11/27/2025
Status: ❌ Closed

Base: main ← Head: fix-rdma-slice-count-race

📝 Commits (3)

d1d1c24 [TE] fix(rdma): pre-compute slice_count to avoid race condition
7165eab [TE] fix(rdma): simplify slice count calculation by merging redundant conditions
de84bfb [TE] fix(rdma): unify slice counting and allocation logic

📊 Changes

3 files changed (+118 additions, -45 deletions)

View changed files

📝 mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h (+20 -0)
📝 mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp (+1 -1)
📝 mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp (+97 -44)

📄 Description

Description

Previously, task.slice_count was incremented inside the slicing loop using __sync_fetch_and_add(). This caused a race condition under USE_EVENT_DRIVEN_COMPLETION (#1053): when some slices completed quickly, completed_slice_count could equal slice_count before all slices were created, causing premature task completion detection. Similarly, any other logic that needs to access slice_count could also be affected by this issue.

This commit:

Pre-computes the total slice count before entering the slicing loop based on request.length, kBlockSize, and kFragmentSize (matching the existing merge logic)
Removes the incremental __sync_fetch_and_add() call inside the loop
Removes incorrect manual slice cleanup in the error path (TransferTask destructor already handles slice deallocation, manual cleanup caused double-free)

Type of Change

Types
- Bug fix
- New feature
  - Transfer Engine
  - Mooncake Store
  - Mooncake EP
  - Integration
  - P2P Store
  - Python Wheel
- Breaking change
- CI/CD
- Documentation update
- Other

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have updated the documentation.
I have added tests to prove my changes are effective.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/Mooncake/pull/1128 **Author:** [@wwq2333](https://github.com/wwq2333) **Created:** 11/27/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix-rdma-slice-count-race` --- ### 📝 Commits (3) - [`d1d1c24`](https://github.com/kvcache-ai/Mooncake/commit/d1d1c2413dd42d95f47658eda3d70b9cc94ebb96) [TE] fix(rdma): pre-compute slice_count to avoid race condition - [`7165eab`](https://github.com/kvcache-ai/Mooncake/commit/7165eab9632a266592ca2b8ff7a94229b94e5a74) [TE] fix(rdma): simplify slice count calculation by merging redundant conditions - [`de84bfb`](https://github.com/kvcache-ai/Mooncake/commit/de84bfbcbafe5b7dc583086a3959b28801a34f76) [TE] fix(rdma): unify slice counting and allocation logic ### 📊 Changes **3 files changed** (+118 additions, -45 deletions) <details> <summary>View changed files</summary> 📝 `mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h` (+20 -0) 📝 `mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp` (+1 -1) 📝 `mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp` (+97 -44) </details> ### 📄 Description ## Description Previously, `task.slice_count` was incremented inside the slicing loop using `__sync_fetch_and_add()`. This caused a race condition under USE_EVENT_DRIVEN_COMPLETION (#1053): when some slices completed quickly, `completed_slice_count` could equal `slice_count` before all slices were created, causing premature task completion detection. Similarly, any other logic that needs to access `slice_count` could also be affected by this issue. This commit: - Pre-computes the total slice count before entering the slicing loop based on request.length, kBlockSize, and kFragmentSize (matching the existing merge logic) - Removes the incremental __sync_fetch_and_add() call inside the loop - Removes incorrect manual slice cleanup in the error path (TransferTask destructor already handles slice deallocation, manual cleanup caused double-free) ## Type of Change * Types - [X] Bug fix - [ ] New feature - [ ] Transfer Engine - [ ] Mooncake Store - [ ] Mooncake EP - [ ] Integration - [ ] P2P Store - [ ] Python Wheel - [ ] Breaking change - [ ] CI/CD - [ ] Documentation update - [ ] Other ## How Has This Been Tested?  ## Checklist - [X] I have performed a self-review of my own code. - [ ] I have updated the documentation. - [ ] I have added tests to prove my changes are effective. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

closed this issue
added the
pull-request
label

Rows
Columns

[PR #1128] [CLOSED] [TE] fix(rdma): pre-compute slice_count to avoid race condition #2028