Skip to main content

redis/aio/
multiplexed_connection.rs

1use super::{AsyncPushSender, ConnectionLike, Runtime, SharedHandleContainer, TaskHandle};
2#[cfg(feature = "cache-aio")]
3use crate::caching::{CacheManager, CacheStatistics, PrepCacheItem};
4use crate::{
5    AsyncConnectionConfig, ProtocolVersion, PushInfo, RedisConnectionInfo, ServerError,
6    ToRedisArgs,
7    aio::setup_connection,
8    check_resp3, cmd,
9    cmd::Cmd,
10    errors::{ErrorKind, RedisError, closed_connection_error},
11    parser::ValueCodec,
12    types::{RedisFuture, RedisResult, Value},
13};
14use ::tokio::{
15    io::{AsyncRead, AsyncWrite},
16    sync::{mpsc, oneshot},
17};
18#[cfg(feature = "token-based-authentication")]
19use {
20    arcstr::ArcStr,
21    log::{debug, error, warn},
22};
23
24use futures_util::{
25    future::{Future, FutureExt},
26    ready,
27    sink::Sink,
28    stream::{self, Stream, StreamExt},
29};
30use pin_project_lite::pin_project;
31use std::collections::VecDeque;
32use std::fmt;
33use std::fmt::Debug;
34use std::pin::Pin;
35use std::sync::Arc;
36use std::task::{self, Poll};
37use std::time::Duration;
38use tokio_util::codec::Decoder;
39
40// Senders which the result of a single request are sent through
41type PipelineOutput = oneshot::Sender<RedisResult<Value>>;
42
43enum ErrorOrErrors {
44    Errors(Vec<(usize, ServerError)>),
45    // only set if we receive a transmission error
46    FirstError(RedisError),
47}
48
49enum ResponseAggregate {
50    SingleCommand,
51    Pipeline {
52        buffer: Vec<Value>,
53        error_or_errors: ErrorOrErrors,
54        expectation: PipelineResponseExpectation,
55    },
56}
57
58// TODO - this is a really bad name.
59struct PipelineResponseExpectation {
60    // The number of responses to skip before starting to save responses in the buffer.
61    skipped_response_count: usize,
62    // The number of responses to keep in the buffer
63    expected_response_count: usize,
64    // whether the pipelined request is a transaction
65    is_transaction: bool,
66    seen_responses: usize,
67}
68
69impl ResponseAggregate {
70    fn new(expectation: Option<PipelineResponseExpectation>) -> Self {
71        match expectation {
72            Some(expectation) => Self::Pipeline {
73                buffer: Vec::new(),
74                error_or_errors: ErrorOrErrors::Errors(Vec::new()),
75                expectation,
76            },
77            None => Self::SingleCommand,
78        }
79    }
80}
81
82struct InFlight {
83    output: Option<PipelineOutput>,
84    response_aggregate: ResponseAggregate,
85}
86
87// A single message sent through the pipeline
88struct PipelineMessage {
89    input: Vec<u8>,
90    // If `output` is None, then the caller doesn't expect to receive an answer.
91    output: Option<PipelineOutput>,
92    // If `None`, this is a single request, not a pipeline of multiple requests.
93    // If `Some`, the first value is the number of responses to skip,
94    // the second is the number of responses to keep, and the third is whether the pipeline is a transaction.
95    expectation: Option<PipelineResponseExpectation>,
96}
97
98/// Wrapper around a `Stream + Sink` where each item sent through the `Sink` results in one or more
99/// items being output by the `Stream` (the number is specified at time of sending). With the
100/// interface provided by `Pipeline` an easy interface of request to response, hiding the `Stream`
101/// and `Sink`.
102#[derive(Clone)]
103struct Pipeline {
104    sender: mpsc::Sender<PipelineMessage>,
105}
106
107impl Debug for Pipeline {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        f.debug_tuple("Pipeline").field(&self.sender).finish()
110    }
111}
112
113#[cfg(feature = "cache-aio")]
114pin_project! {
115    struct PipelineSink<T> {
116        #[pin]
117        sink_stream: T,
118        in_flight: VecDeque<InFlight>,
119        error: Option<RedisError>,
120        push_sender: Option<Arc<dyn AsyncPushSender>>,
121        cache_manager: Option<CacheManager>,
122    }
123}
124
125#[cfg(not(feature = "cache-aio"))]
126pin_project! {
127    struct PipelineSink<T> {
128        #[pin]
129        sink_stream: T,
130        in_flight: VecDeque<InFlight>,
131        error: Option<RedisError>,
132        push_sender: Option<Arc<dyn AsyncPushSender>>,
133    }
134}
135
136fn send_push(push_sender: &Option<Arc<dyn AsyncPushSender>>, info: PushInfo) {
137    if let Some(sender) = push_sender {
138        let _ = sender.send(info);
139    }
140}
141
142pub(crate) fn send_disconnect(push_sender: &Option<Arc<dyn AsyncPushSender>>) {
143    send_push(push_sender, PushInfo::disconnect());
144}
145
146impl<T> PipelineSink<T>
147where
148    T: Stream<Item = RedisResult<Value>> + 'static,
149{
150    fn new(
151        sink_stream: T,
152        push_sender: Option<Arc<dyn AsyncPushSender>>,
153        #[cfg(feature = "cache-aio")] cache_manager: Option<CacheManager>,
154    ) -> Self
155    where
156        T: Sink<Vec<u8>, Error = RedisError> + Stream<Item = RedisResult<Value>> + 'static,
157    {
158        Self {
159            sink_stream,
160            in_flight: VecDeque::new(),
161            error: None,
162            push_sender,
163            #[cfg(feature = "cache-aio")]
164            cache_manager,
165        }
166    }
167
168    // Read messages from the stream and send them back to the caller
169    fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<(), ()>> {
170        loop {
171            let item = ready!(self.as_mut().project().sink_stream.poll_next(cx));
172            let item = match item {
173                Some(result) => result,
174                // The redis response stream is not going to produce any more items so we simulate a disconnection error to break out of the loop.
175                None => Err(closed_connection_error()),
176            };
177
178            let is_unrecoverable = item.as_ref().is_err_and(|err| err.is_unrecoverable_error());
179            self.as_mut().send_result(item);
180            if is_unrecoverable {
181                let self_ = self.project();
182                send_disconnect(self_.push_sender);
183                return Poll::Ready(Err(()));
184            }
185        }
186    }
187
188    fn send_result(self: Pin<&mut Self>, result: RedisResult<Value>) {
189        let self_ = self.project();
190        let result = match result {
191            // If this push message isn't a reply, we'll pass it as-is to the push manager and stop iterating
192            Ok(Value::Push { kind, data }) if !kind.has_reply() => {
193                #[cfg(feature = "cache-aio")]
194                if let Some(cache_manager) = &self_.cache_manager {
195                    cache_manager.handle_push_value(&kind, &data);
196                }
197                send_push(self_.push_sender, PushInfo { kind, data });
198
199                return;
200            }
201            // If this push message is a reply to a query, we'll clone it to the push manager and continue with sending the reply
202            Ok(Value::Push { kind, data }) if kind.has_reply() => {
203                send_push(
204                    self_.push_sender,
205                    PushInfo {
206                        kind: kind.clone(),
207                        data: data.clone(),
208                    },
209                );
210                Ok(Value::Push { kind, data })
211            }
212            _ => result,
213        };
214
215        let mut entry = match self_.in_flight.pop_front() {
216            Some(entry) => entry,
217            None => return,
218        };
219
220        match &mut entry.response_aggregate {
221            ResponseAggregate::SingleCommand => {
222                if let Some(output) = entry.output.take() {
223                    _ = output.send(result);
224                }
225            }
226            ResponseAggregate::Pipeline {
227                buffer,
228                error_or_errors,
229                expectation:
230                    PipelineResponseExpectation {
231                        expected_response_count,
232                        skipped_response_count,
233                        is_transaction,
234                        seen_responses,
235                    },
236            } => {
237                *seen_responses += 1;
238                if *skipped_response_count > 0 {
239                    // server errors in skipped values are still counted for errors in transactions, since they're errors that will cause the transaction to fail,
240                    // and we only skip values in transaction.
241                    if *is_transaction && let ErrorOrErrors::Errors(errs) = error_or_errors {
242                        match result {
243                            Ok(Value::ServerError(err)) => {
244                                errs.push((*seen_responses - 2, err)); // - 1 to offset the early increment, and -1 to offset the added MULTI call.
245                            }
246                            Err(err) => *error_or_errors = ErrorOrErrors::FirstError(err),
247                            _ => {}
248                        }
249                    }
250
251                    *skipped_response_count -= 1;
252                    self_.in_flight.push_front(entry);
253                    return;
254                }
255
256                match result {
257                    Ok(item) => {
258                        buffer.push(item);
259                    }
260                    Err(err) => {
261                        if matches!(error_or_errors, ErrorOrErrors::Errors(_)) {
262                            *error_or_errors = ErrorOrErrors::FirstError(err);
263                        }
264                    }
265                }
266
267                if buffer.len() < *expected_response_count {
268                    // Need to gather more response values
269                    self_.in_flight.push_front(entry);
270                    return;
271                }
272
273                let response =
274                    match std::mem::replace(error_or_errors, ErrorOrErrors::Errors(Vec::new())) {
275                        ErrorOrErrors::Errors(errors) => {
276                            if errors.is_empty() {
277                                Ok(Value::Array(std::mem::take(buffer)))
278                            } else {
279                                Err(RedisError::make_aborted_transaction(errors))
280                            }
281                        }
282                        ErrorOrErrors::FirstError(redis_error) => Err(redis_error),
283                    };
284
285                // `Err` means that the receiver was dropped in which case it does not
286                // care about the output and we can continue by just dropping the value
287                // and sender
288                if let Some(output) = entry.output.take() {
289                    _ = output.send(response);
290                }
291            }
292        }
293    }
294}
295
296impl<T> Sink<PipelineMessage> for PipelineSink<T>
297where
298    T: Sink<Vec<u8>, Error = RedisError> + Stream<Item = RedisResult<Value>> + 'static,
299{
300    type Error = ();
301
302    // Retrieve incoming messages and write them to the sink
303    fn poll_ready(
304        mut self: Pin<&mut Self>,
305        cx: &mut task::Context,
306    ) -> Poll<Result<(), Self::Error>> {
307        // It is crucial that we always try to advance both the read and the write side together.
308        // If we do not, we are susceptible to TCP deadlock.
309        // Here, we do this by advancing reads and then moving on to advance writes regardless of
310        // whether the read was pending or not. This ensures that we are registered to be woken
311        // if either reads or writes become available.
312        // See https://github.com/redis-rs/redis-rs/issues/1955.
313        if matches!(self.as_mut().poll_read(cx), Poll::Ready(Err(()))) {
314            return Poll::Ready(Err(()));
315        }
316        match ready!(self.as_mut().project().sink_stream.poll_ready(cx)) {
317            Ok(()) => Ok(()).into(),
318            Err(err) => {
319                *self.project().error = Some(err);
320                Ok(()).into()
321            }
322        }
323    }
324
325    fn start_send(
326        mut self: Pin<&mut Self>,
327        PipelineMessage {
328            input,
329            mut output,
330            expectation,
331        }: PipelineMessage,
332    ) -> Result<(), Self::Error> {
333        // If initially a receiver was created, but then dropped, there is nothing to receive our output we do not need to send the message as it is
334        // ambiguous whether the message will be sent anyway. Helps shed some load on the
335        // connection.
336        if output.as_ref().is_some_and(|output| output.is_closed()) {
337            return Ok(());
338        }
339
340        let self_ = self.as_mut().project();
341
342        if let Some(err) = self_.error.take() {
343            if let Some(output) = output.take() {
344                _ = output.send(Err(err));
345            }
346            return Err(());
347        }
348
349        match self_.sink_stream.start_send(input) {
350            Ok(()) => {
351                let response_aggregate = ResponseAggregate::new(expectation);
352                let entry = InFlight {
353                    output,
354                    response_aggregate,
355                };
356
357                self_.in_flight.push_back(entry);
358                Ok(())
359            }
360            Err(err) => {
361                if let Some(output) = output.take() {
362                    _ = output.send(Err(err));
363                }
364                Err(())
365            }
366        }
367    }
368
369    fn poll_flush(
370        mut self: Pin<&mut Self>,
371        cx: &mut task::Context,
372    ) -> Poll<Result<(), Self::Error>> {
373        // It is crucial that we always try to advance both the read and the write side together.
374        // If we do not, we are susceptible to TCP deadlock.
375        // Here, we do this by advancing reads and then moving on to advance writes regardless of
376        // whether the read was pending or not. This ensures that we are registered to be woken
377        // if either reads or writes become available.
378        // See https://github.com/redis-rs/redis-rs/issues/1955.
379        if matches!(self.as_mut().poll_read(cx), Poll::Ready(Err(()))) {
380            return Poll::Ready(Err(()));
381        }
382        self.as_mut()
383            .project()
384            .sink_stream
385            .poll_flush(cx)
386            .map_err(|err| {
387                self.as_mut().send_result(Err(err));
388            })
389    }
390
391    fn poll_close(
392        mut self: Pin<&mut Self>,
393        cx: &mut task::Context,
394    ) -> Poll<Result<(), Self::Error>> {
395        // No new requests will come in after the first call to `close` but we need to complete any
396        // in progress requests before closing
397        if !self.in_flight.is_empty() {
398            ready!(self.as_mut().poll_flush(cx))?;
399        }
400        let this = self.as_mut().project();
401        this.sink_stream.poll_close(cx).map_err(|err| {
402            self.send_result(Err(err));
403        })
404    }
405}
406
407impl Pipeline {
408    const DEFAULT_BUFFER_SIZE: usize = 50;
409
410    fn resolve_buffer_size(size: Option<usize>) -> usize {
411        size.unwrap_or(Self::DEFAULT_BUFFER_SIZE)
412    }
413
414    fn new<T>(
415        sink_stream: T,
416        push_sender: Option<Arc<dyn AsyncPushSender>>,
417        #[cfg(feature = "cache-aio")] cache_manager: Option<CacheManager>,
418        buffer_size: usize,
419    ) -> (Self, impl Future<Output = ()>)
420    where
421        T: Sink<Vec<u8>, Error = RedisError>,
422        T: Stream<Item = RedisResult<Value>>,
423        T: Unpin + Send + 'static,
424    {
425        let (sender, mut receiver) = mpsc::channel(buffer_size);
426
427        let sink = PipelineSink::new(
428            sink_stream,
429            push_sender,
430            #[cfg(feature = "cache-aio")]
431            cache_manager,
432        );
433        let f = stream::poll_fn(move |cx| receiver.poll_recv(cx))
434            .map(Ok)
435            .forward(sink)
436            .map(|_| ());
437        (Self { sender }, f)
438    }
439
440    async fn send_recv(
441        &mut self,
442        input: Vec<u8>,
443        // If `None`, this is a single request, not a pipeline of multiple requests.
444        // If `Some`, the value inside defines how the response should look like
445        expectation: Option<PipelineResponseExpectation>,
446        timeout: Option<Duration>,
447        skip_response: bool,
448    ) -> Result<Value, RedisError> {
449        if input.is_empty() {
450            return Err(RedisError::make_empty_command());
451        }
452
453        let request = async {
454            if skip_response {
455                self.sender
456                    .send(PipelineMessage {
457                        input,
458                        expectation,
459                        output: None,
460                    })
461                    .await
462                    .map_err(|_| None)?;
463
464                return Ok(Value::Nil);
465            }
466
467            let (sender, receiver) = oneshot::channel();
468
469            self.sender
470                .send(PipelineMessage {
471                    input,
472                    expectation,
473                    output: Some(sender),
474                })
475                .await
476                .map_err(|_| None)?;
477
478            receiver.await
479            // The `sender` was dropped which likely means that the stream part
480            // failed for one reason or another
481            .map_err(|_| None)
482            .and_then(|res| res.map_err(Some))
483        };
484
485        match timeout {
486            Some(timeout) => match Runtime::locate().timeout(timeout, request).await {
487                Ok(res) => res,
488                Err(elapsed) => Err(Some(elapsed.into())),
489            },
490            None => request.await,
491        }
492        .map_err(|err| err.unwrap_or_else(closed_connection_error))
493    }
494}
495
496/// Semaphore plus the `limit` it was created with, so the limit stays readable.
497#[derive(Clone, Debug)]
498struct ConcurrencyLimiter {
499    limit: usize,
500    semaphore: Arc<async_lock::Semaphore>,
501}
502
503/// Permits a single pipeline can hold: capped at the connection `limit`.
504/// At least one is always acquired.
505fn permit_vec_capacity(count: usize, limit: usize) -> usize {
506    count.min(limit).max(1)
507}
508
509/// Builds the optional concurrency limiter from the configured limit. A limit of `0` is rejected.
510fn build_concurrency_limiter(limit: Option<usize>) -> RedisResult<Option<ConcurrencyLimiter>> {
511    match limit {
512        Some(0) => Err(RedisError::from((
513            ErrorKind::InvalidClientConfig,
514            "concurrency limit must be greater than 0",
515        ))),
516        Some(n) => Ok(Some(ConcurrencyLimiter {
517            limit: n,
518            semaphore: Arc::new(async_lock::Semaphore::new(n)),
519        })),
520        None => Ok(None),
521    }
522}
523
524/// A connection object which can be cloned, allowing requests to be be sent concurrently
525/// on the same underlying connection (tcp/unix socket).
526///
527/// This connection object is cancellation-safe, and the user can drop request future without polling them to completion,
528/// but this doesn't mean that the actual request sent to the server is cancelled.
529/// A side-effect of this is that the underlying connection won't be closed until all sent requests have been answered,
530/// which means that in case of blocking commands, the underlying connection resource might not be released,
531/// even when all clones of the multiplexed connection have been dropped (see <https://github.com/redis-rs/redis-rs/issues/1236>).
532/// This isn't an issue in a connection that was created in a canonical way, which ensures that `_task_handle` is set, so that
533/// once all of the connection's clones are dropped, the task will also be dropped. If the user creates the connection in
534/// another way and `_task_handle` isn't set, they should manually spawn the returned driver function, keep the spawned task's
535/// handle and abort the task whenever they want, at the risk of effectively closing the clones of the multiplexed connection.
536#[derive(Clone)]
537pub struct MultiplexedConnection {
538    pipeline: Pipeline,
539    db: i64,
540    response_timeout: Option<Duration>,
541    protocol: ProtocolVersion,
542    concurrency_limiter: Option<ConcurrencyLimiter>,
543    // This handle ensures that once all the clones of the connection will be dropped, the underlying task will stop.
544    // This handle is only set for connection whose task was spawned by the crate, not for users who spawned their own
545    // task.
546    _task_handle: Option<SharedHandleContainer>,
547    #[cfg(feature = "cache-aio")]
548    pub(crate) cache_manager: Option<CacheManager>,
549    #[cfg(feature = "token-based-authentication")]
550    // This handle ensures that once all the clones of the connection will be dropped, the underlying task will stop.
551    // It is only set for connections that use a credentials provider for token-based authentication.
552    _credentials_subscription_task_handle: Option<SharedHandleContainer>,
553}
554
555impl Debug for MultiplexedConnection {
556    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
557        let Self {
558            pipeline,
559            db,
560            response_timeout,
561            protocol,
562            concurrency_limiter: _,
563            _task_handle,
564            #[cfg(feature = "cache-aio")]
565                cache_manager: _,
566            #[cfg(feature = "token-based-authentication")]
567                _credentials_subscription_task_handle: _,
568        } = self;
569
570        f.debug_struct("MultiplexedConnection")
571            .field("pipeline", &pipeline)
572            .field("db", &db)
573            .field("response_timeout", &response_timeout)
574            .field("protocol", &protocol)
575            .finish()
576    }
577}
578
579impl MultiplexedConnection {
580    /// Constructs a new `MultiplexedConnection` out of a `AsyncRead + AsyncWrite` object
581    /// and a `RedisConnectionInfo`
582    pub async fn new<C>(
583        connection_info: &RedisConnectionInfo,
584        stream: C,
585    ) -> RedisResult<(Self, impl Future<Output = ()>)>
586    where
587        C: Unpin + AsyncRead + AsyncWrite + Send + 'static,
588    {
589        Self::new_with_config(connection_info, stream, AsyncConnectionConfig::default()).await
590    }
591
592    /// Constructs a new `MultiplexedConnection` out of a `AsyncRead + AsyncWrite` object
593    /// , a `RedisConnectionInfo` and a `AsyncConnectionConfig`.
594    pub async fn new_with_config<C>(
595        connection_info: &RedisConnectionInfo,
596        stream: C,
597        config: AsyncConnectionConfig,
598    ) -> RedisResult<(Self, impl Future<Output = ()> + 'static)>
599    where
600        C: Unpin + AsyncRead + AsyncWrite + Send + 'static,
601    {
602        let mut codec = ValueCodec::default().framed(stream);
603        if let Some(boundary) = config.write_backpressure_boundary {
604            codec.set_backpressure_boundary(boundary);
605        }
606        if config.push_sender.is_some() {
607            check_resp3!(
608                connection_info.protocol,
609                "Can only pass push sender to a connection using RESP3"
610            );
611        }
612
613        #[cfg(feature = "cache-aio")]
614        let cache_config = config.cache.as_ref().map(|cache| match cache {
615            crate::client::Cache::Config(cache_config) => *cache_config,
616            #[cfg(any(feature = "connection-manager", feature = "cluster-async"))]
617            crate::client::Cache::Manager(cache_manager) => cache_manager.cache_config,
618        });
619        #[cfg(feature = "cache-aio")]
620        let cache_manager_opt = config
621            .cache
622            .map(|cache| {
623                check_resp3!(
624                    connection_info.protocol,
625                    "Can only enable client side caching in a connection using RESP3"
626                );
627                match cache {
628                    crate::client::Cache::Config(cache_config) => {
629                        Ok(CacheManager::new(cache_config))
630                    }
631                    #[cfg(any(feature = "connection-manager", feature = "cluster-async"))]
632                    crate::client::Cache::Manager(cache_manager) => Ok(cache_manager),
633                }
634            })
635            .transpose()?;
636
637        #[cfg(feature = "token-based-authentication")]
638        let mut connection_info = connection_info.clone();
639        #[cfg(not(feature = "token-based-authentication"))]
640        let connection_info = connection_info.clone();
641
642        #[cfg(feature = "token-based-authentication")]
643        if let Some(ref credentials_provider) = config.credentials_provider {
644            // Retrieve the initial credentials from the provider and apply them to the connection info
645            match credentials_provider.subscribe().next().await {
646                Some(Ok(credentials)) => {
647                    connection_info.username = Some(ArcStr::from(credentials.username));
648                    connection_info.password = Some(ArcStr::from(credentials.password));
649                }
650                Some(Err(err)) => {
651                    error!("Error while receiving credentials from stream: {err}");
652                    return Err(err);
653                }
654                None => {
655                    let err = RedisError::from((
656                        ErrorKind::AuthenticationFailed,
657                        "Credentials stream closed unexpectedly before yielding credentials!",
658                    ));
659                    error!("{err}");
660                    return Err(err);
661                }
662            }
663        }
664
665        setup_connection(
666            &mut codec,
667            &connection_info,
668            #[cfg(feature = "cache-aio")]
669            cache_config,
670        )
671        .await?;
672        if config.push_sender.is_some() {
673            check_resp3!(
674                connection_info.protocol,
675                "Can only pass push sender to a connection using RESP3"
676            );
677        }
678
679        let (pipeline, driver) = Pipeline::new(
680            codec,
681            config.push_sender,
682            #[cfg(feature = "cache-aio")]
683            cache_manager_opt.clone(),
684            Pipeline::resolve_buffer_size(config.pipeline_buffer_size),
685        );
686
687        let concurrency_limiter = build_concurrency_limiter(config.concurrency_limit)?;
688
689        let con = Self {
690            pipeline,
691            db: connection_info.db,
692            response_timeout: config.response_timeout,
693            protocol: connection_info.protocol,
694            concurrency_limiter,
695            _task_handle: None,
696            #[cfg(feature = "cache-aio")]
697            cache_manager: cache_manager_opt,
698            #[cfg(feature = "token-based-authentication")]
699            _credentials_subscription_task_handle: None,
700        };
701
702        // Set up streaming credentials subscription if provider is available
703        #[cfg(feature = "token-based-authentication")]
704        if let Some(streaming_provider) = config.credentials_provider {
705            let mut inner_connection = con.clone();
706            let mut stream = streaming_provider.subscribe();
707
708            let subscription_task_handle = Runtime::locate().spawn(async move {
709                while let Some(result) = stream.next().await {
710                    match result {
711                        Ok(credentials) => {
712                            if let Err(err) = inner_connection
713                                .re_authenticate_with_credentials(&credentials)
714                                .await
715                            {
716                                if err.is_connection_dropped() {
717                                    warn!(
718                                        "Re-authentication task ended, connection is dead: {err}"
719                                    );
720                                    return;
721                                }
722                                error!("Failed to re-authenticate async connection: {err}.");
723                                return;
724                            } else {
725                                debug!("Re-authenticated async connection");
726                            }
727                        }
728                        Err(err) => {
729                            error!("Credentials stream error for async connection: {err}.");
730                        }
731                    }
732                }
733                warn!("Credentials stream ended; no further re-authentication will occur.");
734            });
735            return Ok((
736                Self {
737                    _credentials_subscription_task_handle: Some(SharedHandleContainer::new(
738                        subscription_task_handle,
739                    )),
740                    ..con
741                },
742                driver,
743            ));
744        }
745
746        Ok((con, driver))
747    }
748
749    /// This should be called strictly before the multiplexed connection is cloned - that is, before it is returned to the user.
750    /// Otherwise some clones will be able to kill the backing task, while other clones are still alive.
751    pub(crate) fn set_task_handle(&mut self, handle: TaskHandle) {
752        self._task_handle = Some(SharedHandleContainer::new(handle));
753    }
754
755    /// Sets the time that the multiplexer will wait for responses on operations before failing.
756    pub fn set_response_timeout(&mut self, timeout: std::time::Duration) {
757        self.response_timeout = Some(timeout);
758    }
759
760    /// Sends an already encoded (packed) command into the TCP socket and
761    /// reads the single response from it.
762    pub async fn send_packed_command(&mut self, cmd: &Cmd) -> RedisResult<Value> {
763        let _permit = if cmd.skip_concurrency_limit {
764            None
765        } else if let Some(limiter) = &self.concurrency_limiter {
766            Some(limiter.semaphore.acquire().await)
767        } else {
768            None
769        };
770        #[cfg(feature = "cache-aio")]
771        if let Some(cache_manager) = &self.cache_manager {
772            match cache_manager.get_cached_cmd(cmd) {
773                PrepCacheItem::Cached(value) => return Ok(value),
774                PrepCacheItem::NotCached(cacheable_command) => {
775                    let mut pipeline = crate::Pipeline::new();
776                    cacheable_command.pack_command(cache_manager, &mut pipeline);
777
778                    let result = self
779                        .pipeline
780                        .send_recv(
781                            pipeline.get_packed_pipeline(),
782                            Some(PipelineResponseExpectation {
783                                skipped_response_count: 0,
784                                expected_response_count: pipeline.commands.len(),
785                                is_transaction: false,
786                                seen_responses: 0,
787                            }),
788                            self.response_timeout,
789                            cmd.is_no_response(),
790                        )
791                        .await?;
792                    let replies: Vec<Value> = crate::types::from_redis_value(result)?;
793                    return cacheable_command.resolve(cache_manager, replies.into_iter());
794                }
795                PrepCacheItem::NotCacheable => (),
796            }
797        }
798        self.pipeline
799            .send_recv(
800                cmd.get_packed_command(),
801                None,
802                self.response_timeout,
803                cmd.is_no_response(),
804            )
805            .await
806    }
807
808    /// Sends multiple already encoded (packed) command into the TCP socket
809    /// and reads `count` responses from it.  This is used to implement
810    /// pipelining.
811    pub async fn send_packed_commands(
812        &mut self,
813        cmd: &crate::Pipeline,
814        offset: usize,
815        count: usize,
816    ) -> RedisResult<Vec<Value>> {
817        // Try to acquire 1 permit per command in the pipeline: block on the first to guarantee
818        // progress, then grab as many more as are immediately available without blocking.
819        // This roughly reflects the pipeline's load on the server while avoiding deadlock --
820        // a large pipeline can always proceed even if it can't acquire all permits.
821        let _permits = if let Some(limiter) = &self.concurrency_limiter {
822            let capacity = permit_vec_capacity(count, limiter.limit);
823            let mut permits = Vec::with_capacity(capacity);
824            permits.push(limiter.semaphore.acquire().await);
825            for _ in 1..capacity {
826                match limiter.semaphore.try_acquire() {
827                    Some(permit) => permits.push(permit),
828                    None => break,
829                }
830            }
831            permits
832        } else {
833            Vec::new()
834        };
835        #[cfg(feature = "cache-aio")]
836        if let Some(cache_manager) = &self.cache_manager {
837            let (cacheable_pipeline, pipeline, (skipped_response_count, expected_response_count)) =
838                cache_manager.get_cached_pipeline(cmd);
839            if pipeline.is_empty() {
840                return cacheable_pipeline.resolve(cache_manager, Value::Array(Vec::new()));
841            }
842            let result = self
843                .pipeline
844                .send_recv(
845                    pipeline.get_packed_pipeline(),
846                    Some(PipelineResponseExpectation {
847                        skipped_response_count,
848                        expected_response_count,
849                        is_transaction: cacheable_pipeline.transaction_mode,
850                        seen_responses: 0,
851                    }),
852                    self.response_timeout,
853                    false,
854                )
855                .await?;
856
857            return cacheable_pipeline.resolve(cache_manager, result);
858        }
859        let value = self
860            .pipeline
861            .send_recv(
862                cmd.get_packed_pipeline(),
863                Some(PipelineResponseExpectation {
864                    skipped_response_count: offset,
865                    expected_response_count: count,
866                    is_transaction: cmd.is_transaction(),
867                    seen_responses: 0,
868                }),
869                self.response_timeout,
870                false,
871            )
872            .await?;
873        match value {
874            Value::Array(values) => Ok(values),
875            _ => Ok(vec![value]),
876        }
877    }
878
879    /// Gets [`CacheStatistics`] for current connection if caching is enabled.
880    #[cfg(feature = "cache-aio")]
881    #[cfg_attr(docsrs, doc(cfg(feature = "cache-aio")))]
882    pub fn get_cache_statistics(&self) -> Option<CacheStatistics> {
883        self.cache_manager.as_ref().map(|cm| cm.statistics())
884    }
885}
886
887impl ConnectionLike for MultiplexedConnection {
888    fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
889        (async move { self.send_packed_command(cmd).await }).boxed()
890    }
891
892    fn req_packed_commands<'a>(
893        &'a mut self,
894        pipeline: &'a crate::Pipeline,
895        offset: usize,
896        count: usize,
897    ) -> RedisFuture<'a, Vec<Value>> {
898        (async move { self.send_packed_commands(pipeline, offset, count).await }).boxed()
899    }
900
901    fn get_db(&self) -> i64 {
902        self.db
903    }
904}
905
906impl MultiplexedConnection {
907    /// Subscribes to a new channel(s).    
908    ///
909    /// Updates from the sender will be sent on the push sender that was passed to the connection.
910    /// If the connection was configured without a push sender, the connection won't be able to pass messages back to the user.
911    ///
912    /// This method is only available when the connection is using RESP3 protocol, and will return an error otherwise.
913    ///
914    /// ```rust,no_run
915    /// # async fn func() -> redis::RedisResult<()> {
916    /// let client = redis::Client::open("redis://127.0.0.1/?protocol=resp3").unwrap();
917    /// let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
918    /// let config = redis::AsyncConnectionConfig::new().set_push_sender(tx);
919    /// let mut con = client.get_multiplexed_async_connection_with_config(&config).await?;
920    /// con.subscribe(&["channel_1", "channel_2"]).await?;
921    /// # Ok(()) }
922    /// ```
923    pub async fn subscribe(&mut self, channel_name: impl ToRedisArgs) -> RedisResult<()> {
924        check_resp3!(self.protocol);
925        let mut cmd = cmd("SUBSCRIBE");
926        cmd.arg(channel_name);
927        cmd.exec_async(self).await?;
928        Ok(())
929    }
930
931    /// Unsubscribes from channel(s).
932    ///
933    /// This method is only available when the connection is using RESP3 protocol, and will return an error otherwise.
934    ///
935    /// ```rust,no_run
936    /// # async fn func() -> redis::RedisResult<()> {
937    /// let client = redis::Client::open("redis://127.0.0.1/?protocol=resp3").unwrap();
938    /// let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
939    /// let config = redis::AsyncConnectionConfig::new().set_push_sender(tx);
940    /// let mut con = client.get_multiplexed_async_connection_with_config(&config).await?;
941    /// con.subscribe(&["channel_1", "channel_2"]).await?;
942    /// con.unsubscribe(&["channel_1", "channel_2"]).await?;
943    /// # Ok(()) }
944    /// ```
945    pub async fn unsubscribe(&mut self, channel_name: impl ToRedisArgs) -> RedisResult<()> {
946        check_resp3!(self.protocol);
947        let mut cmd = cmd("UNSUBSCRIBE");
948        cmd.arg(channel_name);
949        cmd.exec_async(self).await?;
950        Ok(())
951    }
952
953    /// Subscribes to new channel(s) with pattern(s).
954    ///
955    /// Updates from the sender will be sent on the push sender that was passed to the connection.
956    /// If the connection was configured without a push sender, the connection won't be able to pass messages back to the user.
957    ///
958    /// This method is only available when the connection is using RESP3 protocol, and will return an error otherwise.
959    ///
960    /// ```rust,no_run
961    /// # async fn func() -> redis::RedisResult<()> {
962    /// let client = redis::Client::open("redis://127.0.0.1/?protocol=resp3").unwrap();
963    /// let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
964    /// let config = redis::AsyncConnectionConfig::new().set_push_sender(tx);
965    /// let mut con = client.get_multiplexed_async_connection_with_config(&config).await?;
966    /// con.psubscribe("channel*_1").await?;
967    /// con.psubscribe(&["channel*_2", "channel*_3"]).await?;
968    /// # Ok(())
969    /// # }
970    /// ```
971    pub async fn psubscribe(&mut self, channel_pattern: impl ToRedisArgs) -> RedisResult<()> {
972        check_resp3!(self.protocol);
973        let mut cmd = cmd("PSUBSCRIBE");
974        cmd.arg(channel_pattern);
975        cmd.exec_async(self).await?;
976        Ok(())
977    }
978
979    /// Unsubscribes from channel pattern(s).
980    ///
981    /// This method is only available when the connection is using RESP3 protocol, and will return an error otherwise.
982    pub async fn punsubscribe(&mut self, channel_pattern: impl ToRedisArgs) -> RedisResult<()> {
983        check_resp3!(self.protocol);
984        let mut cmd = cmd("PUNSUBSCRIBE");
985        cmd.arg(channel_pattern);
986        cmd.exec_async(self).await?;
987        Ok(())
988    }
989}
990
991#[cfg(feature = "token-based-authentication")]
992impl MultiplexedConnection {
993    /// Re-authenticate the connection with new credentials
994    ///
995    /// This method allows existing async connections to update their authentication
996    /// when tokens are refreshed, enabling streaming credential updates.
997    async fn re_authenticate_with_credentials(
998        &mut self,
999        credentials: &crate::auth::BasicAuth,
1000    ) -> RedisResult<()> {
1001        let mut auth_cmd =
1002            crate::connection::authenticate_cmd(Some(&credentials.username), &credentials.password);
1003        auth_cmd.skip_concurrency_limit = true;
1004        self.send_packed_command(&auth_cmd)
1005            .await?
1006            .extract_error()
1007            .map(|_| ())
1008    }
1009}
1010
1011#[cfg(test)]
1012mod tests {
1013    use super::*;
1014
1015    #[test]
1016    fn test_pipeline_resolve_buffer_size_default() {
1017        assert_eq!(Pipeline::resolve_buffer_size(None), 50);
1018    }
1019
1020    #[test]
1021    fn test_pipeline_resolve_buffer_size_custom() {
1022        assert_eq!(Pipeline::resolve_buffer_size(Some(100)), 100);
1023    }
1024
1025    #[test]
1026    fn test_permit_vec_capacity_is_bounded_by_limit() {
1027        // Capacity tracks the limit, never the (unbounded) pipeline length.
1028        assert_eq!(permit_vec_capacity(1_000_000, 4), 4);
1029        assert_eq!(permit_vec_capacity(10, 10), 10);
1030        assert_eq!(permit_vec_capacity(2, 4), 2);
1031        assert_eq!(permit_vec_capacity(0, 4), 1);
1032        assert_eq!(permit_vec_capacity(5, 0), 1);
1033    }
1034
1035    #[test]
1036    fn test_build_concurrency_limiter_rejects_zero() {
1037        // A zero-permit semaphore would deadlock every command on acquire.
1038        let err = build_concurrency_limiter(Some(0)).unwrap_err();
1039        assert_eq!(err.kind(), ErrorKind::InvalidClientConfig);
1040    }
1041
1042    #[test]
1043    fn test_build_concurrency_limiter_allows_none_and_positive() {
1044        assert!(build_concurrency_limiter(None).unwrap().is_none());
1045        let limiter = build_concurrency_limiter(Some(4)).unwrap().unwrap();
1046        assert_eq!(limiter.limit, 4);
1047    }
1048
1049    fn mock_conn_info() -> RedisConnectionInfo {
1050        RedisConnectionInfo {
1051            skip_set_lib_name: true,
1052            ..Default::default()
1053        }
1054    }
1055
1056    async fn create_mock_connection(
1057        concurrency_limit: usize,
1058    ) -> (
1059        MultiplexedConnection,
1060        tokio::sync::mpsc::Receiver<()>,
1061        tokio::sync::mpsc::Sender<()>,
1062    ) {
1063        use futures_util::StreamExt;
1064        use tokio::io::AsyncWriteExt;
1065        use tokio_util::codec::FramedRead;
1066
1067        let (client_half, server_half) = tokio::io::duplex(4096);
1068        let (cmd_received_tx, cmd_received_rx) = tokio::sync::mpsc::channel::<()>(10);
1069        let (send_response_tx, mut send_response_rx) = tokio::sync::mpsc::channel::<()>(10);
1070
1071        let (server_read, mut server_write) = tokio::io::split(server_half);
1072
1073        tokio::spawn(async move {
1074            let mut reader = FramedRead::new(server_read, ValueCodec::default());
1075            while let Some(Ok(_)) = reader.next().await {
1076                let _ = cmd_received_tx.send(()).await;
1077            }
1078        });
1079
1080        tokio::spawn(async move {
1081            while send_response_rx.recv().await.is_some() {
1082                let _ = server_write.write_all(b"+OK\r\n").await;
1083                let _ = server_write.flush().await;
1084            }
1085        });
1086
1087        let config = AsyncConnectionConfig::new()
1088            .set_concurrency_limit(concurrency_limit)
1089            .set_response_timeout(None)
1090            .set_connection_timeout(None);
1091
1092        let (conn, driver) =
1093            MultiplexedConnection::new_with_config(&mock_conn_info(), client_half, config)
1094                .await
1095                .unwrap();
1096        tokio::spawn(driver);
1097
1098        (conn, cmd_received_rx, send_response_tx)
1099    }
1100
1101    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1102    async fn test_concurrency_limit_enforced() {
1103        let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(2).await;
1104
1105        let h1 = tokio::spawn({
1106            let mut c = conn.clone();
1107            async move { c.send_packed_command(&cmd("PING")).await }
1108        });
1109        let h2 = tokio::spawn({
1110            let mut c = conn.clone();
1111            async move { c.send_packed_command(&cmd("PING")).await }
1112        });
1113        let h3 = tokio::spawn({
1114            let mut c = conn.clone();
1115            async move { c.send_packed_command(&cmd("PING")).await }
1116        });
1117
1118        cmd_received_rx.recv().await.unwrap();
1119        cmd_received_rx.recv().await.unwrap();
1120
1121        let third = tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1122        assert!(
1123            third.is_err(),
1124            "3rd request should be blocked by concurrency limit"
1125        );
1126
1127        send_response_tx.send(()).await.unwrap();
1128
1129        cmd_received_rx.recv().await.unwrap();
1130
1131        send_response_tx.send(()).await.unwrap();
1132        send_response_tx.send(()).await.unwrap();
1133
1134        h1.await.unwrap().unwrap();
1135        h2.await.unwrap().unwrap();
1136        h3.await.unwrap().unwrap();
1137    }
1138
1139    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1140    async fn test_no_limit_bypasses_concurrency_limit() {
1141        let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(1).await;
1142
1143        let h1 = tokio::spawn({
1144            let mut c = conn.clone();
1145            async move { c.send_packed_command(&cmd("PING")).await }
1146        });
1147
1148        cmd_received_rx.recv().await.unwrap();
1149
1150        let h2 = tokio::spawn({
1151            let mut c = conn.clone();
1152            async move {
1153                let mut ping = cmd("PING");
1154                ping.skip_concurrency_limit = true;
1155                c.send_packed_command(&ping).await
1156            }
1157        });
1158
1159        let received =
1160            tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1161        assert!(
1162            received.is_ok(),
1163            "no_limit request should bypass concurrency limit"
1164        );
1165
1166        send_response_tx.send(()).await.unwrap();
1167        send_response_tx.send(()).await.unwrap();
1168
1169        h1.await.unwrap().unwrap();
1170        h2.await.unwrap().unwrap();
1171    }
1172
1173    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1174    async fn test_pipeline_acquires_multiple_permits() {
1175        let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(3).await;
1176
1177        let pipeline_handle = tokio::spawn({
1178            let mut c = conn.clone();
1179            async move {
1180                let mut pipe = crate::Pipeline::new();
1181                pipe.cmd("SET").arg("a").arg("1");
1182                pipe.cmd("SET").arg("b").arg("2");
1183                pipe.cmd("SET").arg("c").arg("3");
1184                c.send_packed_commands(&pipe, 0, 3).await
1185            }
1186        });
1187
1188        for _ in 0..3 {
1189            cmd_received_rx.recv().await.unwrap();
1190        }
1191
1192        let single_handle = tokio::spawn({
1193            let mut c = conn.clone();
1194            async move { c.send_packed_command(&cmd("PING")).await }
1195        });
1196
1197        let blocked =
1198            tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1199        assert!(
1200            blocked.is_err(),
1201            "single command should be blocked while pipeline holds all permits"
1202        );
1203
1204        for _ in 0..3 {
1205            send_response_tx.send(()).await.unwrap();
1206        }
1207
1208        cmd_received_rx.recv().await.unwrap();
1209        send_response_tx.send(()).await.unwrap();
1210
1211        pipeline_handle.await.unwrap().unwrap();
1212        single_handle.await.unwrap().unwrap();
1213    }
1214
1215    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1216    async fn test_pipeline_proceeds_with_partial_permits() {
1217        let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(2).await;
1218
1219        let single_handle = tokio::spawn({
1220            let mut c = conn.clone();
1221            async move { c.send_packed_command(&cmd("PING")).await }
1222        });
1223        cmd_received_rx.recv().await.unwrap();
1224
1225        let pipeline_handle = tokio::spawn({
1226            let mut c = conn.clone();
1227            async move {
1228                let mut pipe = crate::Pipeline::new();
1229                for i in 0..5 {
1230                    pipe.cmd("SET").arg(format!("k{i}")).arg(i);
1231                }
1232                c.send_packed_commands(&pipe, 0, 5).await
1233            }
1234        });
1235
1236        let received =
1237            tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1238        assert!(
1239            received.is_ok(),
1240            "pipeline should proceed even with only partial permits"
1241        );
1242
1243        for _ in 1..5 {
1244            cmd_received_rx.recv().await.unwrap();
1245        }
1246
1247        for _ in 0..6 {
1248            send_response_tx.send(()).await.unwrap();
1249        }
1250
1251        single_handle.await.unwrap().unwrap();
1252        pipeline_handle.await.unwrap().unwrap();
1253    }
1254
1255    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1256    async fn test_permit_released_on_cancellation() {
1257        let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(1).await;
1258
1259        let h1 = tokio::spawn({
1260            let mut c = conn.clone();
1261            async move { c.send_packed_command(&cmd("PING")).await }
1262        });
1263        cmd_received_rx.recv().await.unwrap();
1264
1265        // Start a second request that will block on the semaphore, then cancel it
1266        let h2 = tokio::spawn({
1267            let mut c = conn.clone();
1268            async move { c.send_packed_command(&cmd("PING")).await }
1269        });
1270        tokio::time::sleep(Duration::from_millis(50)).await;
1271        h2.abort();
1272        let _ = h2.await;
1273
1274        // Complete the first request
1275        send_response_tx.send(()).await.unwrap();
1276        h1.await.unwrap().unwrap();
1277
1278        // The permit from the cancelled request should have been released,
1279        // so a new request should proceed
1280        let h3 = tokio::spawn({
1281            let mut c = conn.clone();
1282            async move { c.send_packed_command(&cmd("PING")).await }
1283        });
1284
1285        let received =
1286            tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1287        assert!(
1288            received.is_ok(),
1289            "request after cancellation should acquire the permit"
1290        );
1291
1292        send_response_tx.send(()).await.unwrap();
1293        h3.await.unwrap().unwrap();
1294    }
1295
1296    /// Regression test for the TCP buffer deadlock reported in
1297    /// <https://github.com/redis-rs/redis-rs/issues/1955>.
1298    ///
1299    /// # Why this can deadlock
1300    ///
1301    /// A TCP buffer deadlock is a property of *both* peers — neither can
1302    /// produce it alone.
1303    ///
1304    /// **Client side (this bug):** `PipelineSink::poll_flush` polled the read
1305    /// half only *after* the underlying writer's flush returned Ready. When
1306    /// our TCP send buffer fills, the codec's flush is Pending, and we
1307    /// returned Pending without ever registering a read waker — so response
1308    /// bytes already sitting in our recv buffer were never polled. The driver
1309    /// task parked indefinitely.
1310    ///
1311    /// **Server side:** any server that stops reading from a client's
1312    /// socket while its own pending write to that client can't make
1313    /// forward progress is enough to trigger the deadlock. Combined with
1314    /// the client bug, both directions wedge: the client's send buffer
1315    /// can't drain (server isn't reading), the server's send buffer can't
1316    /// drain (client isn't reading because of the bug), and the state is
1317    /// permanent.
1318    ///
1319    /// # What this test simulates
1320    ///
1321    /// `tokio::io::duplex` stands in for a TCP socket pair with tiny
1322    /// kernel buffers. The mini-server is a loop that reads a request,
1323    /// writes a response, repeats — awaiting the write makes the loop
1324    /// stop reading the instant a write becomes Pending, which is the
1325    /// general server-side behavior described above.
1326    ///
1327    /// The test issues a few large concurrent SETs whose request payloads
1328    /// (and pretend responses) each exceed the duplex buffer in both
1329    /// directions. With the bug, the client's codec backs up, the
1330    /// server's writes back up, and both sides park without enough wakers
1331    /// to escape — the test times out and panics. With the fix, the
1332    /// driver drains responses even while the writer is back-pressured,
1333    /// which keeps the duplex moving and lets the SETs complete.
1334    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1335    async fn test_deadlock_when_writes_blocked_with_pending_response() {
1336        use futures_util::StreamExt;
1337        use tokio::io::AsyncWriteExt;
1338        use tokio_util::codec::FramedRead;
1339
1340        // Small duplex buffer + ~4 KiB request/response sizes. The polling
1341        // pathology doesn't depend on scale; a real socket would see the
1342        // same shape with tens of KiB of buffer and MB-scale payloads.
1343        const BUFFER_SIZE: usize = 256;
1344        const PAYLOAD_SIZE: usize = 4096;
1345        const REQUEST_COUNT: usize = 3;
1346
1347        let (client_half, server_half) = tokio::io::duplex(BUFFER_SIZE);
1348        let (server_read, mut server_write) = tokio::io::split(server_half);
1349
1350        // Pretend response: a bulk string the same size as the request
1351        // payload, so server-bound writes also exceed the duplex buffer
1352        // (the server's write pends and its loop stops reading).
1353        let mut response = Vec::with_capacity(PAYLOAD_SIZE + 16);
1354        response.extend_from_slice(format!("${PAYLOAD_SIZE}\r\n").as_bytes());
1355        response.extend(std::iter::repeat_n(b'V', PAYLOAD_SIZE));
1356        response.extend_from_slice(b"\r\n");
1357
1358        // Mini-server: read a request, write a response, loop. Awaiting
1359        // the write makes the loop stop reading the instant a write
1360        // becomes Pending — the general "server stops reading once its
1361        // own write is back-pressured" behavior the deadlock requires.
1362        let server_task = tokio::spawn(async move {
1363            let mut reader = FramedRead::new(server_read, ValueCodec::default());
1364            loop {
1365                match reader.next().await {
1366                    Some(Ok(_)) => {}
1367                    _ => return,
1368                }
1369                if server_write.write_all(&response).await.is_err() {
1370                    return;
1371                }
1372            }
1373        });
1374
1375        let config = AsyncConnectionConfig::new()
1376            .set_response_timeout(None)
1377            .set_connection_timeout(None);
1378        let (conn, driver) =
1379            MultiplexedConnection::new_with_config(&mock_conn_info(), client_half, config)
1380                .await
1381                .unwrap();
1382        let driver_handle = tokio::spawn(driver);
1383
1384        // A handful of concurrent large SETs. Each request exceeds the
1385        // duplex buffer (codec's flush stays backed up); each reply also
1386        // exceeds the duplex buffer (server's write stays backed up). Both
1387        // directions wedge.
1388        let mut handles = Vec::with_capacity(REQUEST_COUNT);
1389        for i in 0..REQUEST_COUNT {
1390            let mut c = conn.clone();
1391            handles.push(tokio::spawn(async move {
1392                let mut set = cmd("SET");
1393                set.arg(format!("k{i}")).arg(vec![b'X'; PAYLOAD_SIZE]);
1394                c.send_packed_command(&set).await
1395            }));
1396        }
1397
1398        let join_all = async move {
1399            let mut results = Vec::with_capacity(handles.len());
1400            for h in handles {
1401                results.push(h.await);
1402            }
1403            results
1404        };
1405
1406        let outcome = tokio::time::timeout(Duration::from_secs(5), join_all).await;
1407
1408        // Clean up before asserting so a panic doesn't leak tasks.
1409        driver_handle.abort();
1410        server_task.abort();
1411
1412        let results = outcome.expect(
1413            "DEADLOCK reproduced: client driver parked in poll_flush with no \
1414             read waker registered. Server has buffered responses in the duplex \
1415             and stopped reading once its own write became Pending; the client \
1416             cannot send the rest of its requests because the server is no \
1417             longer draining the link. Both sides wedged.",
1418        );
1419        for (i, res) in results.into_iter().enumerate() {
1420            let join = res.unwrap_or_else(|e| panic!("SET task {i} panicked: {e}"));
1421            join.unwrap_or_else(|e| panic!("SET task {i} returned an error: {e}"));
1422        }
1423    }
1424
1425    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1426    async fn test_permit_released_on_response_timeout() {
1427        use futures_util::StreamExt;
1428        use tokio::io::AsyncWriteExt;
1429        use tokio_util::codec::FramedRead;
1430
1431        let (client_half, server_half) = tokio::io::duplex(4096);
1432        let (cmd_received_tx, mut cmd_received_rx) = tokio::sync::mpsc::channel::<()>(10);
1433
1434        let (server_read, mut server_write) = tokio::io::split(server_half);
1435
1436        tokio::spawn(async move {
1437            let mut reader = FramedRead::new(server_read, ValueCodec::default());
1438            while let Some(Ok(_)) = reader.next().await {
1439                let _ = cmd_received_tx.send(()).await;
1440            }
1441        });
1442
1443        tokio::spawn(async move {
1444            futures_util::future::pending::<()>().await;
1445            let _ = server_write.write_all(b"").await;
1446        });
1447
1448        let config = AsyncConnectionConfig::new()
1449            .set_concurrency_limit(1)
1450            .set_response_timeout(Some(Duration::from_millis(100)))
1451            .set_connection_timeout(None);
1452
1453        let (conn, driver) =
1454            MultiplexedConnection::new_with_config(&mock_conn_info(), client_half, config)
1455                .await
1456                .unwrap();
1457        tokio::spawn(driver);
1458
1459        // First request times out since the mock never responds
1460        let mut c1 = conn.clone();
1461        let err = c1.send_packed_command(&cmd("PING")).await.unwrap_err();
1462        assert!(err.is_io_error(), "expected IO error from timeout");
1463        cmd_received_rx.recv().await.unwrap();
1464
1465        // Second request should acquire the permit released by the first,
1466        // reach the server, and then also time out
1467        let mut c2 = conn.clone();
1468        let err = c2.send_packed_command(&cmd("PING")).await.unwrap_err();
1469        assert!(err.is_io_error(), "expected IO error from timeout");
1470        cmd_received_rx.recv().await.unwrap();
1471    }
1472}