Refactor ConnectionActor to use single enum for mutually-exclusive response/multi-packet state

[coderabbitai]

Context

The ConnectionActor currently maintains two separate optional fields for handling streaming output:

• response: Option<FrameStream<F, E>>
• multi_packet: MultiPacketContext<F>

These fields are mutually exclusive by design, but this invariant is enforced only through runtime assertions at multiple points in the code:

• src/connection/mod.rs:245: set_response() asserts multi_packet is not active
• src/connection/mod.rs:257: set_multi_packet() asserts response is None
• src/connection/mod.rs:295: set_multi_packet_with_stamp() asserts response is None
• src/connection/mod.rs:366: run() asserts at most one is active

Proposed Solution

Replace the two separate fields with a single enum that encodes the mutual exclusion at the type level:

enum ActiveOutput<F, E> {
    None,
    Response(FrameStream<F, E>),
    MultiPacket(MultiPacketContext<F>),
}

This approach would:

1. Eliminate the runtime invariant checks entirely (making violation a compile-time error)
2. Simplify the tokio::select! logic by removing conditional branches
3. Make the mutual exclusion explicit in the type system
4. Reduce cognitive overhead when reading and maintaining the actor code

Benefits

• Type safety: Impossible states become unrepresentable
• Simplified logic: No need for runtime checks or conditional branches
• Clearer intent: The mutual exclusion is self-documenting in the type
• Easier maintenance: Future modifications won't accidentally violate the invariant

Related

Follow-up to PR #424 which refactored the connection module into submodules.

---

Requested by: @leynos
Source: #424 (comment)

[traycerai]

Plan

Observations

The ConnectionActor currently maintains mutual exclusion between response: Option<FrameStream<F, E>> and multi_packet: MultiPacketContext<F> through runtime debug_assert! checks at lines 343-350, 355-361, 393-399, and 464-468 in file:src/connection.rs. The MultiPacketContext already encapsulates channel and stamping state (lines 200-244), and the Event enum (lines 63-80) already distinguishes between response and multi-packet events. All public API usage occurs in test code, with ConnectionActor exported via file:src/lib.rs line 57.

Approach

Replace the two separate optional fields with a single enum that makes mutual exclusion a compile-time guarantee. This leverages the existing MultiPacketContext structure and Event enum separation. The refactoring will introduce a new ActiveOutput enum, update the constructor signatures to accept this enum, and modify setter methods to work with the unified type. Internal polling and processing logic will be simplified by eliminating conditional branches that check for mutual exclusion.

Implementation Steps

1. Define the ActiveOutput Enum

In file:src/connection.rs, add a new enum after the MultiPacketContext struct (around line 245):

enum ActiveOutput<F, E> {
    None,
    Response(FrameStream<F, E>),
    MultiPacket(MultiPacketContext<F>),
}

Add helper methods to ActiveOutput:

• is_none() → returns true for None variant
• is_response() → returns true for Response variant
• is_multi_packet() → returns true for MultiPacket variant
• as_response_mut() → returns Option<&mut FrameStream<F, E>>
• as_multi_packet_mut() → returns Option<&mut MultiPacketContext<F>>
• take() → consumes and returns the inner value, replacing with None

2. Update ConnectionActor Struct

In file:src/connection.rs lines 137-153, replace:

response: Option<FrameStream<F, E>>,
multi_packet: MultiPacketContext<F>,

With:

active_output: ActiveOutput<F, E>,

3. Update Constructor Methods

For new() method (lines 277-289):

• Change parameter from response: Option<FrameStream<F, E>> to active_output: ActiveOutput<F, E>
• Update the call to with_hooks() to pass active_output instead of response

For with_hooks() method (lines 293-323):

• Change parameter from response: Option<FrameStream<F, E>> to active_output: ActiveOutput<F, E>
• Initialize the struct with active_output field instead of separate response and multi_packet fields
• Update the ActorState::new() call (line 469) to use active_output.is_response() and active_output.is_multi_packet()

4. Refactor Setter Methods

