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