Skip to content

[Transfer Engine] Bound TCP connection pools with FIFO backpressure#2974

Draft
jacklin78911-collab wants to merge 2 commits into
kvcache-ai:mainfrom
jacklin78911-collab:te-2930-backpressure
Draft

[Transfer Engine] Bound TCP connection pools with FIFO backpressure#2974
jacklin78911-collab wants to merge 2 commits into
kvcache-ai:mainfrom
jacklin78911-collab:te-2930-backpressure

Conversation

@jacklin78911-collab

@jacklin78911-collab jacklin78911-collab commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR addresses the bounded-pool and backpressure part of #2930.

The existing TCP connection pool can create an unbounded number of
connections to the same peer when no idle connection is available. This
change introduces:

  • a configurable per-peer connection cap;
  • FIFO backpressure for acquisition requests above the cap;
  • explicit tracking of in-progress connection attempts;
  • deterministic shutdown handling for queued and pending acquisitions.

The cap counts both established and in-progress connections:

established connections + pending connects <= per-peer cap

Scope

This PR implements Phase A of #2930:

  • bounded per-peer TCP connection pools;
  • FIFO queuing when the per-peer cap is reached;
  • resuming queued acquisitions when a connection becomes reusable or
    capacity is released;
  • safe pending-connect ownership detachment during shutdown;
  • acquisition callbacks executed outside the pool mutex.

This PR intentionally does not implement:

  • stale or half-closed pooled-socket detection after a peer restart;
  • transfer or payload progress deadlines;
  • timeout of queued waiters;
  • a global bound on the number of peers;
  • changes to transfer-status counter synchronization.

Design

Each peer pool contains:

  • established pooled connections;
  • FIFO acquisition waiters;
  • explicit PendingConnect tokens for in-progress resolve/connect attempts.

When a pooled connection is requested:

  1. An idle open connection is reused if available.
  2. Otherwise, if the combined established and pending count is below the cap,
    a pending-connect slot is reserved and an asynchronous resolve/connect is
    started.
  3. If the cap is already reached, the acquisition callback is appended to the
    per-peer FIFO waiter queue.
  4. Returning, discarding, or failing a connection releases capacity and drives
    the oldest queued waiter.

Each pending-connect token records its peer, attempt ID, callback, resolver,
socket, and lifecycle stage. Resolve/connect handlers validate token identity,
stage, and shutdown state before completing the attempt.

Acquisition callbacks are never invoked while holding the connection-pool
mutex.

Waiter follow-up work is dispatched through an iterative trampoline to avoid
deep synchronous callback re-entry when completing an acquisition immediately
releases capacity or drives another waiter.

Shutdown behavior

During shutdown, the pool first marks itself as shutting down and detaches:

  • queued acquisition callbacks;
  • pending-connect callbacks;
  • pending resolvers;
  • pending and established sockets.

These objects are moved out while holding the pool mutex. The mutex is then
released before any observer or acquisition callback is invoked.

The worker shutdown ordering preserves the baseline stop-then-join pattern:

stop io_context
join worker

After the worker exits, pending resolvers are cancelled, sockets are closed,
and detached acquisition callbacks are completed with failure outside the
pool mutex.

This strengthens acquisition-path shutdown guarantees without introducing a
new drain guarantee for transfers that have already started.

Backpressure limitation

This phase assumes that a connection occupying a slot eventually returns or
surfaces an error.

If an established transfer makes no progress and does not produce an I/O
error, it may continue occupying a capped slot, and queued acquisitions may
continue waiting behind the cap.

Progress deadlines and stale-connection detection are intentionally left for
a follow-up phase of #2930.

Configuration

The per-peer limit is configured through:

MC_TCP_MAX_CONNECTIONS_PER_PEER

The default value is 64.

The value is parsed only when the TCP connection pool is enabled. Valid values
are from 1 to 65535; invalid values fall back to the default.

The default is intended as an initial configurable guardrail rather than a
claim that 64 is optimal for every workload. Raising the value provides
behavior closer to the previous effectively unbounded pool.

