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
142pub(crate) fn send_disconnect(push_sender: &Option<Arc<dyn AsyncPushSender>>) {
143 send_push(push_sender, PushInfo::disconnect());
144}
145
146impl<T> PipelineSink<T>
147where
148 T: Stream<Item = RedisResult<Value>> + 'static,
149{
150 fn new(
151 sink_stream: T,
152 push_sender: Option<Arc<dyn AsyncPushSender>>,
153 #[cfg(feature = "cache-aio")] cache_manager: Option<CacheManager>,
154 ) -> Self
155 where
156 T: Sink<Vec<u8>, Error = RedisError> + Stream<Item = RedisResult<Value>> + 'static,
157 {
158 Self {
159 sink_stream,
160 in_flight: VecDeque::new(),
161 error: None,
162 push_sender,
163 #[cfg(feature = "cache-aio")]
164 cache_manager,
165 }
166 }
167
168 fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<(), ()>> {
170 loop {
171 let item = ready!(self.as_mut().project().sink_stream.poll_next(cx));
172 let item = match item {
173 Some(result) => result,
174 None => Err(closed_connection_error()),
176 };
177
178 let is_unrecoverable = item.as_ref().is_err_and(|err| err.is_unrecoverable_error());
179 self.as_mut().send_result(item);
180 if is_unrecoverable {
181 let self_ = self.project();
182 send_disconnect(self_.push_sender);
183 return Poll::Ready(Err(()));
184 }
185 }
186 }
187
188 fn send_result(self: Pin<&mut Self>, result: RedisResult<Value>) {
189 let self_ = self.project();
190 let result = match result {
191 Ok(Value::Push { kind, data }) if !kind.has_reply() => {
193 #[cfg(feature = "cache-aio")]
194 if let Some(cache_manager) = &self_.cache_manager {
195 cache_manager.handle_push_value(&kind, &data);
196 }
197 send_push(self_.push_sender, PushInfo { kind, data });
198
199 return;
200 }
201 Ok(Value::Push { kind, data }) if kind.has_reply() => {
203 send_push(
204 self_.push_sender,
205 PushInfo {
206 kind: kind.clone(),
207 data: data.clone(),
208 },
209 );
210 Ok(Value::Push { kind, data })
211 }
212 _ => result,
213 };
214
215 let mut entry = match self_.in_flight.pop_front() {
216 Some(entry) => entry,
217 None => return,
218 };
219
220 match &mut entry.response_aggregate {
221 ResponseAggregate::SingleCommand => {
222 if let Some(output) = entry.output.take() {
223 _ = output.send(result);
224 }
225 }
226 ResponseAggregate::Pipeline {
227 buffer,
228 error_or_errors,
229 expectation:
230 PipelineResponseExpectation {
231 expected_response_count,
232 skipped_response_count,
233 is_transaction,
234 seen_responses,
235 },
236 } => {
237 *seen_responses += 1;
238 if *skipped_response_count > 0 {
239 if *is_transaction && let ErrorOrErrors::Errors(errs) = error_or_errors {
242 match result {
243 Ok(Value::ServerError(err)) => {
244 errs.push((*seen_responses - 2, err)); }
246 Err(err) => *error_or_errors = ErrorOrErrors::FirstError(err),
247 _ => {}
248 }
249 }
250
251 *skipped_response_count -= 1;
252 self_.in_flight.push_front(entry);
253 return;
254 }
255
256 match result {
257 Ok(item) => {
258 buffer.push(item);
259 }
260 Err(err) => {
261 if matches!(error_or_errors, ErrorOrErrors::Errors(_)) {
262 *error_or_errors = ErrorOrErrors::FirstError(err);
263 }
264 }
265 }
266
267 if buffer.len() < *expected_response_count {
268 self_.in_flight.push_front(entry);
270 return;
271 }
272
273 let response =
274 match std::mem::replace(error_or_errors, ErrorOrErrors::Errors(Vec::new())) {
275 ErrorOrErrors::Errors(errors) => {
276 if errors.is_empty() {
277 Ok(Value::Array(std::mem::take(buffer)))
278 } else {
279 Err(RedisError::make_aborted_transaction(errors))
280 }
281 }
282 ErrorOrErrors::FirstError(redis_error) => Err(redis_error),
283 };
284
285 if let Some(output) = entry.output.take() {
289 _ = output.send(response);
290 }
291 }
292 }
293 }
294}
295
296impl<T> Sink<PipelineMessage> for PipelineSink<T>
297where
298 T: Sink<Vec<u8>, Error = RedisError> + Stream<Item = RedisResult<Value>> + 'static,
299{
300 type Error = ();
301
302 fn poll_ready(
304 mut self: Pin<&mut Self>,
305 cx: &mut task::Context,
306 ) -> Poll<Result<(), Self::Error>> {
307 if matches!(self.as_mut().poll_read(cx), Poll::Ready(Err(()))) {
314 return Poll::Ready(Err(()));
315 }
316 match ready!(self.as_mut().project().sink_stream.poll_ready(cx)) {
317 Ok(()) => Ok(()).into(),
318 Err(err) => {
319 *self.project().error = Some(err);
320 Ok(()).into()
321 }
322 }
323 }
324
325 fn start_send(
326 mut self: Pin<&mut Self>,
327 PipelineMessage {
328 input,
329 mut output,
330 expectation,
331 }: PipelineMessage,
332 ) -> Result<(), Self::Error> {
333 if output.as_ref().is_some_and(|output| output.is_closed()) {
337 return Ok(());
338 }
339
340 let self_ = self.as_mut().project();
341
342 if let Some(err) = self_.error.take() {
343 if let Some(output) = output.take() {
344 _ = output.send(Err(err));
345 }
346 return Err(());
347 }
348
349 match self_.sink_stream.start_send(input) {
350 Ok(()) => {
351 let response_aggregate = ResponseAggregate::new(expectation);
352 let entry = InFlight {
353 output,
354 response_aggregate,
355 };
356
357 self_.in_flight.push_back(entry);
358 Ok(())
359 }
360 Err(err) => {
361 if let Some(output) = output.take() {
362 _ = output.send(Err(err));
363 }
364 Err(())
365 }
366 }
367 }
368
369 fn poll_flush(
370 mut self: Pin<&mut Self>,
371 cx: &mut task::Context,
372 ) -> Poll<Result<(), Self::Error>> {
373 if matches!(self.as_mut().poll_read(cx), Poll::Ready(Err(()))) {
380 return Poll::Ready(Err(()));
381 }
382 self.as_mut()
383 .project()
384 .sink_stream
385 .poll_flush(cx)
386 .map_err(|err| {
387 self.as_mut().send_result(Err(err));
388 })
389 }
390
391 fn poll_close(
392 mut self: Pin<&mut Self>,
393 cx: &mut task::Context,
394 ) -> Poll<Result<(), Self::Error>> {
395 if !self.in_flight.is_empty() {
398 ready!(self.as_mut().poll_flush(cx))?;
399 }
400 let this = self.as_mut().project();
401 this.sink_stream.poll_close(cx).map_err(|err| {
402 self.send_result(Err(err));
403 })
404 }
405}
406
407impl Pipeline {
408 const DEFAULT_BUFFER_SIZE: usize = 50;
409
410 fn resolve_buffer_size(size: Option<usize>) -> usize {
411 size.unwrap_or(Self::DEFAULT_BUFFER_SIZE)
412 }
413
414 fn new<T>(
415 sink_stream: T,
416 push_sender: Option<Arc<dyn AsyncPushSender>>,
417 #[cfg(feature = "cache-aio")] cache_manager: Option<CacheManager>,
418 buffer_size: usize,
419 ) -> (Self, impl Future<Output = ()>)
420 where
421 T: Sink<Vec<u8>, Error = RedisError>,
422 T: Stream<Item = RedisResult<Value>>,
423 T: Unpin + Send + 'static,
424 {
425 let (sender, mut receiver) = mpsc::channel(buffer_size);
426
427 let sink = PipelineSink::new(
428 sink_stream,
429 push_sender,
430 #[cfg(feature = "cache-aio")]
431 cache_manager,
432 );
433 let f = stream::poll_fn(move |cx| receiver.poll_recv(cx))
434 .map(Ok)
435 .forward(sink)
436 .map(|_| ());
437 (Self { sender }, f)
438 }
439
440 async fn send_recv(
441 &mut self,
442 input: Vec<u8>,
443 expectation: Option<PipelineResponseExpectation>,
446 timeout: Option<Duration>,
447 skip_response: bool,
448 ) -> Result<Value, RedisError> {
449 if input.is_empty() {
450 return Err(RedisError::make_empty_command());
451 }
452
453 let request = async {
454 if skip_response {
455 self.sender
456 .send(PipelineMessage {
457 input,
458 expectation,
459 output: None,
460 })
461 .await
462 .map_err(|_| None)?;
463
464 return Ok(Value::Nil);
465 }
466
467 let (sender, receiver) = oneshot::channel();
468
469 self.sender
470 .send(PipelineMessage {
471 input,
472 expectation,
473 output: Some(sender),
474 })
475 .await
476 .map_err(|_| None)?;
477
478 receiver.await
479 .map_err(|_| None)
482 .and_then(|res| res.map_err(Some))
483 };
484
485 match timeout {
486 Some(timeout) => match Runtime::locate().timeout(timeout, request).await {
487 Ok(res) => res,
488 Err(elapsed) => Err(Some(elapsed.into())),
489 },
490 None => request.await,
491 }
492 .map_err(|err| err.unwrap_or_else(closed_connection_error))
493 }
494}
495
496#[derive(Clone, Debug)]
498struct ConcurrencyLimiter {
499 limit: usize,
500 semaphore: Arc<async_lock::Semaphore>,
501}
502
503fn permit_vec_capacity(count: usize, limit: usize) -> usize {
506 count.min(limit).max(1)
507}
508
509fn build_concurrency_limiter(limit: Option<usize>) -> RedisResult<Option<ConcurrencyLimiter>> {
511 match limit {
512 Some(0) => Err(RedisError::from((
513 ErrorKind::InvalidClientConfig,
514 "concurrency limit must be greater than 0",
515 ))),
516 Some(n) => Ok(Some(ConcurrencyLimiter {
517 limit: n,
518 semaphore: Arc::new(async_lock::Semaphore::new(n)),
519 })),
520 None => Ok(None),
521 }
522}
523
524#[derive(Clone)]
537pub struct MultiplexedConnection {
538 pipeline: Pipeline,
539 db: i64,
540 response_timeout: Option<Duration>,
541 protocol: ProtocolVersion,
542 concurrency_limiter: Option<ConcurrencyLimiter>,
543 _task_handle: Option<SharedHandleContainer>,
547 #[cfg(feature = "cache-aio")]
548 pub(crate) cache_manager: Option<CacheManager>,
549 #[cfg(feature = "token-based-authentication")]
550 _credentials_subscription_task_handle: Option<SharedHandleContainer>,
553}
554
555impl Debug for MultiplexedConnection {
556 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
557 let Self {
558 pipeline,
559 db,
560 response_timeout,
561 protocol,
562 concurrency_limiter: _,
563 _task_handle,
564 #[cfg(feature = "cache-aio")]
565 cache_manager: _,
566 #[cfg(feature = "token-based-authentication")]
567 _credentials_subscription_task_handle: _,
568 } = self;
569
570 f.debug_struct("MultiplexedConnection")
571 .field("pipeline", &pipeline)
572 .field("db", &db)
573 .field("response_timeout", &response_timeout)
574 .field("protocol", &protocol)
575 .finish()
576 }
577}
578
579impl MultiplexedConnection {
580 pub async fn new<C>(
583 connection_info: &RedisConnectionInfo,
584 stream: C,
585 ) -> RedisResult<(Self, impl Future<Output = ()>)>
586 where
587 C: Unpin + AsyncRead + AsyncWrite + Send + 'static,
588 {
589 Self::new_with_config(connection_info, stream, AsyncConnectionConfig::default()).await
590 }
591
592 pub async fn new_with_config<C>(
595 connection_info: &RedisConnectionInfo,
596 stream: C,
597 config: AsyncConnectionConfig,
598 ) -> RedisResult<(Self, impl Future<Output = ()> + 'static)>
599 where
600 C: Unpin + AsyncRead + AsyncWrite + Send + 'static,
601 {
602 let mut codec = ValueCodec::default().framed(stream);
603 if let Some(boundary) = config.write_backpressure_boundary {
604 codec.set_backpressure_boundary(boundary);
605 }
606 if config.push_sender.is_some() {
607 check_resp3!(
608 connection_info.protocol,
609 "Can only pass push sender to a connection using RESP3"
610 );
611 }
612
613 #[cfg(feature = "cache-aio")]
614 let cache_config = config.cache.as_ref().map(|cache| match cache {
615 crate::client::Cache::Config(cache_config) => *cache_config,
616 #[cfg(any(feature = "connection-manager", feature = "cluster-async"))]
617 crate::client::Cache::Manager(cache_manager) => cache_manager.cache_config,
618 });
619 #[cfg(feature = "cache-aio")]
620 let cache_manager_opt = config
621 .cache
622 .map(|cache| {
623 check_resp3!(
624 connection_info.protocol,
625 "Can only enable client side caching in a connection using RESP3"
626 );
627 match cache {
628 crate::client::Cache::Config(cache_config) => {
629 Ok(CacheManager::new(cache_config))
630 }
631 #[cfg(any(feature = "connection-manager", feature = "cluster-async"))]
632 crate::client::Cache::Manager(cache_manager) => Ok(cache_manager),
633 }
634 })
635 .transpose()?;
636
637 #[cfg(feature = "token-based-authentication")]
638 let mut connection_info = connection_info.clone();
639 #[cfg(not(feature = "token-based-authentication"))]
640 let connection_info = connection_info.clone();
641
642 #[cfg(feature = "token-based-authentication")]
643 if let Some(ref credentials_provider) = config.credentials_provider {
644 match credentials_provider.subscribe().next().await {
646 Some(Ok(credentials)) => {
647 connection_info.username = Some(ArcStr::from(credentials.username));
648 connection_info.password = Some(ArcStr::from(credentials.password));
649 }
650 Some(Err(err)) => {
651 error!("Error while receiving credentials from stream: {err}");
652 return Err(err);
653 }
654 None => {
655 let err = RedisError::from((
656 ErrorKind::AuthenticationFailed,
657 "Credentials stream closed unexpectedly before yielding credentials!",
658 ));
659 error!("{err}");
660 return Err(err);
661 }
662 }
663 }
664
665 setup_connection(
666 &mut codec,
667 &connection_info,
668 #[cfg(feature = "cache-aio")]
669 cache_config,
670 )
671 .await?;
672 if config.push_sender.is_some() {
673 check_resp3!(
674 connection_info.protocol,
675 "Can only pass push sender to a connection using RESP3"
676 );
677 }
678
679 let (pipeline, driver) = Pipeline::new(
680 codec,
681 config.push_sender,
682 #[cfg(feature = "cache-aio")]
683 cache_manager_opt.clone(),
684 Pipeline::resolve_buffer_size(config.pipeline_buffer_size),
685 );
686
687 let concurrency_limiter = build_concurrency_limiter(config.concurrency_limit)?;
688
689 let con = Self {
690 pipeline,
691 db: connection_info.db,
692 response_timeout: config.response_timeout,
693 protocol: connection_info.protocol,
694 concurrency_limiter,
695 _task_handle: None,
696 #[cfg(feature = "cache-aio")]
697 cache_manager: cache_manager_opt,
698 #[cfg(feature = "token-based-authentication")]
699 _credentials_subscription_task_handle: None,
700 };
701
702 #[cfg(feature = "token-based-authentication")]
704 if let Some(streaming_provider) = config.credentials_provider {
705 let mut inner_connection = con.clone();
706 let mut stream = streaming_provider.subscribe();
707
708 let subscription_task_handle = Runtime::locate().spawn(async move {
709 while let Some(result) = stream.next().await {
710 match result {
711 Ok(credentials) => {
712 if let Err(err) = inner_connection
713 .re_authenticate_with_credentials(&credentials)
714 .await
715 {
716 if err.is_connection_dropped() {
717 warn!(
718 "Re-authentication task ended, connection is dead: {err}"
719 );
720 return;
721 }
722 error!("Failed to re-authenticate async connection: {err}.");
723 return;
724 } else {
725 debug!("Re-authenticated async connection");
726 }
727 }
728 Err(err) => {
729 error!("Credentials stream error for async connection: {err}.");
730 }
731 }
732 }
733 warn!("Credentials stream ended; no further re-authentication will occur.");
734 });
735 return Ok((
736 Self {
737 _credentials_subscription_task_handle: Some(SharedHandleContainer::new(
738 subscription_task_handle,
739 )),
740 ..con
741 },
742 driver,
743 ));
744 }
745
746 Ok((con, driver))
747 }
748
749 pub(crate) fn set_task_handle(&mut self, handle: TaskHandle) {
752 self._task_handle = Some(SharedHandleContainer::new(handle));
753 }
754
755 pub fn set_response_timeout(&mut self, timeout: std::time::Duration) {
757 self.response_timeout = Some(timeout);
758 }
759
760 pub async fn send_packed_command(&mut self, cmd: &Cmd) -> RedisResult<Value> {
763 let _permit = if cmd.skip_concurrency_limit {
764 None
765 } else if let Some(limiter) = &self.concurrency_limiter {
766 Some(limiter.semaphore.acquire().await)
767 } else {
768 None
769 };
770 #[cfg(feature = "cache-aio")]
771 if let Some(cache_manager) = &self.cache_manager {
772 match cache_manager.get_cached_cmd(cmd) {
773 PrepCacheItem::Cached(value) => return Ok(value),
774 PrepCacheItem::NotCached(cacheable_command) => {
775 let mut pipeline = crate::Pipeline::new();
776 cacheable_command.pack_command(cache_manager, &mut pipeline);
777
778 let result = self
779 .pipeline
780 .send_recv(
781 pipeline.get_packed_pipeline(),
782 Some(PipelineResponseExpectation {
783 skipped_response_count: 0,
784 expected_response_count: pipeline.commands.len(),
785 is_transaction: false,
786 seen_responses: 0,
787 }),
788 self.response_timeout,
789 cmd.is_no_response(),
790 )
791 .await?;
792 let replies: Vec<Value> = crate::types::from_redis_value(result)?;
793 return cacheable_command.resolve(cache_manager, replies.into_iter());
794 }
795 PrepCacheItem::NotCacheable => (),
796 }
797 }
798 self.pipeline
799 .send_recv(
800 cmd.get_packed_command(),
801 None,
802 self.response_timeout,
803 cmd.is_no_response(),
804 )
805 .await
806 }
807
808 pub async fn send_packed_commands(
812 &mut self,
813 cmd: &crate::Pipeline,
814 offset: usize,
815 count: usize,
816 ) -> RedisResult<Vec<Value>> {
817 let _permits = if let Some(limiter) = &self.concurrency_limiter {
822 let capacity = permit_vec_capacity(count, limiter.limit);
823 let mut permits = Vec::with_capacity(capacity);
824 permits.push(limiter.semaphore.acquire().await);
825 for _ in 1..capacity {
826 match limiter.semaphore.try_acquire() {
827 Some(permit) => permits.push(permit),
828 None => break,
829 }
830 }
831 permits
832 } else {
833 Vec::new()
834 };
835 #[cfg(feature = "cache-aio")]
836 if let Some(cache_manager) = &self.cache_manager {
837 let (cacheable_pipeline, pipeline, (skipped_response_count, expected_response_count)) =
838 cache_manager.get_cached_pipeline(cmd);
839 if pipeline.is_empty() {
840 return cacheable_pipeline.resolve(cache_manager, Value::Array(Vec::new()));
841 }
842 let result = self
843 .pipeline
844 .send_recv(
845 pipeline.get_packed_pipeline(),
846 Some(PipelineResponseExpectation {
847 skipped_response_count,
848 expected_response_count,
849 is_transaction: cacheable_pipeline.transaction_mode,
850 seen_responses: 0,
851 }),
852 self.response_timeout,
853 false,
854 )
855 .await?;
856
857 return cacheable_pipeline.resolve(cache_manager, result);
858 }
859 let value = self
860 .pipeline
861 .send_recv(
862 cmd.get_packed_pipeline(),
863 Some(PipelineResponseExpectation {
864 skipped_response_count: offset,
865 expected_response_count: count,
866 is_transaction: cmd.is_transaction(),
867 seen_responses: 0,
868 }),
869 self.response_timeout,
870 false,
871 )
872 .await?;
873 match value {
874 Value::Array(values) => Ok(values),
875 _ => Ok(vec![value]),
876 }
877 }
878
879 #[cfg(feature = "cache-aio")]
881 #[cfg_attr(docsrs, doc(cfg(feature = "cache-aio")))]
882 pub fn get_cache_statistics(&self) -> Option<CacheStatistics> {
883 self.cache_manager.as_ref().map(|cm| cm.statistics())
884 }
885}
886
887impl ConnectionLike for MultiplexedConnection {
888 fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
889 (async move { self.send_packed_command(cmd).await }).boxed()
890 }
891
892 fn req_packed_commands<'a>(
893 &'a mut self,
894 pipeline: &'a crate::Pipeline,
895 offset: usize,
896 count: usize,
897 ) -> RedisFuture<'a, Vec<Value>> {
898 (async move { self.send_packed_commands(pipeline, offset, count).await }).boxed()
899 }
900
901 fn get_db(&self) -> i64 {
902 self.db
903 }
904}
905
906impl MultiplexedConnection {
907 pub async fn subscribe(&mut self, channel_name: impl ToRedisArgs) -> RedisResult<()> {
924 check_resp3!(self.protocol);
925 let mut cmd = cmd("SUBSCRIBE");
926 cmd.arg(channel_name);
927 cmd.exec_async(self).await?;
928 Ok(())
929 }
930
931 pub async fn unsubscribe(&mut self, channel_name: impl ToRedisArgs) -> RedisResult<()> {
946 check_resp3!(self.protocol);
947 let mut cmd = cmd("UNSUBSCRIBE");
948 cmd.arg(channel_name);
949 cmd.exec_async(self).await?;
950 Ok(())
951 }
952
953 pub async fn psubscribe(&mut self, channel_pattern: impl ToRedisArgs) -> RedisResult<()> {
972 check_resp3!(self.protocol);
973 let mut cmd = cmd("PSUBSCRIBE");
974 cmd.arg(channel_pattern);
975 cmd.exec_async(self).await?;
976 Ok(())
977 }
978
979 pub async fn punsubscribe(&mut self, channel_pattern: impl ToRedisArgs) -> RedisResult<()> {
983 check_resp3!(self.protocol);
984 let mut cmd = cmd("PUNSUBSCRIBE");
985 cmd.arg(channel_pattern);
986 cmd.exec_async(self).await?;
987 Ok(())
988 }
989}
990
991#[cfg(feature = "token-based-authentication")]
992impl MultiplexedConnection {
993 async fn re_authenticate_with_credentials(
998 &mut self,
999 credentials: &crate::auth::BasicAuth,
1000 ) -> RedisResult<()> {
1001 let mut auth_cmd =
1002 crate::connection::authenticate_cmd(Some(&credentials.username), &credentials.password);
1003 auth_cmd.skip_concurrency_limit = true;
1004 self.send_packed_command(&auth_cmd)
1005 .await?
1006 .extract_error()
1007 .map(|_| ())
1008 }
1009}
1010
1011#[cfg(test)]
1012mod tests {
1013 use super::*;
1014
1015 #[test]
1016 fn test_pipeline_resolve_buffer_size_default() {
1017 assert_eq!(Pipeline::resolve_buffer_size(None), 50);
1018 }
1019
1020 #[test]
1021 fn test_pipeline_resolve_buffer_size_custom() {
1022 assert_eq!(Pipeline::resolve_buffer_size(Some(100)), 100);
1023 }
1024
1025 #[test]
1026 fn test_permit_vec_capacity_is_bounded_by_limit() {
1027 assert_eq!(permit_vec_capacity(1_000_000, 4), 4);
1029 assert_eq!(permit_vec_capacity(10, 10), 10);
1030 assert_eq!(permit_vec_capacity(2, 4), 2);
1031 assert_eq!(permit_vec_capacity(0, 4), 1);
1032 assert_eq!(permit_vec_capacity(5, 0), 1);
1033 }
1034
1035 #[test]
1036 fn test_build_concurrency_limiter_rejects_zero() {
1037 let err = build_concurrency_limiter(Some(0)).unwrap_err();
1039 assert_eq!(err.kind(), ErrorKind::InvalidClientConfig);
1040 }
1041
1042 #[test]
1043 fn test_build_concurrency_limiter_allows_none_and_positive() {
1044 assert!(build_concurrency_limiter(None).unwrap().is_none());
1045 let limiter = build_concurrency_limiter(Some(4)).unwrap().unwrap();
1046 assert_eq!(limiter.limit, 4);
1047 }
1048
1049 fn mock_conn_info() -> RedisConnectionInfo {
1050 RedisConnectionInfo {
1051 skip_set_lib_name: true,
1052 ..Default::default()
1053 }
1054 }
1055
1056 async fn create_mock_connection(
1057 concurrency_limit: usize,
1058 ) -> (
1059 MultiplexedConnection,
1060 tokio::sync::mpsc::Receiver<()>,
1061 tokio::sync::mpsc::Sender<()>,
1062 ) {
1063 use futures_util::StreamExt;
1064 use tokio::io::AsyncWriteExt;
1065 use tokio_util::codec::FramedRead;
1066
1067 let (client_half, server_half) = tokio::io::duplex(4096);
1068 let (cmd_received_tx, cmd_received_rx) = tokio::sync::mpsc::channel::<()>(10);
1069 let (send_response_tx, mut send_response_rx) = tokio::sync::mpsc::channel::<()>(10);
1070
1071 let (server_read, mut server_write) = tokio::io::split(server_half);
1072
1073 tokio::spawn(async move {
1074 let mut reader = FramedRead::new(server_read, ValueCodec::default());
1075 while let Some(Ok(_)) = reader.next().await {
1076 let _ = cmd_received_tx.send(()).await;
1077 }
1078 });
1079
1080 tokio::spawn(async move {
1081 while send_response_rx.recv().await.is_some() {
1082 let _ = server_write.write_all(b"+OK\r\n").await;
1083 let _ = server_write.flush().await;
1084 }
1085 });
1086
1087 let config = AsyncConnectionConfig::new()
1088 .set_concurrency_limit(concurrency_limit)
1089 .set_response_timeout(None)
1090 .set_connection_timeout(None);
1091
1092 let (conn, driver) =
1093 MultiplexedConnection::new_with_config(&mock_conn_info(), client_half, config)
1094 .await
1095 .unwrap();
1096 tokio::spawn(driver);
1097
1098 (conn, cmd_received_rx, send_response_tx)
1099 }
1100
1101 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1102 async fn test_concurrency_limit_enforced() {
1103 let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(2).await;
1104
1105 let h1 = tokio::spawn({
1106 let mut c = conn.clone();
1107 async move { c.send_packed_command(&cmd("PING")).await }
1108 });
1109 let h2 = tokio::spawn({
1110 let mut c = conn.clone();
1111 async move { c.send_packed_command(&cmd("PING")).await }
1112 });
1113 let h3 = tokio::spawn({
1114 let mut c = conn.clone();
1115 async move { c.send_packed_command(&cmd("PING")).await }
1116 });
1117
1118 cmd_received_rx.recv().await.unwrap();
1119 cmd_received_rx.recv().await.unwrap();
1120
1121 let third = tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1122 assert!(
1123 third.is_err(),
1124 "3rd request should be blocked by concurrency limit"
1125 );
1126
1127 send_response_tx.send(()).await.unwrap();
1128
1129 cmd_received_rx.recv().await.unwrap();
1130
1131 send_response_tx.send(()).await.unwrap();
1132 send_response_tx.send(()).await.unwrap();
1133
1134 h1.await.unwrap().unwrap();
1135 h2.await.unwrap().unwrap();
1136 h3.await.unwrap().unwrap();
1137 }
1138
1139 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1140 async fn test_no_limit_bypasses_concurrency_limit() {
1141 let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(1).await;
1142
1143 let h1 = tokio::spawn({
1144 let mut c = conn.clone();
1145 async move { c.send_packed_command(&cmd("PING")).await }
1146 });
1147
1148 cmd_received_rx.recv().await.unwrap();
1149
1150 let h2 = tokio::spawn({
1151 let mut c = conn.clone();
1152 async move {
1153 let mut ping = cmd("PING");
1154 ping.skip_concurrency_limit = true;
1155 c.send_packed_command(&ping).await
1156 }
1157 });
1158
1159 let received =
1160 tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1161 assert!(
1162 received.is_ok(),
1163 "no_limit request should bypass concurrency limit"
1164 );
1165
1166 send_response_tx.send(()).await.unwrap();
1167 send_response_tx.send(()).await.unwrap();
1168
1169 h1.await.unwrap().unwrap();
1170 h2.await.unwrap().unwrap();
1171 }
1172
1173 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1174 async fn test_pipeline_acquires_multiple_permits() {
1175 let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(3).await;
1176
1177 let pipeline_handle = tokio::spawn({
1178 let mut c = conn.clone();
1179 async move {
1180 let mut pipe = crate::Pipeline::new();
1181 pipe.cmd("SET").arg("a").arg("1");
1182 pipe.cmd("SET").arg("b").arg("2");
1183 pipe.cmd("SET").arg("c").arg("3");
1184 c.send_packed_commands(&pipe, 0, 3).await
1185 }
1186 });
1187
1188 for _ in 0..3 {
1189 cmd_received_rx.recv().await.unwrap();
1190 }
1191
1192 let single_handle = tokio::spawn({
1193 let mut c = conn.clone();
1194 async move { c.send_packed_command(&cmd("PING")).await }
1195 });
1196
1197 let blocked =
1198 tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1199 assert!(
1200 blocked.is_err(),
1201 "single command should be blocked while pipeline holds all permits"
1202 );
1203
1204 for _ in 0..3 {
1205 send_response_tx.send(()).await.unwrap();
1206 }
1207
1208 cmd_received_rx.recv().await.unwrap();
1209 send_response_tx.send(()).await.unwrap();
1210
1211 pipeline_handle.await.unwrap().unwrap();
1212 single_handle.await.unwrap().unwrap();
1213 }
1214
1215 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1216 async fn test_pipeline_proceeds_with_partial_permits() {
1217 let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(2).await;
1218
1219 let single_handle = tokio::spawn({
1220 let mut c = conn.clone();
1221 async move { c.send_packed_command(&cmd("PING")).await }
1222 });
1223 cmd_received_rx.recv().await.unwrap();
1224
1225 let pipeline_handle = tokio::spawn({
1226 let mut c = conn.clone();
1227 async move {
1228 let mut pipe = crate::Pipeline::new();
1229 for i in 0..5 {
1230 pipe.cmd("SET").arg(format!("k{i}")).arg(i);
1231 }
1232 c.send_packed_commands(&pipe, 0, 5).await
1233 }
1234 });
1235
1236 let received =
1237 tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1238 assert!(
1239 received.is_ok(),
1240 "pipeline should proceed even with only partial permits"
1241 );
1242
1243 for _ in 1..5 {
1244 cmd_received_rx.recv().await.unwrap();
1245 }
1246
1247 for _ in 0..6 {
1248 send_response_tx.send(()).await.unwrap();
1249 }
1250
1251 single_handle.await.unwrap().unwrap();
1252 pipeline_handle.await.unwrap().unwrap();
1253 }
1254
1255 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1256 async fn test_permit_released_on_cancellation() {
1257 let (conn, mut cmd_received_rx, send_response_tx) = create_mock_connection(1).await;
1258
1259 let h1 = tokio::spawn({
1260 let mut c = conn.clone();
1261 async move { c.send_packed_command(&cmd("PING")).await }
1262 });
1263 cmd_received_rx.recv().await.unwrap();
1264
1265 let h2 = tokio::spawn({
1267 let mut c = conn.clone();
1268 async move { c.send_packed_command(&cmd("PING")).await }
1269 });
1270 tokio::time::sleep(Duration::from_millis(50)).await;
1271 h2.abort();
1272 let _ = h2.await;
1273
1274 send_response_tx.send(()).await.unwrap();
1276 h1.await.unwrap().unwrap();
1277
1278 let h3 = tokio::spawn({
1281 let mut c = conn.clone();
1282 async move { c.send_packed_command(&cmd("PING")).await }
1283 });
1284
1285 let received =
1286 tokio::time::timeout(Duration::from_millis(100), cmd_received_rx.recv()).await;
1287 assert!(
1288 received.is_ok(),
1289 "request after cancellation should acquire the permit"
1290 );
1291
1292 send_response_tx.send(()).await.unwrap();
1293 h3.await.unwrap().unwrap();
1294 }
1295
1296 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1335 async fn test_deadlock_when_writes_blocked_with_pending_response() {
1336 use futures_util::StreamExt;
1337 use tokio::io::AsyncWriteExt;
1338 use tokio_util::codec::FramedRead;
1339
1340 const BUFFER_SIZE: usize = 256;
1344 const PAYLOAD_SIZE: usize = 4096;
1345 const REQUEST_COUNT: usize = 3;
1346
1347 let (client_half, server_half) = tokio::io::duplex(BUFFER_SIZE);
1348 let (server_read, mut server_write) = tokio::io::split(server_half);
1349
1350 let mut response = Vec::with_capacity(PAYLOAD_SIZE + 16);
1354 response.extend_from_slice(format!("${PAYLOAD_SIZE}\r\n").as_bytes());
1355 response.extend(std::iter::repeat_n(b'V', PAYLOAD_SIZE));
1356 response.extend_from_slice(b"\r\n");
1357
1358 let server_task = tokio::spawn(async move {
1363 let mut reader = FramedRead::new(server_read, ValueCodec::default());
1364 loop {
1365 match reader.next().await {
1366 Some(Ok(_)) => {}
1367 _ => return,
1368 }
1369 if server_write.write_all(&response).await.is_err() {
1370 return;
1371 }
1372 }
1373 });
1374
1375 let config = AsyncConnectionConfig::new()
1376 .set_response_timeout(None)
1377 .set_connection_timeout(None);
1378 let (conn, driver) =
1379 MultiplexedConnection::new_with_config(&mock_conn_info(), client_half, config)
1380 .await
1381 .unwrap();
1382 let driver_handle = tokio::spawn(driver);
1383
1384 let mut handles = Vec::with_capacity(REQUEST_COUNT);
1389 for i in 0..REQUEST_COUNT {
1390 let mut c = conn.clone();
1391 handles.push(tokio::spawn(async move {
1392 let mut set = cmd("SET");
1393 set.arg(format!("k{i}")).arg(vec![b'X'; PAYLOAD_SIZE]);
1394 c.send_packed_command(&set).await
1395 }));
1396 }
1397
1398 let join_all = async move {
1399 let mut results = Vec::with_capacity(handles.len());
1400 for h in handles {
1401 results.push(h.await);
1402 }
1403 results
1404 };
1405
1406 let outcome = tokio::time::timeout(Duration::from_secs(5), join_all).await;
1407
1408 driver_handle.abort();
1410 server_task.abort();
1411
1412 let results = outcome.expect(
1413 "DEADLOCK reproduced: client driver parked in poll_flush with no \
1414 read waker registered. Server has buffered responses in the duplex \
1415 and stopped reading once its own write became Pending; the client \
1416 cannot send the rest of its requests because the server is no \
1417 longer draining the link. Both sides wedged.",
1418 );
1419 for (i, res) in results.into_iter().enumerate() {
1420 let join = res.unwrap_or_else(|e| panic!("SET task {i} panicked: {e}"));
1421 join.unwrap_or_else(|e| panic!("SET task {i} returned an error: {e}"));
1422 }
1423 }
1424
1425 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1426 async fn test_permit_released_on_response_timeout() {
1427 use futures_util::StreamExt;
1428 use tokio::io::AsyncWriteExt;
1429 use tokio_util::codec::FramedRead;
1430
1431 let (client_half, server_half) = tokio::io::duplex(4096);
1432 let (cmd_received_tx, mut cmd_received_rx) = tokio::sync::mpsc::channel::<()>(10);
1433
1434 let (server_read, mut server_write) = tokio::io::split(server_half);
1435
1436 tokio::spawn(async move {
1437 let mut reader = FramedRead::new(server_read, ValueCodec::default());
1438 while let Some(Ok(_)) = reader.next().await {
1439 let _ = cmd_received_tx.send(()).await;
1440 }
1441 });
1442
1443 tokio::spawn(async move {
1444 futures_util::future::pending::<()>().await;
1445 let _ = server_write.write_all(b"").await;
1446 });
1447
1448 let config = AsyncConnectionConfig::new()
1449 .set_concurrency_limit(1)
1450 .set_response_timeout(Some(Duration::from_millis(100)))
1451 .set_connection_timeout(None);
1452
1453 let (conn, driver) =
1454 MultiplexedConnection::new_with_config(&mock_conn_info(), client_half, config)
1455 .await
1456 .unwrap();
1457 tokio::spawn(driver);
1458
1459 let mut c1 = conn.clone();
1461 let err = c1.send_packed_command(&cmd("PING")).await.unwrap_err();
1462 assert!(err.is_io_error(), "expected IO error from timeout");
1463 cmd_received_rx.recv().await.unwrap();
1464
1465 let mut c2 = conn.clone();
1468 let err = c2.send_packed_command(&cmd("PING")).await.unwrap_err();
1469 assert!(err.is_io_error(), "expected IO error from timeout");
1470 cmd_received_rx.recv().await.unwrap();
1471 }
1472}