Replace set_response() method (lines 342-351):

• Remove the debug_assert! check
• Change signature to pub fn set_active_output(&mut self, output: ActiveOutput<F, E>)
• Simply assign: self.active_output = output;

Remove set_multi_packet() method (lines 354-368):

• This functionality is now handled by set_active_output()

Replace set_multi_packet_with_correlation() method (lines 388-406):

• Remove the debug_assert! check
• Create a helper method or keep this as a convenience method that constructs ActiveOutput::MultiPacket with the appropriate stamping configuration
• Signature: pub fn set_multi_packet_with_correlation(&mut self, channel: Option<mpsc::Receiver<F>>, correlation_id: Option<u64>)
• Implementation creates MultiPacketContext, calls install(), then wraps in ActiveOutput::MultiPacket

Update clear_multi_packet() method (line 408):

• Change to self.active_output = ActiveOutput::None;

5. Update Internal Methods

Update apply_multi_packet_correlation() method (lines 410-434):

• Access stamping state via self.active_output.as_multi_packet_mut()

Update run() method (lines 450-480):

• Remove the debug_assert! at lines 464-468 (mutual exclusion is now enforced by the type system)
• Update ActorState::new() call to use self.active_output.is_response() and self.active_output.is_multi_packet()

Update next_event() method (lines 497-513):

• Change multi_available condition to check self.active_output.is_multi_packet()
• Change resp_available condition to check self.active_output.is_response()
• Update the poll_queue call for multi-packet to use self.active_output.as_multi_packet_mut().and_then(|ctx| ctx.channel_mut())
• Update the poll_response call to use self.active_output.as_response_mut()

Update process_response() method (lines 699-714):

• Change self.response = None; to self.active_output = ActiveOutput::None;

Update start_shutdown() method (lines 717-738):

• Replace the multi-packet handling block (lines 724-733) to work with self.active_output.as_multi_packet_mut()
• Replace the response handling block (lines 735-737) to work with self.active_output.take()

Update handle_multi_packet_closed() method (lines 617-635):

• Access multi-packet context via self.active_output.as_multi_packet_mut()
• After clearing, set self.active_output = ActiveOutput::None;

Update try_opportunistic_drain() method (lines 770-824):

• For QueueKind::Multi case (lines 800-823), access channel via self.active_output.as_multi_packet_mut()

6. Update Test Support Module

In file:src/connection/test_support.rs:

Update create_test_actor_with_hooks() function (lines 48-61):

• Change the None parameter to ActiveOutput::None

Update ActorHarness::set_multi_queue() method (lines 116-118):

• Change implementation to construct ActiveOutput::MultiPacket or ActiveOutput::None based on the parameter
• Call self.actor.set_active_output() instead of self.actor.set_multi_packet()

Update ActorHarness::has_multi_queue() method (lines 125-126):

• Change to self.actor.active_output.is_multi_packet()

7. Update All Test Files

For each test file that constructs ConnectionActor:

Files to update:

• file:tests/advanced/interaction_fuzz.rs
• file:tests/async_stream.rs
• file:tests/connection_actor_errors.rs
• file:tests/connection_actor_fairness.rs
• file:tests/connection_actor_shutdown.rs
• file:tests/connection_fragmentation.rs
• file:tests/correlation_id.rs
• file:tests/fixtures/correlation.rs
• file:tests/fixtures/multi_packet.rs
• file:tests/fixtures/stream_end.rs
• file:tests/multi_packet.rs
• file:tests/multi_packet_streaming.rs
• file:tests/stream_end.rs
• file:tests/wireframe_protocol.rs

Changes needed:

• Replace ConnectionActor::new(queues, handle, None, shutdown) with ConnectionActor::new(queues, handle, ActiveOutput::None, shutdown)
• Replace ConnectionActor::new(queues, handle, Some(stream), shutdown) with ConnectionActor::new(queues, handle, ActiveOutput::Response(stream), shutdown)
• Replace ConnectionActor::with_hooks(..., None, ...) with ConnectionActor::with_hooks(..., ActiveOutput::None, ...)
• Replace ConnectionActor::with_hooks(..., Some(stream), ...) with ConnectionActor::with_hooks(..., ActiveOutput::Response(stream), ...)
• Replace calls to actor.set_multi_packet(Some(rx)) with actor.set_multi_packet_with_correlation(Some(rx), None) or create the ActiveOutput::MultiPacket directly
• Replace calls to actor.set_multi_packet(None) with actor.set_active_output(ActiveOutput::None)

8. Update Documentation

Update doc comments in file:src/connection.rs:

• Constructor examples (lines 124-136, 263-275) to show ActiveOutput usage
• set_multi_packet_with_correlation() example (lines 372-387) to clarify it's a convenience method
• Add module-level documentation explaining the ActiveOutput enum and its purpose

9. Verify Compilation and Tests

After all changes:

• Run cargo build to ensure compilation succeeds
• Run cargo test to verify all tests pass
• Run cargo clippy to check for any new warnings
• Run cargo doc to ensure documentation builds correctly

Benefits Achieved

sequenceDiagram
    participant Client as Test/Client Code
    participant Actor as ConnectionActor
    participant Output as ActiveOutput Enum
    participant State as Type System
    
    Client->>Actor: new(queues, handle, ActiveOutput::Response(stream), shutdown)
    Actor->>Output: Store as active_output
    Note over State: Compile-time guarantee:<br/>Only one output type active
    
    Client->>Actor: Attempt to set multi-packet
    Actor->>Output: Replace with ActiveOutput::MultiPacket
    Note over State: No runtime checks needed<br/>Impossible states prevented
    
    Actor->>Output: Poll active output
    Output-->>Actor: Return appropriate stream/channel
    Note over Actor: Simplified select! logic<br/>No conditional branches

Loading

The refactoring eliminates all four runtime debug_assert! checks, making mutual exclusion violations a compile-time error. The tokio::select! logic in next_event() becomes simpler with clearer intent. Future modifications cannot accidentally violate the invariant, and the code is more maintainable with self-documenting types.

Import In IDE

🤖 Prompt for AI Agents

## Observations

The `ConnectionActor` currently maintains mutual exclusion between `response: Option<FrameStream<F, E>>` and `multi_packet: MultiPacketContext<F>` through runtime `debug_assert!` checks at lines 343-350, 355-361, 393-399, and 464-468 in `file:src/connection.rs`. The `MultiPacketContext` already encapsulates channel and stamping state (lines 200-244), and the `Event` enum (lines 63-80) already distinguishes between response and multi-packet events. All public API usage occurs in test code, with `ConnectionActor` exported via `file:src/lib.rs` line 57.

## Approach

Replace the two separate optional fields with a single enum that makes mutual exclusion a compile-time guarantee. This leverages the existing `MultiPacketContext` structure and `Event` enum separation. The refactoring will introduce a new `ActiveOutput` enum, update the constructor signatures to accept this enum, and modify setter methods to work with the unified type. Internal polling and processing logic will be simplified by eliminating conditional branches that check for mutual exclusion.

## Implementation Steps

### 1. Define the `ActiveOutput` Enum

In `file:src/connection.rs`, add a new enum after the `MultiPacketContext` struct (around line 245):

```rust
enum ActiveOutput<F, E> {
    None,
    Response(FrameStream<F, E>),
    MultiPacket(MultiPacketContext<F>),
}
```

Add helper methods to `ActiveOutput`:
- `is_none()` → returns true for `None` variant
- `is_response()` → returns true for `Response` variant
- `is_multi_packet()` → returns true for `MultiPacket` variant
- `as_response_mut()` → returns `Option<&mut FrameStream<F, E>>`
- `as_multi_packet_mut()` → returns `Option<&mut MultiPacketContext<F>>`
- `take()` → consumes and returns the inner value, replacing with `None`

### 2. Update `ConnectionActor` Struct

In `file:src/connection.rs` lines 137-153, replace:
```rust
response: Option<FrameStream<F, E>>,
multi_packet: MultiPacketContext<F>,
```

With:
```rust
active_output: ActiveOutput<F, E>,
```

### 3. Update Constructor Methods

**For `new()` method** (lines 277-289):
- Change parameter from `response: Option<FrameStream<F, E>>` to `active_output: ActiveOutput<F, E>`
- Update the call to `with_hooks()` to pass `active_output` instead of `response`

**For `with_hooks()` method** (lines 293-323):
- Change parameter from `response: Option<FrameStream<F, E>>` to `active_output: ActiveOutput<F, E>`
- Initialize the struct with `active_output` field instead of separate `response` and `multi_packet` fields
- Update the `ActorState::new()` call (line 469) to use `active_output.is_response()` and `active_output.is_multi_packet()`

### 4. Refactor Setter Methods

**Replace `set_response()` method** (lines 342-351):
- Remove the `debug_assert!` check
- Change signature to `pub fn set_active_output(&mut self, output: ActiveOutput<F, E>)`
- Simply assign: `self.active_output = output;`

**Remove `set_multi_packet()` method** (lines 354-368):
- This functionality is now handled by `set_active_output()`

**Replace `set_multi_packet_with_correlation()` method** (lines 388-406):
- Remove the `debug_assert!` check
- Create a helper method or keep this as a convenience method that constructs `ActiveOutput::MultiPacket` with the appropriate stamping configuration
- Signature: `pub fn set_multi_packet_with_correlation(&mut self, channel: Option<mpsc::Receiver<F>>, correlation_id: Option<u64>)`
- Implementation creates `MultiPacketContext`, calls `install()`, then wraps in `ActiveOutput::MultiPacket`

**Update `clear_multi_packet()` method** (line 408):
- Change to `self.active_output = ActiveOutput::None;`

### 5. Update Internal Methods

**Update `apply_multi_packet_correlation()` method** (lines 410-434):
- Access stamping state via `self.active_output.as_multi_packet_mut()`

**Update `run()` method** (lines 450-480):
- Remove the `debug_assert!` at lines 464-468 (mutual exclusion is now enforced by the type system)
- Update `ActorState::new()` call to use `self.active_output.is_response()` and `self.active_output.is_multi_packet()`

**Update `next_event()` method** (lines 497-513):
- Change `multi_available` condition to check `self.active_output.is_multi_packet()`
- Change `resp_available` condition to check `self.active_output.is_response()`
- Update the `poll_queue` call for multi-packet to use `self.active_output.as_multi_packet_mut().and_then(|ctx| ctx.channel_mut())`
- Update the `poll_response` call to use `self.active_output.as_response_mut()`

**Update `process_response()` method** (lines 699-714):
- Change `self.response = None;` to `self.active_output = ActiveOutput::None;`

**Update `start_shutdown()` method** (lines 717-738):
- Replace the multi-packet handling block (lines 724-733) to work with `self.active_output.as_multi_packet_mut()`
- Replace the response handling block (lines 735-737) to work with `self.active_output.take()`

**Update `handle_multi_packet_closed()` method** (lines 617-635):
- Access multi-packet context via `self.active_output.as_multi_packet_mut()`
- After clearing, set `self.active_output = ActiveOutput::None;`

**Update `try_opportunistic_drain()` method** (lines 770-824):
- For `QueueKind::Multi` case (lines 800-823), access channel via `self.active_output.as_multi_packet_mut()`

### 6. Update Test Support Module

In `file:src/connection/test_support.rs`:

**Update `create_test_actor_with_hooks()` function** (lines 48-61):
- Change the `None` parameter to `ActiveOutput::None`

**Update `ActorHarness::set_multi_queue()` method** (lines 116-118):
- Change implementation to construct `ActiveOutput::MultiPacket` or `ActiveOutput::None` based on the parameter
- Call `self.actor.set_active_output()` instead of `self.actor.set_multi_packet()`