Feedback on the initial default is welcome.

Test-only instrumentation

The deterministic pending-connect shutdown test uses narrowly scoped test
hooks.

The hook registry, setters, and invocation points are compiled only when
BUILD_UNIT_TESTS=ON through:

MOONCAKE_TCP_TRANSPORT_TEST_HOOKS

A production-check build with BUILD_UNIT_TESTS=OFF contained no hook
definition or hook symbol.

The shutdown detachment observer is also invoked only after releasing the pool
mutex.

Tests

The test suite includes coverage for:

  • enforcing the per-peer cap under concurrent writes;
  • FIFO backpressure above the cap;
  • reusing a returned connection with cap=1;
  • destroying the engine with active transfers and queued waiters;
  • destroying the engine during an in-progress resolve/connect;
  • detaching pending-connect ownership during shutdown;
  • completing acquisition callbacks exactly once on ordinary shutdown paths.

Validation results:

  • normal test suite: 11/11 passed;

  • ASan/UBSan/LeakSanitizer: 11/11 passed with no sanitizer diagnostics;

  • TSan:

    • PendingConnectShutdownDetachesOwnershipAndCompletesOnce: passed;
    • QueuedWaitersAndActiveTransfersDoNotMakeDestructionHang: passed;
    • both shutdown tests passed together in the same process;
    • the shutdown tests passed in both WSL and an independent AutoDL Linux
      environment.

A separate TSan run of
CapOneReusesConnectionForQueuedValidWrites exposed a pre-existing race
between:

Transport::Slice::markSuccess()

and:

TcpTransport::getTransferStatus()

The write side performs an atomic counter update while the read side reads the
same counter non-atomically. Neither access site is modified by this PR, so
that issue is intentionally kept outside this patch.

Correctness boundaries

The ordinary network, cancellation, shutdown, and callback-exception paths are
designed to transfer and complete acquisition callbacks exactly once.

Allocation failure through std::bad_alloc is outside the stated
exactly-once correctness contract of this PR.

Lifetime safety continues to rely on the existing Transfer Engine contract
that submission does not race destruction and that destruction is not
initiated from the TCP worker thread.

Suggested review order

The patch is easiest to review in this order:

  1. tcp_transport.h

    • PeerPool
    • PendingConnect
    • connection accounting and action structures
  2. tcp_transport.cpp

    • FIFO waiter driving
    • action dispatch and trampoline
    • pending resolve/connect completion
    • connection return/discard paths
    • shutdown extraction and callback completion
  3. tcp_write_visibility_test.cpp

    • cap enforcement
    • cap-one reuse
    • queued-waiter shutdown
    • pending-connect shutdown

The patch is intentionally test-heavy. Of the +1437/-164 total, roughly
1000 touched lines are in tcp_transport.h/.cpp (including replacement of
the previous synchronous pool path); the remainder is tests and build wiring.

The main changed files are:

tcp_transport.h
tcp_transport.cpp
tcp_transport/CMakeLists.txt
tcp_write_visibility_test.cpp
tests/CMakeLists.txt

Feedback requested

This is currently a Draft PR. In particular, feedback would be helpful on:

  • whether bounded pooling and stale-transfer recovery should remain separate
    phases;
  • whether the initial default cap of 64 is appropriate;
  • whether the deterministic test-hook approach is acceptable;
  • whether any additional failure-path tests are required before marking the PR
    ready for review.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@codecov-commenter

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 77.11268% with 195 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...gine/src/transport/tcp_transport/tcp_transport.cpp 63.74% 178 Missing ⚠️
...ransfer-engine/tests/tcp_write_visibility_test.cpp 95.26% 17 Missing ⚠️

📢 Thoughts on this report? Let us know!

