Struct StreamMultiplexer 

pub struct StreamMultiplexer { /* private fields */ }

Multiplexes multiple concurrent inference requests over a shared pipeline transport pair.

The PipelineStageRuntime send/recv methods process one message at a time. When multiple requests are in flight (e.g. batched API requests), their responses must be routed back to the correct caller by nonce.

StreamMultiplexer provides two complementary APIs:

• Request-response — send_and_await sends an ActivationMessage and suspends the caller until the matching reply (same nonce) arrives.
• Async dispatch — register_handler registers a one-shot oneshot::Sender that is fired when a response with the matching nonce is delivered by dispatch_incoming.

Typical usage

// Spawn one background task that continuously calls dispatch_incoming:
tokio::spawn(async move {
    loop {
        mux.dispatch_incoming(&mut stage).await.unwrap();
    }
});

// Each request task calls send_and_await:
let response = mux.send_and_await(msg, &mut stage).await?;

Implementations

impl StreamMultiplexer

pub fn new() -> Self

Create a new, empty multiplexer.

pub fn register_handler(&mut self, nonce: u64) -> Receiver<ActivationMessage>

Register a one-shot handler that fires when a response with nonce arrives.

Returns the corresponding oneshot::Receiver which the caller can await to get the response. If a handler for the same nonce is already registered, it is replaced and the old receiver will never fire.

pub async fn send_and_await( &mut self, msg: ActivationMessage, stage: &mut PipelineStageRuntime, ) -> DistributedResult<ActivationMessage>

Send msg via stage and await the matching response.

This is a higher-level convenience that calls register_handler, forwards the message to the next shard, and then awaits the registered one-shot receiver.

Errors

Returns DistributedError::Cancelled if the dispatch task drops the sender before delivering a response (e.g. on transport error).

pub async fn dispatch_incoming( &mut self, stage: &mut PipelineStageRuntime, ) -> DistributedResult<()>

Receive one incoming message from stage and route it to the waiting caller identified by its nonce.

This should be called repeatedly from a dedicated background task:

loop {
    mux.dispatch_incoming(&mut stage).await?;
}

Messages whose nonce is not in pending (e.g. unsolicited or already cancelled) are silently dropped.

pub fn pending_count(&self) -> usize

Number of requests currently awaiting a response.

Trait Implementations

impl Default for StreamMultiplexer

fn default() -> Self

Returns the “default value” for a type.