**Update `ActorHarness::has_multi_queue()` method** (lines 125-126):
- Change to `self.actor.active_output.is_multi_packet()`

### 7. Update All Test Files

For each test file that constructs `ConnectionActor`:

**Files to update**:
- `file:tests/advanced/interaction_fuzz.rs`
- `file:tests/async_stream.rs`
- `file:tests/connection_actor_errors.rs`
- `file:tests/connection_actor_fairness.rs`
- `file:tests/connection_actor_shutdown.rs`
- `file:tests/connection_fragmentation.rs`
- `file:tests/correlation_id.rs`
- `file:tests/fixtures/correlation.rs`
- `file:tests/fixtures/multi_packet.rs`
- `file:tests/fixtures/stream_end.rs`
- `file:tests/multi_packet.rs`
- `file:tests/multi_packet_streaming.rs`
- `file:tests/stream_end.rs`
- `file:tests/wireframe_protocol.rs`

**Changes needed**:
- Replace `ConnectionActor::new(queues, handle, None, shutdown)` with `ConnectionActor::new(queues, handle, ActiveOutput::None, shutdown)`
- Replace `ConnectionActor::new(queues, handle, Some(stream), shutdown)` with `ConnectionActor::new(queues, handle, ActiveOutput::Response(stream), shutdown)`
- Replace `ConnectionActor::with_hooks(..., None, ...)` with `ConnectionActor::with_hooks(..., ActiveOutput::None, ...)`
- Replace `ConnectionActor::with_hooks(..., Some(stream), ...)` with `ConnectionActor::with_hooks(..., ActiveOutput::Response(stream), ...)`
- Replace calls to `actor.set_multi_packet(Some(rx))` with `actor.set_multi_packet_with_correlation(Some(rx), None)` or create the `ActiveOutput::MultiPacket` directly
- Replace calls to `actor.set_multi_packet(None)` with `actor.set_active_output(ActiveOutput::None)`

### 8. Update Documentation

Update doc comments in `file:src/connection.rs`:
- Constructor examples (lines 124-136, 263-275) to show `ActiveOutput` usage
- `set_multi_packet_with_correlation()` example (lines 372-387) to clarify it's a convenience method
- Add module-level documentation explaining the `ActiveOutput` enum and its purpose

### 9. Verify Compilation and Tests

After all changes:
- Run `cargo build` to ensure compilation succeeds
- Run `cargo test` to verify all tests pass
- Run `cargo clippy` to check for any new warnings
- Run `cargo doc` to ensure documentation builds correctly

## Benefits Achieved

```mermaid
sequenceDiagram
    participant Client as Test/Client Code
    participant Actor as ConnectionActor
    participant Output as ActiveOutput Enum
    participant State as Type System
    
    Client->>Actor: new(queues, handle, ActiveOutput::Response(stream), shutdown)
    Actor->>Output: Store as active_output
    Note over State: Compile-time guarantee:<br/>Only one output type active
    
    Client->>Actor: Attempt to set multi-packet
    Actor->>Output: Replace with ActiveOutput::MultiPacket
    Note over State: No runtime checks needed<br/>Impossible states prevented
    
    Actor->>Output: Poll active output
    Output-->>Actor: Return appropriate stream/channel
    Note over Actor: Simplified select! logic<br/>No conditional branches
```

The refactoring eliminates all four runtime `debug_assert!` checks, making mutual exclusion violations a compile-time error. The `tokio::select!` logic in `next_event()` becomes simpler with clearer intent. Future modifications cannot accidentally violate the invariant, and the code is more maintainable with self-documenting types.

---

Execution Information

Branch: main
Commit: 9b4ea6e

---

💡 Tips

Supported Commands (Inside Comments)

• Use @traycerai generate to iterate on the previous version of the implementation plan.

Supported Commands (Inside Description)

• Add @traycerai ignore anywhere in the ticket description to prevent this ticket from being processed.
• Add @traycerai branch:<branch-name> anywhere in the ticket description to specify the target branch for the implementation plan.

Community

• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.