@he-yufeng he-yufeng left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the connection-pool concurrency design here (I read it against #2930 rather than building the TCP transport locally, so this is a design-level review, not a runtime one).

The core shape is sound. The cap counting connections.size() + pending_connects.size() is the right call, since counting only established sockets would let many in-flight connects blow past the cap before any of them finish. FIFO waiters with driveWaitersLocked gives fair backpressure. The part I looked at hardest is callback re-entrancy, and it holds up: every acquisition callback is invoked outside state->mutex (the reject_for_shutdown path returns and calls back after the lock scope, and the normal path collects PoolActions under the lock and runs them via enqueuePoolActions after releasing it), so a callback that re-enters the pool cannot self-deadlock. Rejecting new acquisitions once shutting_down and failing queued waiters when the runtime is gone are the right deterministic-shutdown behaviors.

The one property I'd most want the added test to pin down (since it is exactly what an unbounded->bounded pool tends to get subtly wrong): under a shutdown that races in-flight connects, every entry in both pending_connects and waiters has its callback invoked exactly once, with none leaked (callback dropped) or double-invoked. A stress variant of tcp_write_visibility_test that spawns N over-cap acquisitions and calls shutdownConnectionPool() mid-connect, asserting the completion count equals N, would lock that in against future refactors.

Nice fix for #2930.

@alogfans

Copy link
Copy Markdown
Collaborator

Thanks for working on this issue. Bounding the number of TCP connections is a valuable first step, especially for preventing FD exhaustion and excessive TIME_WAIT accumulation under heavy transfers.

One concern is that the current implementation may move part of the resource pressure rather than fully eliminate it. A large transfer can still generate many acquisition callbacks, which then accumulate in PeerPool::waiters after the connection limit is reached. This could lead to significant user-space memory consumption and extended task lifetimes, and the resulting backlog may be harder to observe than socket growth.

It may be worth considering whether the connection lifecycle can be modeled differently. Since the number of destination peers is generally bounded, we could maintain a small set of long-lived connection lanes for each peer and let them consume work from a bounded per-peer transfer queue:

bounded transfer queue -> fixed long-lived connection lanes

This would apply backpressure when transfer work enters the system, instead of creating one connection-acquisition callback for every slice. It would also make queue depth, waiting time, cancellation, and peer health easier to manage and observe.

Overall, this PR is a useful step toward controlling connection growth. I’d just like us to ensure that the pressure is genuinely bounded rather than shifted from sockets to waiter state.

@jacklin78911-collab

jacklin78911-collab commented Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

@alogfans
Thanks — I agree with this assessment.

The current implementation bounds sockets, including in-flight connects, but shifts the remaining unbounded state to PeerPool::waiters. Boundedness is therefore not genuinely established, and as you noted, the resulting backlog is harder to observe than socket growth.

A bounded per-peer transfer queue consumed by a small fixed set of long-lived connection lanes looks like the cleaner end-state.

Before changing the branch, could you confirm the preferred scope for this Draft PR?

  1. Keep [Transfer Engine] Bound TCP connection pools with FIFO backpressure #2974 as a first-step guardrail against FD exhaustion and excessive TIME_WAIT, explicitly document the unbounded-waiter limitation, and implement the bounded queue / connection-lane model as a follow-up.

  2. Revise this Draft PR around the bounded queue / fixed-lane model before it is considered for merge.

If you prefer the second direction, I will first post a short design note covering the queued work-item granularity, queue-full behavior, transfer failure semantics, reconnect handling, and shutdown invariants before changing the implementation.

I will also carry @he-yufeng's suggested shutdown-race stress test into either design, since the exactly-once completion property remains important.

@jacklin78911-collab

jacklin78911-collab commented Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

@he-yufeng Thank you for the careful review, especially for validating the re-entrancy reasoning around PoolActions.

The exactly-once completion guarantee during a racing shutdown is the key invariant I want to preserve. With the direction discussed above—a bounded per-peer transfer queue consumed by a fixed set of long-lived lanes—I’ll adapt your suggested test to the revised design: submit N tasks beyond the queue/lane capacity, initiate shutdown while some tasks are connecting and others remain queued, and assert that the total completion count is exactly N, with no dropped or duplicate callbacks.

Although the abstraction changes from acquisition callbacks to queued transfer tasks, the invariant remains the same. I’ll include both the deterministic case and a stress variant in the revised test suite.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants