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
40type PipelineOutput = oneshot::Sender<RedisResult<Value>>;
42
43enum ErrorOrErrors {
44 Errors(Vec<(usize, ServerError)>),
45 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
58struct PipelineResponseExpectation {
60 skipped_response_count: usize,
62 expected_response_count: usize,
64 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
87struct PipelineMessage {
89 input: Vec<u8>,
90 output: Option<PipelineOutput>,
92 expectation: Option<PipelineResponseExpectation>,
96}
97
98#[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 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 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 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 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 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 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)); }
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 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 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 fn poll_ready(
313 mut self: Pin<&mut Self>,
314 cx: &mut task::Context,
315 ) -> Poll<Result<(), Self::Error>> {
316 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 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 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 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 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 .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#[derive(Clone, Debug)]
509struct ConcurrencyLimiter {
510 limit: usize,
511 semaphore: Arc<async_lock::Semaphore>,
512}
513
514fn permit_vec_capacity(count: usize, limit: usize) -> usize {
517 count.min(limit).max(1)
518}
519
520fn 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#[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 _task_handle: Option<SharedHandleContainer>,
558 #[cfg(feature = "cache-aio")]
559 pub(crate) cache_manager: Option<CacheManager>,
560 #[cfg(feature = "token-based-authentication")]
561 _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 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 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 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 #[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 pub(crate) fn set_task_handle(&mut self, handle: TaskHandle) {
763 self._task_handle = Some(SharedHandleContainer::new(handle));
764 }
765
766 pub fn set_response_timeout(&mut self, timeout: std::time::Duration) {
768 self.response_timeout = Some(timeout);
769 }
770
771 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 pub async fn send_packed_commands(
823 &mut self,
824 cmd: &crate::Pipeline,
825 offset: usize,
826 count: usize,
827 ) -> RedisResult<Vec<Value>> {
828 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 #[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 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 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 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 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 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 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 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 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 send_response_tx.send(()).await.unwrap();
1287 h1.await.unwrap().unwrap();
1288
1289 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 #[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 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 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 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 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 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 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 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}