Coverage Report

Created: 2026-07-21 15:28

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/build/source/nativelink-store/src/redis_store.rs
Line
Count
Source
1
// Copyright 2024-2026 The NativeLink Authors. All rights reserved.
2
//
3
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//    See LICENSE file for details
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
15
use core::cmp;
16
use core::fmt::Debug;
17
use core::marker::PhantomData;
18
use core::ops::{Bound, RangeBounds};
19
use core::pin::Pin;
20
use core::str::FromStr;
21
use core::time::Duration;
22
use std::borrow::Cow;
23
use std::collections::HashSet;
24
use std::sync::{Arc, Weak};
25
use std::time::Instant;
26
27
use async_trait::async_trait;
28
use bytes::Bytes;
29
use const_format::formatcp;
30
use futures::stream::FuturesUnordered;
31
use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt, future};
32
use itertools::izip;
33
use nativelink_config::stores::{RedisMode, RedisSpec};
34
use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err};
35
use nativelink_metric::MetricsComponent;
36
use nativelink_redis_tester::SubscriptionManagerNotify;
37
use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
38
use nativelink_util::common::DigestInfo;
39
use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};
40
use nativelink_util::store_trait::{
41
    BoolValue, RemoveCallback, SchedulerCurrentVersionProvider, SchedulerIndexProvider,
42
    SchedulerStore, SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider,
43
    SchedulerSubscription, SchedulerSubscriptionManager, StoreDriver, StoreKey, UploadSizeInfo,
44
};
45
use nativelink_util::task::JoinHandleDropGuard;
46
use nativelink_util::{background_spawn, spawn};
47
use parking_lot::{Mutex, RwLock};
48
use patricia_tree::StringPatriciaMap;
49
use redis::aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig};
50
use redis::cluster::ClusterClient;
51
use redis::cluster_async::ClusterConnection;
52
use redis::sentinel::{SentinelClient, SentinelNodeConnectionInfo, SentinelServerType};
53
use redis::{
54
    AsyncCommands, AsyncIter, Client, IntoConnectionInfo, PushInfo, ScanOptions, Script, Value,
55
    pipe,
56
};
57
use serde::Deserialize;
58
use serde::de::IntoDeserializer;
59
use tokio::select;
60
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
61
use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore};
62
use tokio::time::{sleep, timeout};
63
use tokio_stream::wrappers::UnboundedReceiverStream;
64
use tracing::{debug, error, info, trace, warn};
65
use url::Url;
66
use uuid::Uuid;
67
68
use crate::cas_utils::is_zero_digest;
69
use crate::redis_utils::{
70
    FtAggregateCursor, FtAggregateOptions, FtCreateOptions, SearchSchema, ft_aggregate, ft_create,
71
};
72
73
/// The default size of the read chunk when reading data from Redis.
74
/// Note: If this changes it should be updated in the config documentation.
75
const DEFAULT_READ_CHUNK_SIZE: usize = 64 * 1024;
76
77
/// The default size of the connection pool if not specified.
78
/// Note: If this changes it should be updated in the config documentation.
79
const DEFAULT_CONNECTION_POOL_SIZE: usize = 3;
80
81
/// The default delay between retries if not specified.
82
/// Note: If this changes it should be updated in the config documentation.
83
const DEFAULT_RETRY_DELAY: f32 = 0.1;
84
85
/// Maximum attempts (re-resolving the master each time) for an operation to
86
/// ride out a transient Redis topology change — a failover, a dropped
87
/// connection, or brief replica lag. Used by `update`'s verify and by the read
88
/// paths so a transient error doesn't fail an otherwise healthy request.
89
const MAX_REDIS_RETRY_ATTEMPTS: u32 = 5;
90
91
/// Whether a Redis error is worth re-resolving the master and retrying: a
92
/// dropped/refused connection, an IO error, or a write that hit a freshly
93
/// demoted replica. Anything else (e.g. a real protocol/logic error) is
94
/// returned as-is.
95
13
fn is_retryable_redis_error(err: &redis::RedisError) -> bool {
96
13
    err.is_connection_dropped()
97
9
        || err.is_connection_refusal()
98
9
        || err.is_io_error()
99
8
        || err.kind() == redis::ErrorKind::Server(redis::ServerErrorKind::ReadOnly)
100
13
}
101
102
/// The default connection timeout in milliseconds if not specified.
103
/// Note: If this changes it should be updated in the config documentation.
104
const DEFAULT_CONNECTION_TIMEOUT_MS: u64 = 3000;
105
106
/// The default command timeout in milliseconds if not specified.
107
/// Note: If this changes it should be updated in the config documentation.
108
const DEFAULT_COMMAND_TIMEOUT_MS: u64 = 10_000;
109
110
/// The default `check_health` PING ceiling in milliseconds if not specified.
111
/// Note: If this changes it should be updated in the config documentation.
112
const DEFAULT_HEALTH_CHECK_TIMEOUT_MS: u64 = 4000;
113
114
/// The default maximum number of chunk uploads per update.
115
/// Note: If this changes it should be updated in the config documentation.
116
pub const DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE: usize = 10;
117
118
/// The default COUNT value passed when scanning keys in Redis.
119
/// Note: If this changes it should be updated in the config documentation.
120
const DEFAULT_SCAN_COUNT: usize = 10_000;
121
122
/// The default COUNT value passed when scanning search indexes
123
/// Note: If this changes it should be updated in the config documentation.
124
pub const DEFAULT_MAX_COUNT_PER_CURSOR: u64 = 1_500;
125
126
const DEFAULT_CLIENT_PERMITS: usize = 500;
127
128
/// A wrapper around Redis to allow it to be reconnected.
129
pub trait RedisManager<C>
130
where
131
    C: ConnectionLike + Clone,
132
{
133
    /// Get a connection manager and a unique identifier for this connection
134
    /// which may be used to issue a reconnect later.
135
    fn get_connection(&self) -> impl Future<Output = Result<(C, Uuid), Error>> + Send;
136
137
    /// Reconnect if the uuid matches the uuid returned from `get_connection()`.
138
    fn reconnect(&self, uuid: Uuid) -> impl Future<Output = Result<(C, Uuid), Error>> + Send;
139
140
    /// Get an invocation of the update version script for a given `key`.
141
    fn update_script(&self, key: &str) -> redis::ScriptInvocation<'_>;
142
143
    /// Configure the connection to have a psubscribe on it and perform the
144
    /// subscription on reconnect.
145
    fn psubscribe(&self, pattern: &str) -> impl Future<Output = Result<(), Error>> + Send;
146
}
147
148
#[derive(Debug)]
149
pub struct ClusterRedisManager<C>
150
where
151
    C: ConnectionLike + Clone,
152
{
153
    /// A constant Uuid, we never reconnect.
154
    uuid: Uuid,
155
156
    /// Redis script used to update a value in redis if the version matches.
157
    /// This is done by incrementing the version number and then setting the new
158
    /// data only if the version number matches the existing version number.
159
    update_if_version_matches_script: Script,
160
161
    /// The client pool connecting to the backing Redis instance(s).
162
    connection_manager: C,
163
}
164
165
impl<C> ClusterRedisManager<C>
166
where
167
    C: ConnectionLike + Clone,
168
{
169
26
    pub async fn new(mut connection_manager: C) -> Result<Self, Error> {
170
26
        let update_if_version_matches_script = Script::new(LUA_VERSION_SET_SCRIPT);
171
26
        update_if_version_matches_script
172
26
            .load_async(&mut connection_manager)
173
26
            .await
?0
;
174
26
        Ok(Self {
175
26
            uuid: Uuid::new_v4(),
176
26
            update_if_version_matches_script,
177
26
            connection_manager,
178
26
        })
179
26
    }
180
}
181
182
impl<C> RedisManager<C> for ClusterRedisManager<C>
183
where
184
    C: ConnectionLike + Clone + Send + Sync,
185
{
186
56
    fn get_connection(&self) -> impl Future<Output = Result<(C, Uuid), Error>> + Send {
187
56
        future::ready(Ok((self.connection_manager.clone(), self.uuid)))
188
56
    }
189
190
9
    fn reconnect(&self, _uuid: Uuid) -> impl Future<Output = Result<(C, Uuid), Error>> + Send {
191
9
        self.get_connection()
192
9
    }
193
194
0
    fn update_script(&self, key: &str) -> redis::ScriptInvocation<'_> {
195
0
        self.update_if_version_matches_script.key(key)
196
0
    }
197
198
2
    fn psubscribe(&self, _pattern: &str) -> impl Future<Output = Result<(), Error>> + Send {
199
        // This is a no-op for cluster connections.
200
2
        future::ready(Ok(()))
201
2
    }
202
}
203
204
type RedisConnectFuture<C> = dyn Future<Output = Result<C, Error>> + Send;
205
type RedisConnectFn<C> = dyn Fn() -> Pin<Box<RedisConnectFuture<C>>> + Send + Sync;
206
207
pub struct StandardRedisManager<C>
208
where
209
    C: ConnectionLike + Clone,
210
{
211
    /// Function used to re-connect to Redis.
212
    connect_func: Box<RedisConnectFn<C>>,
213
214
    /// Redis script used to update a value in redis if the version matches.
215
    /// This is done by incrementing the version number and then setting the new
216
    /// data only if the version number matches the existing version number.
217
    update_if_version_matches_script: Script,
218
219
    /// The client pool connecting to the backing Redis instance(s) and a Uuid
220
    /// for this connection in order to avoid multiple reconnection attempts.
221
    connection_manager: tokio::sync::RwLock<(C, Uuid)>,
222
223
    /// Serializes reconnect attempts so a Sentinel master failover triggers a
224
    /// single re-resolution instead of a thundering herd. Kept separate from
225
    /// `connection_manager` on purpose: the (potentially multi-second) connect
226
    /// runs while holding only this lock, so in-flight `get_connection` readers
227
    /// keep using the existing handle and are never frozen behind a slow
228
    /// reconnect. See [`Self::reconnect`].
229
    reconnect_lock: tokio::sync::Mutex<()>,
230
231
    /// A list of subscription that should be performed on reconnect.
232
    subscriptions: Mutex<HashSet<String>>,
233
}
234
235
impl<C> Debug for StandardRedisManager<C>
236
where
237
    C: ConnectionLike + Clone,
238
{
239
0
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
240
0
        f.debug_struct("StandardRedisManager")
241
0
            .field(
242
0
                "update_if_version_matches_script",
243
0
                &self.update_if_version_matches_script,
244
0
            )
245
0
            .field("subscriptions", &self.subscriptions)
246
0
            .finish()
247
0
    }
248
}
249
250
impl<C> StandardRedisManager<C>
251
where
252
    C: ConnectionLike + Clone + Send + Sync,
253
{
254
23
    async fn configure(&self, connection_manager: &mut C) -> Result<(), Error> {
255
23
        self.update_if_version_matches_script
256
23
            .load_async(connection_manager)
257
23
            .await
?0
;
258
23
        Ok(())
259
23
    }
260
261
23
    async fn new(connect_func: Box<RedisConnectFn<C>>) -> Result<Self, Error> {
262
23
        let 
connection_manager20
= connect_func().await
?3
;
263
20
        let update_if_version_matches_script = Script::new(LUA_VERSION_SET_SCRIPT);
264
20
        let connection = Self {
265
20
            connect_func,
266
20
            update_if_version_matches_script,
267
20
            connection_manager: tokio::sync::RwLock::new((connection_manager, Uuid::new_v4())),
268
20
            reconnect_lock: tokio::sync::Mutex::new(()),
269
20
            subscriptions: Mutex::new(HashSet::new()),
270
20
        };
271
        {
272
20
            let mut connection_manager = connection.connection_manager.write().await;
273
20
            connection.configure(&mut connection_manager.0).await
?0
;
274
        }
275
20
        Ok(connection)
276
23
    }
277
}
278
279
impl RedisManager<ConnectionManager> for StandardRedisManager<ConnectionManager> {
280
159
    async fn get_connection(&self) -> Result<(ConnectionManager, Uuid), Error> {
281
159
        Ok(self.connection_manager.read().await.clone())
282
159
    }
283
284
3
    async fn reconnect(&self, uuid: Uuid) -> Result<(ConnectionManager, Uuid), Error> {
285
        // Fast path: another caller already reconnected past this generation,
286
        // so the handle is fresh — hand it back without re-resolving.
287
        {
288
3
            let guard = self.connection_manager.read().await;
289
3
            if guard.1 != uuid {
290
0
                return Ok(guard.clone());
291
3
            }
292
        }
293
        // Serialize reconnect attempts on a dedicated lock so a failover storm
294
        // resolves the new master once. Critically this is NOT the
295
        // `connection_manager` lock, so the slow connect below does not block
296
        // the `get_connection` readers that every Redis op needs — they keep
297
        // using the old (now-failing) handle and fail fast on their own
298
        // command timeout instead of all freezing behind one reconnect.
299
3
        let _reconnect_guard = self.reconnect_lock.lock().await;
300
        // Re-check: a reconnect may have completed while we waited for the lock.
301
        {
302
3
            let guard = self.connection_manager.read().await;
303
3
            if guard.1 != uuid {
304
0
                return Ok(guard.clone());
305
3
            }
306
        }
307
3
        let mut connection_manager = (self.connect_func)().await
?0
;
308
3
        let new_uuid = Uuid::new_v4();
309
3
        self.configure(&mut connection_manager).await
?0
;
310
3
        let subscriptions = {
311
3
            let guard = self.subscriptions.lock();
312
3
            guard.iter().map(Clone::clone).collect::<Vec<_>>()
313
        };
314
3
        for 
subscription0
in subscriptions {
315
0
            connection_manager.psubscribe(&subscription).await?;
316
        }
317
        // Publish the new handle under a brief exclusive lock.
318
        {
319
3
            let mut guard = self.connection_manager.write().await;
320
3
            *guard = (connection_manager.clone(), new_uuid);
321
        }
322
3
        info!(old = %uuid, new = %new_uuid, "StandardRedisManager re-resolved the Redis master");
323
3
        Ok((connection_manager, new_uuid))
324
3
    }
325
326
19
    fn update_script(&self, key: &str) -> redis::ScriptInvocation<'_> {
327
19
        self.update_if_version_matches_script.key(key)
328
19
    }
329
330
8
    async fn psubscribe(&self, pattern: &str) -> Result<(), Error> {
331
8
        debug!(pattern, "new psubscribe");
332
8
        let mut connection = self.get_connection().await
?0
.0;
333
8
        let new_subscription = self.subscriptions.lock().insert(String::from(pattern));
334
8
        if new_subscription {
335
8
            let result = connection.psubscribe(pattern).await;
336
8
            if result.is_err() {
337
0
                self.subscriptions.lock().remove(pattern);
338
8
            }
339
8
            result
?0
;
340
0
        }
341
8
        debug!(pattern, new_subscription, "new psubscribe complete");
342
8
        Ok(())
343
8
    }
344
}
345
346
/// A [`StoreDriver`] implementation that uses Redis as a backing store.
347
#[derive(MetricsComponent)]
348
pub struct RedisStore<C, M>
349
where
350
    C: ConnectionLike + Clone,
351
    M: RedisManager<C>,
352
{
353
    /// The client pool connecting to the backing Redis instance(s).
354
    connection_manager: M,
355
356
    /// The underlying connection type in the connection manager.
357
    _connection_type: PhantomData<C>,
358
359
    /// A channel to publish updates to when a key is added, removed, or modified.
360
    #[metric(
361
        help = "The pubsub channel to publish updates to when a key is added, removed, or modified"
362
    )]
363
    pub_sub_channel: Option<String>,
364
365
    /// A function used to generate names for temporary keys.
366
    temp_name_generator_fn: fn() -> String,
367
368
    /// A common prefix to append to all keys before they are sent to Redis.
369
    ///
370
    /// See [`RedisStore::key_prefix`](`nativelink_config::stores::RedisStore::key_prefix`).
371
    #[metric(help = "Prefix to append to all keys before sending to Redis")]
372
    key_prefix: String,
373
374
    /// The amount of data to read from Redis at a time.
375
    #[metric(help = "The amount of data to read from Redis at a time")]
376
    read_chunk_size: usize,
377
378
    /// The maximum number of chunk uploads per update.
379
    /// This is used to limit the number of chunk uploads per update to prevent
380
    /// overloading when uploading large blocks of data
381
    #[metric(help = "The maximum number of chunk uploads per update")]
382
    max_chunk_uploads_per_update: usize,
383
384
    /// The COUNT value passed when scanning keys in Redis.
385
    /// This is used to hint the amount of work that should be done per response.
386
    #[metric(help = "The COUNT value passed when scanning keys in Redis")]
387
    scan_count: usize,
388
389
    /// The COUNT value used with search indexes
390
    #[metric(help = "The maximum number of results to return per cursor")]
391
    max_count_per_cursor: u64,
392
393
    /// A manager for subscriptions to keys in Redis.
394
    subscription_manager: Arc<RedisSubscriptionManager>,
395
396
    /// Permits to limit inflight Redis requests. Technically only
397
    /// limits the calls to `get_client()`, but the requests per client
398
    /// are small enough that it works well enough.
399
    client_permits: Arc<Semaphore>,
400
401
    /// Per-call ceiling for `check_health` PING.
402
    health_check_timeout: Duration,
403
404
    /// Have we done a subscribe for messages for `remove_callback` subscribes?
405
    has_remove_callback_subscribe: OnceCell<()>,
406
407
    remove_callbacks: Arc<async_lock::Mutex<Vec<RemoveCallback>>>,
408
}
409
410
impl<C, M> Debug for RedisStore<C, M>
411
where
412
    C: ConnectionLike + Clone,
413
    M: RedisManager<C>,
414
{
415
0
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
416
0
        f.debug_struct("RedisStore")
417
0
            .field("temp_name_generator_fn", &self.temp_name_generator_fn)
418
0
            .field("key_prefix", &self.key_prefix)
419
0
            .field("read_chunk_size", &self.read_chunk_size)
420
0
            .field(
421
0
                "max_chunk_uploads_per_update",
422
0
                &self.max_chunk_uploads_per_update,
423
0
            )
424
0
            .field("scan_count", &self.scan_count)
425
0
            .field("subscription_manager", &self.subscription_manager)
426
0
            .field("client_permits", &self.client_permits)
427
0
            .finish()
428
0
    }
429
}
430
431
struct ClientWithPermit<C: ConnectionLike> {
432
    connection_manager: C,
433
    uuid: Uuid,
434
435
    // here so it sticks around with the client and doesn't get dropped until that does
436
    #[allow(dead_code)]
437
    semaphore_permit: OwnedSemaphorePermit,
438
}
439
440
impl<C: ConnectionLike + Clone> ClientWithPermit<C> {
441
3
    async fn reconnect<M: RedisManager<C> + Sync>(&mut self, manager: &M) -> Result<(), Error> {
442
3
        (self.connection_manager, self.uuid) = manager.reconnect(self.uuid).await
?0
;
443
3
        Ok(())
444
3
    }
445
}
446
447
impl<C: ConnectionLike> Drop for ClientWithPermit<C> {
448
157
    fn drop(&mut self) {
449
157
        trace!(
450
157
            remaining = self.semaphore_permit.semaphore().available_permits(),
451
            "Dropping a client permit"
452
        );
453
157
    }
454
}
455
456
/// Decode a `StoreKey` coming from Redis
457
33
pub fn decode_key<'a>(
458
33
    key_prefix: &String,
459
33
    encoded_key: Cow<'a, str>,
460
33
) -> Result<StoreKey<'a>, Error> {
461
33
    let key_no_prefix = if key_prefix.is_empty() {
462
29
        encoded_key
463
    } else {
464
4
        match encoded_key.strip_prefix(key_prefix) {
465
4
            Some(r) => Cow::from(r.to_string()),
466
            None => {
467
0
                return Err(make_err!(
468
0
                    Code::InvalidArgument,
469
0
                    "Redis key ({}) is missing prefix ({})",
470
0
                    encoded_key,
471
0
                    key_prefix
472
0
                ));
473
            }
474
        }
475
    };
476
33
    let maybe_digest_info: Result<_, serde::de::value::Error> =
477
33
        DigestInfo::deserialize(key_no_prefix.clone().into_deserializer());
478
33
    if let Ok(
digest_info6
) = maybe_digest_info {
479
6
        Ok(StoreKey::Digest(digest_info))
480
    } else {
481
27
        Ok(StoreKey::Str(key_no_prefix))
482
    }
483
33
}
484
485
impl<C, M> RedisStore<C, M>
486
where
487
    C: ConnectionLike + Clone + Sync,
488
    M: RedisManager<C> + Sync,
489
{
490
    /// Used for testing when determinism is required.
491
    #[expect(clippy::too_many_arguments)]
492
46
    pub async fn new_from_builder_and_parts(
493
46
        pub_sub_channel: Option<String>,
494
46
        temp_name_generator_fn: fn() -> String,
495
46
        key_prefix: String,
496
46
        read_chunk_size: usize,
497
46
        max_chunk_uploads_per_update: usize,
498
46
        scan_count: usize,
499
46
        max_client_permits: usize,
500
46
        max_count_per_cursor: u64,
501
46
        health_check_timeout: Duration,
502
46
        subscriber_channel: UnboundedReceiver<PushInfo>,
503
46
        connection_manager: M,
504
46
    ) -> Result<Self, Error> {
505
46
        info!("Redis index fingerprint: {FINGERPRINT_CREATE_INDEX_HEX}");
506
46
        let remove_callbacks = Arc::new(async_lock::Mutex::new(Vec::new()));
507
46
        let subscription_manager = Arc::new(RedisSubscriptionManager::new(
508
46
            subscriber_channel,
509
46
            remove_callbacks.clone(),
510
46
            key_prefix.clone(),
511
        ));
512
46
        if let Some(
channel4
) = &pub_sub_channel {
513
4
            connection_manager.psubscribe(channel).await
?0
;
514
42
        }
515
516
46
        Ok(Self {
517
46
            connection_manager,
518
46
            _connection_type: PhantomData,
519
46
            pub_sub_channel,
520
46
            temp_name_generator_fn,
521
46
            key_prefix,
522
46
            read_chunk_size,
523
46
            max_chunk_uploads_per_update,
524
46
            scan_count,
525
46
            subscription_manager,
526
46
            client_permits: Arc::new(Semaphore::new(max_client_permits)),
527
46
            max_count_per_cursor,
528
46
            health_check_timeout,
529
46
            has_remove_callback_subscribe: OnceCell::const_new(),
530
46
            remove_callbacks,
531
46
        })
532
46
    }
533
534
157
    async fn get_client(&self) -> Result<ClientWithPermit<C>, Error> {
535
157
        let local_client_permits = self.client_permits.clone();
536
157
        let remaining = local_client_permits.available_permits();
537
157
        let semaphore_permit = local_client_permits.acquire_owned().await
?0
;
538
157
        trace!(remaining, "Got a client permit");
539
157
        let (connection_manager, uuid) = self.connection_manager.get_connection().await
?0
;
540
157
        Ok(ClientWithPermit {
541
157
            connection_manager,
542
157
            uuid,
543
157
            semaphore_permit,
544
157
        })
545
157
    }
546
547
    /// Encode a `StoreKey` so it can be sent to Redis.
548
150
    pub fn encode_key<'a>(&self, key: &'a StoreKey<'a>) -> Cow<'a, str> {
549
150
        let key_body = key.as_str();
550
150
        if self.key_prefix.is_empty() {
551
142
            key_body
552
        } else {
553
            // This is in the hot path for all redis operations, so we try to reuse the allocation
554
            // from `key.as_str()` if possible.
555
8
            match key_body {
556
7
                Cow::Owned(mut encoded_key) => {
557
7
                    encoded_key.insert_str(0, &self.key_prefix);
558
7
                    Cow::Owned(encoded_key)
559
                }
560
1
                Cow::Borrowed(body) => {
561
1
                    let mut encoded_key = String::with_capacity(self.key_prefix.len() + body.len());
562
1
                    encoded_key.push_str(&self.key_prefix);
563
1
                    encoded_key.push_str(body);
564
1
                    Cow::Owned(encoded_key)
565
                }
566
            }
567
        }
568
150
    }
569
570
23
    fn set_spec_defaults(spec: &mut RedisSpec) -> Result<(), Error> {
571
23
        if spec.addresses.is_empty() {
572
0
            return Err(make_err!(
573
0
                Code::InvalidArgument,
574
0
                "No addresses were specified in redis store configuration."
575
0
            ));
576
23
        }
577
578
23
        if spec.broadcast_channel_capacity != 0 {
579
1
            warn!("broadcast_channel_capacity in Redis spec is deprecated and ignored");
580
22
        }
581
23
        if spec.response_timeout_s != 0 {
582
0
            warn!(
583
                "response_timeout_s in Redis spec is deprecated and ignored, use command_timeout_ms"
584
            );
585
23
        }
586
23
        if spec.connection_timeout_s != 0 {
587
0
            if spec.connection_timeout_ms != 0 {
588
0
                return Err(make_err!(
589
0
                    Code::InvalidArgument,
590
0
                    "Both connection_timeout_s and connection_timeout_ms were set, can only have one!"
591
0
                ));
592
0
            }
593
0
            warn!("connection_timeout_s in Redis spec is deprecated, use connection_timeout_ms");
594
0
            spec.connection_timeout_ms = spec.connection_timeout_s * 1000;
595
23
        }
596
23
        if spec.connection_timeout_ms == 0 {
597
18
            spec.connection_timeout_ms = DEFAULT_CONNECTION_TIMEOUT_MS;
598
18
        
}5
599
23
        if spec.command_timeout_ms == 0 {
600
21
            spec.command_timeout_ms = DEFAULT_COMMAND_TIMEOUT_MS;
601
21
        
}2
602
23
        if spec.health_check_timeout_ms == 0 {
603
23
            spec.health_check_timeout_ms = DEFAULT_HEALTH_CHECK_TIMEOUT_MS;
604
23
        
}0
605
23
        if spec.connection_pool_size == 0 {
606
23
            spec.connection_pool_size = DEFAULT_CONNECTION_POOL_SIZE;
607
23
        
}0
608
23
        if spec.read_chunk_size == 0 {
609
23
            spec.read_chunk_size = DEFAULT_READ_CHUNK_SIZE;
610
23
        
}0
611
23
        if spec.max_count_per_cursor == 0 {
612
23
            spec.max_count_per_cursor = DEFAULT_MAX_COUNT_PER_CURSOR;
613
23
        
}0
614
23
        if spec.max_chunk_uploads_per_update == 0 {
615
23
            spec.max_chunk_uploads_per_update = DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE;
616
23
        
}0
617
23
        if spec.scan_count == 0 {
618
23
            spec.scan_count = DEFAULT_SCAN_COUNT;
619
23
        
}0
620
23
        if spec.max_client_permits == 0 {
621
23
            spec.max_client_permits = DEFAULT_CLIENT_PERMITS;
622
23
        
}0
623
23
        if spec.retry.delay == 0.0 {
624
23
            spec.retry.delay = DEFAULT_RETRY_DELAY;
625
23
        
}0
626
23
        if spec.retry.max_retries == 0 {
627
23
            spec.retry.max_retries = 1;
628
23
        
}0
629
23
        trace!(?spec, "redis spec is after setting defaults");
630
23
        Ok(())
631
23
    }
632
633
    // Only used by tests, because we need to make a real redis connection, then fix this to get fixed values
634
10
    pub fn replace_temp_name_generator(&mut self, replacement: fn() -> String) {
635
10
        self.temp_name_generator_fn = replacement;
636
10
    }
637
}
638
639
impl RedisStore<ClusterConnection, ClusterRedisManager<ClusterConnection>> {
640
0
    pub async fn new_cluster(mut spec: RedisSpec) -> Result<Arc<Self>, Error> {
641
0
        if spec.mode != RedisMode::Cluster {
642
0
            return Err(Error::new(
643
0
                Code::InvalidArgument,
644
0
                "new_cluster only works for Cluster mode".to_string(),
645
0
            ));
646
0
        }
647
0
        Self::set_spec_defaults(&mut spec)?;
648
649
0
        let parsed_addrs: Vec<_> = spec
650
0
            .addresses
651
0
            .iter_mut()
652
0
            .map(|addr| {
653
0
                addr.clone().into_connection_info().map(|connection_info| {
654
0
                    let redis_settings = connection_info
655
0
                        .redis_settings()
656
0
                        .clone()
657
                        // We need RESP3 here because the cluster mode doesn't support RESP2 pubsub
658
                        // See also https://docs.rs/redis/latest/redis/cluster_async/index.html#pubsub
659
0
                        .set_protocol(redis::ProtocolVersion::RESP3);
660
0
                    connection_info.set_redis_settings(redis_settings)
661
0
                })
662
0
            })
663
0
            .collect::<Result<Vec<_>, _>>()?;
664
665
0
        let connection_timeout = Duration::from_millis(spec.connection_timeout_ms);
666
0
        let command_timeout = Duration::from_millis(spec.command_timeout_ms);
667
0
        let (tx, subscriber_channel) = unbounded_channel();
668
669
0
        let builder = ClusterClient::builder(parsed_addrs)
670
0
            .connection_timeout(connection_timeout)
671
0
            .response_timeout(command_timeout)
672
0
            .push_sender(tx)
673
0
            .retries(u32::try_from(spec.retry.max_retries)?);
674
675
0
        let client = builder.build()?;
676
677
0
        Self::new_from_builder_and_parts(
678
0
            spec.experimental_pub_sub_channel,
679
0
            || Uuid::new_v4().to_string(),
680
0
            spec.key_prefix.clone(),
681
0
            spec.read_chunk_size,
682
0
            spec.max_chunk_uploads_per_update,
683
0
            spec.scan_count,
684
0
            spec.max_client_permits,
685
0
            spec.max_count_per_cursor,
686
0
            Duration::from_millis(spec.health_check_timeout_ms),
687
0
            subscriber_channel,
688
0
            ClusterRedisManager::new(client.get_async_connection().await?).await?,
689
        )
690
0
        .await
691
0
        .map(Arc::new)
692
0
    }
693
}
694
695
impl RedisStore<ConnectionManager, StandardRedisManager<ConnectionManager>> {
696
26
    async fn connect(
697
26
        spec: RedisSpec,
698
26
        tx: UnboundedSender<PushInfo>,
699
26
    ) -> Result<ConnectionManager, Error> {
700
26
        let connection_timeout = Duration::from_millis(spec.connection_timeout_ms);
701
26
        let command_timeout = Duration::from_millis(spec.command_timeout_ms);
702
703
26
        let addr = &spec.addresses[0];
704
26
        let local_addr = addr.clone();
705
26
        let mut parsed_addr = local_addr
706
26
            .replace("redis+sentinel://", "redis://")
707
26
            .into_connection_info()
?0
;
708
709
26
        let redis_settings = parsed_addr
710
26
            .redis_settings()
711
26
            .clone()
712
            // We need RESP3 here because we want to do set_push_sender
713
26
            .set_protocol(redis::ProtocolVersion::RESP3);
714
26
        parsed_addr = parsed_addr.set_redis_settings(redis_settings);
715
26
        debug!(?parsed_addr, "Parsed redis addr");
716
717
26
        let 
client25
= timeout(
718
26
            connection_timeout,
719
26
            spawn!("connect", async move {
720
26
                match spec.mode {
721
17
                    RedisMode::Standard => Client::open(parsed_addr).map_err(Into::<Error>::into),
722
                    RedisMode::Cluster => {
723
0
                        return Err(Error::new(
724
0
                            Code::Internal,
725
0
                            "Use RedisStore::new_cluster for cluster connections".to_owned(),
726
0
                        ));
727
                    }
728
9
                    RedisMode::Sentinel => async {
729
9
                        let url_parsing = Url::parse(&local_addr)
?0
;
730
9
                        let master_name = url_parsing
731
9
                            .query_pairs()
732
9
                            .find(|(key, _)| 
key1
==
"sentinelServiceName"1
)
733
9
                            .map_or_else(|| 
"master"8
.
into8
(), |(_, value)|
value1
.
to_string1
());
734
735
9
                        let redis_connection_info = parsed_addr.redis_settings().clone();
736
9
                        let sentinel_connection_info = SentinelNodeConnectionInfo::default()
737
9
                            .set_redis_connection_info(redis_connection_info);
738
739
                        // We fish this out because sentinels don't support db, we need to set it
740
                        // on the client only. See also https://github.com/redis-rs/redis-rs/issues/1950
741
9
                        let original_db = parsed_addr.redis_settings().db();
742
9
                        if original_db != 0 {
743
1
                            // sentinel_connection_info has the actual DB set
744
1
                            let revised_settings = parsed_addr.redis_settings().clone().set_db(0);
745
1
                            parsed_addr = parsed_addr.set_redis_settings(revised_settings);
746
8
                        }
747
748
9
                        SentinelClient::build(
749
9
                            vec![parsed_addr],
750
9
                            master_name,
751
9
                            Some(sentinel_connection_info),
752
9
                            SentinelServerType::Master,
753
                        )
754
9
                        .map_err(Into::<Error>::into)
755
9
                    }
756
18
                    .
and_then9
(|mut s| async move
{9
Ok(
s9
.async_get_client().await) })
757
9
                    .await
?0
758
9
                    .map_err(Into::<Error>::into),
759
                }
760
26
                .err_tip_with_code(|_e| 
{1
761
1
                    (
762
1
                        Code::InvalidArgument,
763
1
                        format!("While connecting to redis with url: {local_addr}"),
764
1
                    )
765
1
                })
766
26
            }),
767
        )
768
26
        .await
769
26
        .err_tip(|| 
format!0
("Timeout while connecting to redis with url: {addr}"))
?0
?0
?1
;
770
771
25
        let connection_manager_config = {
772
25
            ConnectionManagerConfig::new()
773
25
                .set_number_of_retries(spec.retry.max_retries)
774
25
                .set_connection_timeout(Some(connection_timeout))
775
25
                .set_response_timeout(Some(command_timeout))
776
25
                .set_push_sender(tx)
777
        };
778
779
23
        let mut connection_manager =
780
25
            ConnectionManager::new_with_config(client, connection_manager_config)
781
25
                .await
782
25
                .err_tip(|| 
format!2
("While connecting to redis with url: {addr}"))
?2
;
783
784
23
        if let Some(
pub_sub_channel4
) = spec.experimental_pub_sub_channel {
785
4
            connection_manager.psubscribe(pub_sub_channel).await
?0
;
786
19
        }
787
788
23
        Ok(connection_manager)
789
26
    }
790
791
    /// Create a new `RedisStore` from the given configuration.
792
23
    pub async fn new_standard(mut spec: RedisSpec) -> Result<Arc<Self>, Error> {
793
23
        Self::set_spec_defaults(&mut spec)
?0
;
794
795
23
        if spec.addresses.len() != 1 {
796
0
            return Err(make_err!(
797
0
                Code::Unimplemented,
798
0
                "Connecting directly to multiple redis nodes in a cluster is currently unsupported. Please specify a single URL to a single node, and nativelink will use cluster discover to find the other nodes."
799
0
            ));
800
23
        }
801
802
23
        let (tx, subscriber_channel) = unbounded_channel();
803
804
20
        Self::new_from_builder_and_parts(
805
23
            spec.experimental_pub_sub_channel.clone(),
806
1
            || Uuid::new_v4().to_string(),
807
23
            spec.key_prefix.clone(),
808
23
            spec.read_chunk_size,
809
23
            spec.max_chunk_uploads_per_update,
810
23
            spec.scan_count,
811
23
            spec.max_client_permits,
812
23
            spec.max_count_per_cursor,
813
23
            Duration::from_millis(spec.health_check_timeout_ms),
814
23
            subscriber_channel,
815
26
            
StandardRedisManager::new23
(
Box::new23
(move || {
816
26
                Box::pin(Self::connect(spec.clone(), tx.clone()))
817
26
            }))
818
23
            .await
?3
,
819
        )
820
20
        .await
821
20
        .map(Arc::new)
822
23
    }
823
}
824
825
#[async_trait]
826
impl<C, M> StoreDriver for RedisStore<C, M>
827
where
828
    C: ConnectionLike + Clone + Send + Sync + Unpin + 'static,
829
    M: RedisManager<C> + Unpin + Send + Sync + 'static,
830
{
831
0
    async fn post_init(self: Arc<Self>) -> Result<(), Error> {
832
        Ok(())
833
0
    }
834
835
    async fn has_with_results(
836
        self: Pin<&Self>,
837
        keys: &[StoreKey<'_>],
838
        results: &mut [Option<u64>],
839
9
    ) -> Result<(), Error> {
840
        // TODO(palfrey) We could use pipeline here, but it makes retry more
841
        // difficult and it doesn't work very well in cluster mode.
842
        // If we wanted to optimize this with pipeline be careful to
843
        // implement retry and to support cluster mode.
844
845
        izip!(keys.iter(), results.iter_mut(),)
846
9
            .map(|(key, result)| async move {
847
                // We need to do a special pass to ensure our zero key exist.
848
9
                if is_zero_digest(key.borrow()) {
849
2
                    *result = Some(0);
850
2
                    return Ok::<_, Error>(());
851
7
                }
852
7
                let encoded_key = self.encode_key(key);
853
854
7
                let mut client = self.get_client().await
?0
;
855
856
                // Redis returns 0 when the key doesn't exist
857
                // AND when the key exists with value of length 0.
858
                // Therefore, we need to check both length and existence
859
                // and do it in a pipeline for efficiency
860
                // Re-resolve the master and retry on a transient failover so a
861
                // topology change doesn't fail the existence check.
862
7
                let (blob_len, exists) = {
863
7
                    let mut attempt: u32 = 0;
864
                    loop {
865
8
                        attempt += 1;
866
8
                        match pipe()
867
8
                            .strlen(encoded_key.as_ref())
868
8
                            .exists(encoded_key.as_ref())
869
8
                            .query_async::<(u64, bool)>(&mut client.connection_manager)
870
8
                            .await
871
                        {
872
7
                            Ok(v) => break v,
873
1
                            Err(err)
874
1
                                if attempt < MAX_REDIS_RETRY_ATTEMPTS
875
1
                                    && is_retryable_redis_error(&err) =>
876
                            {
877
1
                                client.reconnect(&self.connection_manager).await
?0
;
878
1
                                sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;
879
                            }
880
0
                            Err(err) => {
881
0
                                return Err(
882
0
                                    Error::from(err).append("In RedisStore::has_with_results::all")
883
0
                                );
884
                            }
885
                        }
886
                    }
887
                };
888
889
7
                *result = if exists { Some(blob_len) } else { 
None0
};
890
891
7
                Ok::<_, Error>(())
892
18
            })
893
            .collect::<FuturesUnordered<_>>()
894
            .try_collect()
895
            .await
896
9
    }
897
898
    async fn list(
899
        self: Pin<&Self>,
900
        range: (Bound<StoreKey<'_>>, Bound<StoreKey<'_>>),
901
        handler: &mut (dyn for<'a> FnMut(&'a StoreKey) -> bool + Send + Sync + '_),
902
8
    ) -> Result<u64, Error> {
903
        let range = (
904
            range.0.map(StoreKey::into_owned),
905
            range.1.map(StoreKey::into_owned),
906
        );
907
        let pattern = match range.0 {
908
            Bound::Included(ref start) | Bound::Excluded(ref start) => match range.1 {
909
                Bound::Included(ref end) | Bound::Excluded(ref end) => {
910
                    let start = start.as_str();
911
                    let end = end.as_str();
912
                    let max_length = start.len().min(end.len());
913
                    let length = start
914
                        .chars()
915
                        .zip(end.chars())
916
20
                        .position(|(a, b)| a != b)
917
                        .unwrap_or(max_length);
918
                    format!("{}{}*", self.key_prefix, &start[..length])
919
                }
920
                Bound::Unbounded => format!("{}*", self.key_prefix),
921
            },
922
            Bound::Unbounded => format!("{}*", self.key_prefix),
923
        };
924
        let mut client = self.get_client().await?;
925
        trace!(%pattern, count=self.scan_count, "Running SCAN");
926
        // Restart the scan on a transient failover. Redis SCAN may re-emit keys
927
        // even without failures, so callers already tolerate duplicates;
928
        // re-resolving the master and rescanning is within that contract and
929
        // avoids failing on a topology change mid-iteration.
930
        let mut attempt: u32 = 0;
931
        loop {
932
            attempt += 1;
933
            let opts = ScanOptions::default()
934
                .with_pattern(pattern.clone())
935
                .with_count(self.scan_count);
936
            // Scan via a cloned connection handle so the iterator doesn't hold a
937
            // borrow on `client` across a reconnect.
938
            let mut conn = client.connection_manager.clone();
939
            let mut scan_stream: AsyncIter<Value> = match conn.scan_options(opts).await {
940
                Ok(s) => s,
941
                Err(err)
942
                    if attempt < MAX_REDIS_RETRY_ATTEMPTS && is_retryable_redis_error(&err) =>
943
                {
944
                    client.reconnect(&self.connection_manager).await?;
945
                    sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;
946
                    continue;
947
                }
948
                Err(err) => return Err(Error::from(err).append("During scan_options")),
949
            };
950
            let mut iterations = 0;
951
            let mut errors = vec![];
952
            let mut transient_err = false;
953
            while let Some(key) = scan_stream.next_item().await {
954
                match key {
955
                    Ok(Value::BulkString(raw_key)) => {
956
                        let Ok(str_key) = str::from_utf8(&raw_key) else {
957
                            error!(?raw_key, "Non-utf8 key");
958
                            errors.push(format!("Non-utf8 key {raw_key:?}"));
959
                            continue;
960
                        };
961
                        match decode_key(&self.key_prefix, Cow::from(str_key)) {
962
                            Ok(key) => {
963
                                if range.contains(&key) {
964
                                    iterations += 1;
965
                                    if !handler(&key) {
966
                                        error!("Issue in handler");
967
                                        errors.push("Issue in handler".to_string());
968
                                    }
969
                                } else {
970
                                    trace!(%key, ?range, "Key not in range");
971
                                }
972
                            }
973
                            Err(e) => {
974
                                errors.push(e.to_string());
975
                            }
976
                        }
977
                    }
978
                    Err(err)
979
                        if attempt < MAX_REDIS_RETRY_ATTEMPTS && is_retryable_redis_error(&err) =>
980
                    {
981
                        // Connection dropped mid-scan; restart from scratch.
982
                        transient_err = true;
983
                        break;
984
                    }
985
                    other => {
986
                        error!(?other, "Non-string in key");
987
                        errors.push("Non-string in key".to_string());
988
                    }
989
                }
990
            }
991
            if transient_err {
992
                drop(scan_stream);
993
                client.reconnect(&self.connection_manager).await?;
994
                sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;
995
                continue;
996
            }
997
            return if errors.is_empty() {
998
                Ok(iterations)
999
            } else {
1000
                error!(?errors, "Errors in scan stream");
1001
                Err(Error::new(Code::Internal, format!("Errors: {errors:?}")))
1002
            };
1003
        }
1004
8
    }
1005
1006
    async fn update(
1007
        self: Pin<&Self>,
1008
        key: StoreKey<'_>,
1009
        mut reader: DropCloserReadHalf,
1010
        _upload_size: UploadSizeInfo,
1011
12
    ) -> Result<u64, Error> {
1012
        let final_key = self.encode_key(&key);
1013
1014
        // While the name generation function can be supplied by the user, we need to have the curly
1015
        // braces in place in order to manage redis' hashing behavior and make sure that the temporary
1016
        // key name and the final key name are directed to the same cluster node. See
1017
        // https://redis.io/blog/redis-clustering-best-practices-with-keys/
1018
        //
1019
        // The TL;DR is that if we're in cluster mode and the names hash differently, we can't use request
1020
        // pipelining. By using these braces, we tell redis to only hash the part of the temporary key that's
1021
        // identical to the final key -- so they will always hash to the same node.
1022
        let temp_key = format!(
1023
            "temp-{}-{{{}}}",
1024
            (self.temp_name_generator_fn)(),
1025
            &final_key
1026
        );
1027
1028
        if is_zero_digest(key.borrow()) {
1029
            let chunk = reader
1030
                .peek()
1031
                .await
1032
                .err_tip(|| "Failed to peek in RedisStore::update")?;
1033
            if chunk.is_empty() {
1034
                reader
1035
                    .drain()
1036
                    .await
1037
                    .err_tip(|| "Failed to drain in RedisStore::update")?;
1038
                // Zero-digest keys are special -- we don't need to do anything with it.
1039
                return Ok(0);
1040
            }
1041
        }
1042
1043
        let mut client = self.get_client().await?;
1044
1045
        let mut read_stream = reader
1046
11
            .scan(0u32, |bytes_read, chunk_res| {
1047
11
                future::ready(Some(
1048
11
                    chunk_res
1049
11
                        .err_tip(|| "Failed to read chunk in update in redis store")
1050
11
                        .and_then(|chunk| 
{10
1051
10
                            let offset = isize::try_from(*bytes_read).err_tip(|| "Could not convert offset to isize in RedisStore::update")
?0
;
1052
10
                            let chunk_len = u32::try_from(chunk.len()).err_tip(
1053
                                || "Could not convert chunk length to u32 in RedisStore::update",
1054
0
                            )?;
1055
10
                            let new_bytes_read = bytes_read
1056
10
                                .checked_add(chunk_len)
1057
10
                                .err_tip(|| "Overflow protection in RedisStore::update")
?0
;
1058
10
                            *bytes_read = new_bytes_read;
1059
10
                            Ok::<_, Error>((offset, *bytes_read, chunk))
1060
10
                        }),
1061
                ))
1062
11
            })
1063
11
            .map(|res| {
1064
11
                let (
offset10
,
end_pos10
,
chunk10
) = res
?1
;
1065
10
                let temp_key_ref = &temp_key;
1066
10
                Ok(async move {
1067
10
                    let (mut connection_manager, connect_id) = self.connection_manager.get_connection().await
?0
;
1068
10
                    match connection_manager
1069
10
                        .setrange::<_, _, usize>(temp_key_ref, offset, chunk.to_vec())
1070
10
                        .await {
1071
7
                        Ok(_) => {},
1072
3
                        Err(err)
1073
3
                            if is_retryable_redis_error(&err) =>
1074
                        {
1075
3
                            let (mut connection_manager, _connect_id) = self.connection_manager.reconnect(connect_id).await
?0
;
1076
3
                            connection_manager
1077
3
                                .setrange::<_, _, usize>(temp_key_ref, offset, chunk.to_vec())
1078
3
                                .await
1079
3
                                .err_tip(
1080
1
                                    || format!("(after reconnect) while appending to temp key ({temp_key_ref}) in RedisStore::update. offset = {offset}. end_pos = {end_pos}"),
1081
1
                                )?;
1082
                        }
1083
0
                        Err(err) => {
1084
0
                            let mut error: Error = err.into();
1085
0
                            error
1086
0
                                .messages
1087
0
                                .push(format!("While appending to temp key ({temp_key_ref}) in RedisStore::update. offset = {offset}. end_pos = {end_pos}"));
1088
0
                            return Err(error);
1089
                        }
1090
                    }
1091
9
                    Ok::<u32, Error>(end_pos)
1092
10
                })
1093
11
            })
1094
            .try_buffer_unordered(self.max_chunk_uploads_per_update);
1095
1096
        let mut total_len: u32 = 0;
1097
        while let Some(last_pos) = read_stream.try_next().await? {
1098
            if last_pos > total_len {
1099
                total_len = last_pos;
1100
            }
1101
        }
1102
1103
        let expected_len = usize::try_from(total_len).unwrap_or(usize::MAX);
1104
1105
        // The chunk writes above reconnect on any transient failover error, so on a mid-write Redis
1106
        // failover the data lands on the *current* master. The length check and
1107
        // rename below must run against that same master: the connection
1108
        // captured before the writes may now point at a demoted replica, where
1109
        // strlen reads 0 and would fail an otherwise healthy write. Re-resolve
1110
        // the master and retry so a transient topology change (failover or brief
1111
        // replica lag) doesn't drop the value.
1112
        let mut attempt: u32 = 0;
1113
        let blob_len = loop {
1114
            attempt += 1;
1115
            let blob_len: usize = client
1116
                .connection_manager
1117
                .strlen(&temp_key)
1118
                .await
1119
0
                .err_tip(|| format!("In RedisStore::update strlen check for {temp_key}"))?;
1120
            // Safety check: reject if a retried append double-wrote the data.
1121
            if blob_len > expected_len {
1122
                return Err(make_input_err!(
1123
                    "Data length mismatch in RedisStore::update for {}({}) - expected {} bytes, got {} bytes",
1124
                    key.borrow().as_str(),
1125
                    temp_key,
1126
                    total_len,
1127
                    blob_len,
1128
                ));
1129
            }
1130
            if blob_len == expected_len {
1131
                break blob_len;
1132
            }
1133
            // blob_len < expected (typically 0): the temp key isn't visible on
1134
            // this connection yet — re-resolve the master and retry.
1135
            if attempt >= MAX_REDIS_RETRY_ATTEMPTS {
1136
                return Err(make_input_err!(
1137
                    "Data length mismatch in RedisStore::update for {}({}) - expected {} bytes, got {} bytes after {} attempts",
1138
                    key.borrow().as_str(),
1139
                    temp_key,
1140
                    total_len,
1141
                    blob_len,
1142
                    attempt,
1143
                ));
1144
            }
1145
            let (connection_manager, uuid) = self.connection_manager.reconnect(client.uuid).await?;
1146
            client.connection_manager = connection_manager;
1147
            client.uuid = uuid;
1148
            sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;
1149
        };
1150
1151
        // Rename the temp key so that the data appears under the real key. Any data already present in the real key is lost.
1152
        // Reconnect once on a transient failover error in case the master moved between the verify and here.
1153
        match client
1154
            .connection_manager
1155
            .rename::<_, _, ()>(&temp_key, final_key.as_ref())
1156
            .await
1157
        {
1158
            Ok(()) => {}
1159
            Err(err) if is_retryable_redis_error(&err) => {
1160
                let (connection_manager, uuid) =
1161
                    self.connection_manager.reconnect(client.uuid).await?;
1162
                client.connection_manager = connection_manager;
1163
                client.uuid = uuid;
1164
                client
1165
                    .connection_manager
1166
                    .rename::<_, _, ()>(&temp_key, final_key.as_ref())
1167
                    .await
1168
                    .err_tip(
1169
                        || "While queueing key rename (after reconnect) in RedisStore::update()",
1170
                    )?;
1171
            }
1172
            Err(err) => {
1173
                return Err(
1174
                    Error::from(err).append("While queueing key rename in RedisStore::update()")
1175
                );
1176
            }
1177
        }
1178
1179
        // If we have a publish channel configured, send a notice that the key has been set.
1180
        if let Some(pub_sub_channel) = &self.pub_sub_channel {
1181
            client
1182
                .connection_manager
1183
                .publish::<_, _, ()>(pub_sub_channel, final_key.as_ref())
1184
                .await?;
1185
        }
1186
1187
        Ok(blob_len.try_into().unwrap_or(0))
1188
12
    }
1189
1190
    async fn get_part(
1191
        self: Pin<&Self>,
1192
        key: StoreKey<'_>,
1193
        writer: &mut DropCloserWriteHalf,
1194
        offset: u64,
1195
        length: Option<u64>,
1196
5
    ) -> Result<(), Error> {
1197
        let offset = isize::try_from(offset).err_tip(|| "Could not convert offset to isize")?;
1198
        let length = length
1199
4
            .map(|v| usize::try_from(v).err_tip(|| "Could not convert length to usize"))
1200
            .transpose()?;
1201
1202
        // To follow RBE spec we need to consider any digest's with
1203
        // zero size to be existing.
1204
        if is_zero_digest(key.borrow()) {
1205
            return writer
1206
                .send_eof()
1207
                .err_tip(|| "Failed to send zero EOF in redis store get_part");
1208
        }
1209
1210
        let encoded_key = self.encode_key(&key);
1211
        let encoded_key = encoded_key.as_ref();
1212
1213
        // N.B. the `-1`'s you see here are because redis GETRANGE is inclusive at both the start and end, so when we
1214
        // do math with indices we change them to be exclusive at the end.
1215
1216
        // We want to read the data at the key from `offset` to `offset + length`.
1217
        let data_start = offset;
1218
        let data_end = data_start
1219
            .saturating_add(length.unwrap_or(isize::MAX as usize) as isize)
1220
            .saturating_sub(1);
1221
1222
        // And we don't ever want to read more than `read_chunk_size` bytes at a time, so we'll need to iterate.
1223
        let mut chunk_start = data_start;
1224
        let mut chunk_end = cmp::min(
1225
            data_start.saturating_add(self.read_chunk_size as isize) - 1,
1226
            data_end,
1227
        );
1228
1229
        let mut client = self.get_client().await?;
1230
        loop {
1231
            // getrange is position-based and idempotent, so re-resolve the
1232
            // master and retry on a transient failover without re-sending
1233
            // already-written chunks.
1234
            let chunk: Bytes = {
1235
                let mut attempt: u32 = 0;
1236
                loop {
1237
                    attempt += 1;
1238
                    match client
1239
                        .connection_manager
1240
                        .getrange(encoded_key, chunk_start, chunk_end)
1241
                        .await
1242
                    {
1243
                        Ok(v) => break v,
1244
                        Err(err)
1245
                            if attempt < MAX_REDIS_RETRY_ATTEMPTS
1246
                                && is_retryable_redis_error(&err) =>
1247
                        {
1248
                            client.reconnect(&self.connection_manager).await?;
1249
                            sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;
1250
                        }
1251
                        Err(err) => {
1252
                            return Err(
1253
                                Error::from(err).append("In RedisStore::get_part::getrange")
1254
                            );
1255
                        }
1256
                    }
1257
                }
1258
            };
1259
1260
            let didnt_receive_full_chunk = chunk.len() < self.read_chunk_size;
1261
            let reached_end_of_data = chunk_end == data_end;
1262
1263
            if didnt_receive_full_chunk || reached_end_of_data {
1264
                if !chunk.is_empty() {
1265
                    writer
1266
                        .send(chunk)
1267
                        .await
1268
                        .err_tip(|| "Failed to write data in RedisStore::get_part")?;
1269
                }
1270
1271
                break; // No more data to read.
1272
            }
1273
1274
            // We received a full chunk's worth of data, so write it...
1275
            writer
1276
                .send(chunk)
1277
                .await
1278
                .err_tip(|| "Failed to write data in RedisStore::get_part")?;
1279
1280
            // ...and go grab the next chunk.
1281
            chunk_start = chunk_end + 1;
1282
            chunk_end = cmp::min(
1283
                chunk_start.saturating_add(self.read_chunk_size as isize) - 1,
1284
                data_end,
1285
            );
1286
        }
1287
1288
        // If we didn't write any data, check if the key exists, if not return a NotFound error.
1289
        // This is required by spec.
1290
        if writer.get_bytes_written() == 0 {
1291
            // We're supposed to read 0 bytes, so just check if the key exists.
1292
            let exists: bool = {
1293
                let mut attempt: u32 = 0;
1294
                loop {
1295
                    attempt += 1;
1296
                    match client.connection_manager.exists(encoded_key).await {
1297
                        Ok(v) => break v,
1298
                        Err(err)
1299
                            if attempt < MAX_REDIS_RETRY_ATTEMPTS
1300
                                && is_retryable_redis_error(&err) =>
1301
                        {
1302
                            client.reconnect(&self.connection_manager).await?;
1303
                            sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;
1304
                        }
1305
                        Err(err) => {
1306
                            return Err(
1307
                                Error::from(err).append("In RedisStore::get_part::zero_exists")
1308
                            );
1309
                        }
1310
                    }
1311
                }
1312
            };
1313
1314
            if !exists {
1315
                return Err(make_err!(
1316
                    Code::NotFound,
1317
                    "Data not found in Redis store for digest: {key:?}"
1318
                ));
1319
            }
1320
        }
1321
1322
        writer
1323
            .send_eof()
1324
            .err_tip(|| "Failed to write EOF in redis store get_part")
1325
5
    }
1326
1327
0
    fn inner_store(&self, _digest: Option<StoreKey>) -> &dyn StoreDriver {
1328
0
        self
1329
0
    }
1330
1331
0
    fn as_any(&self) -> &(dyn core::any::Any + Sync + Send) {
1332
0
        self
1333
0
    }
1334
1335
0
    fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send> {
1336
0
        self
1337
0
    }
1338
1339
0
    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
1340
0
        registry.register_indicator(self);
1341
0
    }
1342
1343
6
    fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> {
1344
6
        debug!(?callback, "New callback");
1345
6
        let local_self = self.clone();
1346
6
        background_spawn!("remove_callback_subscribe", async move {
1347
6
            self.remove_callbacks.lock().await.push(callback);
1348
6
            if let Err(
err0
) = local_self.clone().has_remove_callback_subscribe
1349
6
                .get_or_try_init(|| async move {
1350
6
                    let mut client = local_self.get_client().await
?0
;
1351
6
                    let cfg = redis::cmd("CONFIG").arg("GET").arg("notify-keyspace-events").to_owned().query_async::<Vec<(String,String)>>(&mut client.connection_manager).await.map_err(|e| 
Error::from0
(
e0
).
append0
("Parsing notify-keyspace-events"))
?0
;
1352
6
                    if cfg.len() != 1 {
1353
0
                        warn!(?cfg, "Got multiple items for CONFIG GET, expected one");
1354
0
                        return Err(make_input_err!("Got multiple items for CONFIG GET, expected one"));
1355
6
                    }
1356
6
                    let events_cfg = &cfg.first().ok_or_else(|| 
make_err!0
(
Code::InvalidArgument0
, "Only one item"))
?0
.1;
1357
6
                    if events_cfg.is_empty() {
1358
1
                        error!("notify-keyspace-events not enabled for Redis, will fail to get remove callbacks");
1359
5
                    } else if !events_cfg.contains('K') {
1360
1
                        error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'K' so won't get keyspace events we need for eviction events");
1361
4
                    } else if !events_cfg.contains('A') {
1362
1
                        error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'A' so we won't get eviction events");
1363
3
                    }
1364
                    // FIXME: Redis events spec appears unreliable, so we subscribe anyways
1365
                    // It should just need Ke as per https://redis.io/docs/latest/develop/pubsub/keyspace-notifications/
1366
                    // but I'm yet to get reliable eviction events out of that
1367
6
                    info!(notify_keyspace_events=events_cfg, "Attempting to subscribe to eviction events");
1368
6
                    self.connection_manager.psubscribe("__key*__:*").await
?0
;
1369
6
                    Ok::<(), Error>(())
1370
12
                 })
1371
6
                .await {
1372
0
                    error!(?err, "Error while trying to initialise remove_callback_subscribe");
1373
6
                }
1374
6
        });
1375
6
        Ok(())
1376
6
    }
1377
}
1378
1379
#[async_trait]
1380
impl<C, M> HealthStatusIndicator for RedisStore<C, M>
1381
where
1382
    C: ConnectionLike + Clone + Send + Sync + Unpin + 'static,
1383
    M: RedisManager<C> + Send + Sync + Unpin + 'static,
1384
{
1385
1
    fn get_name(&self) -> &'static str {
1386
1
        "RedisStore"
1387
1
    }
1388
1389
    /// Lightweight health check: just `PING` the master, bounded by a
1390
    /// short physical timeout. The default `StoreDriver::check_health`
1391
    /// performs a full `update_oneshot` + `has` + `get_part_unchunked`
1392
    /// roundtrip, which queues behind real production traffic on the
1393
    /// same connection-permit semaphore and Redis master. When the
1394
    /// store is even moderately loaded that easily exceeds the
1395
    /// `HealthServer` per-indicator budget (default 5 s), each
1396
    /// RedisStore-backed indicator (AC, small-blob CAS, scheduler)
1397
    /// reports `HealthStatus::Timeout`, and `/status` returns 503 —
1398
    /// surfaced as a readiness-probe failure that sheds traffic from
1399
    /// an otherwise-functional pod. A `PING` proves the connection
1400
    /// is reachable and the master is accepting commands; that is
1401
    /// the only invariant a kubelet probe needs.
1402
1
    async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus {
1403
        let mut client = match self.get_client().await {
1404
            Ok(c) => c,
1405
            Err(e) => {
1406
                return HealthStatus::new_failed(
1407
                    self,
1408
                    format!("RedisStore::check_health: failed to acquire connection: {e}").into(),
1409
                );
1410
            }
1411
        };
1412
1413
        // Hold the `ClientWithPermit` for the duration of the call so
1414
        // its `Drop` releases the semaphore permit on exit. We just
1415
        // need a `&mut` to the connection manager underneath.
1416
1
        let ping = async {
1417
1
            redis::cmd("PING")
1418
1
                .query_async::<()>(&mut client.connection_manager)
1419
1
                .await
1420
1
        };
1421
        let retry_reason = match timeout(self.health_check_timeout, ping).await {
1422
            Ok(Ok(())) => {
1423
                return HealthStatus::new_ok(self, "RedisStore::check_health: PING ok".into());
1424
            }
1425
            // A PING that errors connection-wise or times out means the handle
1426
            // points at a master that went away (a Sentinel failover). Without
1427
            // re-resolving here, a store with no other traffic stays wedged on
1428
            // the dead handle and reports unhealthy forever — shedding readiness
1429
            // traffic from an otherwise-recovered pod until it is restarted.
1430
            Ok(Err(e)) if is_retryable_redis_error(&e) => format!("PING errored: {e}"),
1431
            Ok(Err(e)) => {
1432
                return HealthStatus::new_failed(
1433
                    self,
1434
                    format!("RedisStore::check_health: PING errored: {e}").into(),
1435
                );
1436
            }
1437
            Err(_) => format!(
1438
                "PING exceeded {}ms timeout",
1439
                self.health_check_timeout.as_millis()
1440
            ),
1441
        };
1442
1443
        // The reconnect (re-resolving the master via Sentinel) is logged by
1444
        // `StandardRedisManager::reconnect`; `retry_reason` is surfaced in the
1445
        // failure message below if the second PING still doesn't come back.
1446
        if let Err(e) = client.reconnect(&self.connection_manager).await {
1447
            return HealthStatus::new_failed(
1448
                self,
1449
                format!("RedisStore::check_health: {retry_reason}; reconnect failed: {e}").into(),
1450
            );
1451
        }
1452
1
        let ping = async {
1453
1
            redis::cmd("PING")
1454
1
                .query_async::<()>(&mut client.connection_manager)
1455
1
                .await
1456
1
        };
1457
        match timeout(self.health_check_timeout, ping).await {
1458
            Ok(Ok(())) => HealthStatus::new_ok(
1459
                self,
1460
                "RedisStore::check_health: PING ok after re-resolving master".into(),
1461
            ),
1462
            Ok(Err(e)) => HealthStatus::new_failed(
1463
                self,
1464
                format!("RedisStore::check_health: PING still errored after reconnect: {e}").into(),
1465
            ),
1466
            Err(_) => HealthStatus::new_failed(
1467
                self,
1468
                format!(
1469
                    "RedisStore::check_health: PING still exceeded {}ms timeout after reconnect",
1470
                    self.health_check_timeout.as_millis()
1471
                )
1472
                .into(),
1473
            ),
1474
        }
1475
1
    }
1476
}
1477
1478
// -------------------------------------------------------------------
1479
// Below this line are specific to the redis scheduler implementation.
1480
// -------------------------------------------------------------------
1481
1482
/// The time in milliseconds that a redis cursor can be idle before it is closed.
1483
const CURSOR_IDLE_MS: u64 = 30_000;
1484
/// The name of the field in the Redis hash that stores the data.
1485
const DATA_FIELD_NAME: &str = "data";
1486
/// The name of the field in the Redis hash that stores the version.
1487
const VERSION_FIELD_NAME: &str = "version";
1488
/// The time to live of indexes in seconds. After this time redis may delete the index.
1489
const INDEX_TTL_S: u64 = 60 * 60 * 24; // 24 hours.
1490
1491
#[allow(rustdoc::broken_intra_doc_links)]
1492
/// Lua script to set a key if the version matches.
1493
/// Args:
1494
///   KEYS[1]: The key where the version is stored.
1495
///   ARGV[1]: The expected version.
1496
///   ARGV[2]: TTL in seconds, or 0 for forever
1497
///   ARGV[3]: The new data.
1498
///   ARGV[4*]: Key-value pairs of additional data to include.
1499
/// Returns:
1500
///   The new version if the version matches. nil is returned if the
1501
///   value was not set.
1502
pub const LUA_VERSION_SET_SCRIPT: &str = formatcp!(
1503
    r"
1504
local key = KEYS[1]
1505
local expected_version = tonumber(ARGV[1])
1506
local ttl = tonumber(ARGV[2])
1507
local new_data = ARGV[3]
1508
local new_version = redis.call('HINCRBY', key, '{VERSION_FIELD_NAME}', 1)
1509
local i
1510
local indexes = {{}}
1511
1512
if new_version-1 ~= expected_version then
1513
    redis.call('HINCRBY', key, '{VERSION_FIELD_NAME}', -1)
1514
    return {{ 0, new_version-1 }}
1515
end
1516
-- Skip first 3 argvs, as they are known inputs.
1517
-- Remember: Lua is 1-indexed.
1518
for i=4, #ARGV do
1519
    indexes[i-3] = ARGV[i]
1520
end
1521
1522
-- In testing we witnessed redis sometimes not update our FT indexes
1523
-- resulting in stale data. It appears if we delete our keys then insert
1524
-- them again it works and reduces risk significantly.
1525
redis.call('DEL', key)
1526
redis.call('HSET', key, '{DATA_FIELD_NAME}', new_data, '{VERSION_FIELD_NAME}', new_version, unpack(indexes))
1527
1528
if ttl ~= 0 then
1529
    redis.call('EXPIRE', key, ttl)
1530
end
1531
return {{ 1, new_version }}
1532
"
1533
);
1534
1535
/// This is the output of the calculations below hardcoded into the executable.
1536
const FINGERPRINT_CREATE_INDEX_HEX: &str = "3e762c15";
1537
1538
#[cfg(test)]
1539
mod test {
1540
    use super::FINGERPRINT_CREATE_INDEX_HEX;
1541
1542
    /// String of the `FT.CREATE` command used to create the index template.
1543
    const CREATE_INDEX_TEMPLATE: &str = "FT.CREATE {} ON HASH PREFIX 1 {} NOOFFSETS NOHL NOFIELDS NOFREQS SCHEMA {} TAG CASESENSITIVE SORTABLE";
1544
1545
    /// Compile-time fingerprint of the `FT.CREATE` command used to create the
1546
    /// index template. This is a simple CRC32 checksum of the command string.
1547
    /// We don't care about it actually being a valid CRC32 checksum, just that
1548
    /// it's a unique identifier with a low chance of collision.
1549
1
    const fn fingerprint_create_index_template() -> u32 {
1550
        const POLY: u32 = 0xEDB8_8320;
1551
        const DATA: &[u8] = CREATE_INDEX_TEMPLATE.as_bytes();
1552
1
        let mut crc = 0xFFFF_FFFF;
1553
1
        let mut i = 0;
1554
102
        while i < DATA.len() {
1555
101
            let byte = DATA[i];
1556
101
            crc ^= byte as u32;
1557
1558
101
            let mut j = 0;
1559
909
            while j < 8 {
1560
808
                crc = if crc & 1 != 0 {
1561
386
                    (crc >> 1) ^ POLY
1562
                } else {
1563
422
                    crc >> 1
1564
                };
1565
808
                j += 1;
1566
            }
1567
101
            i += 1;
1568
        }
1569
1
        crc
1570
1
    }
1571
1572
    /// Verify that our calculation always evaluates to this fixed value.
1573
    #[test]
1574
1
    fn test_fingerprint_value() {
1575
1
        assert_eq!(
1576
1
            format!("{:08x}", &fingerprint_create_index_template()),
1577
            FINGERPRINT_CREATE_INDEX_HEX,
1578
        );
1579
1
    }
1580
}
1581
1582
/// Get the name of the index to create for the given field.
1583
/// This will add some prefix data to the name to try and ensure
1584
/// if the index definition changes, the name will get a new name.
1585
macro_rules! get_index_name {
1586
    ($prefix:expr, $field:expr, $maybe_sort:expr) => {
1587
        format_args!(
1588
            "{}_{}_{}_{}",
1589
            $prefix,
1590
            $field,
1591
            $maybe_sort.unwrap_or(""),
1592
            FINGERPRINT_CREATE_INDEX_HEX
1593
        )
1594
    };
1595
}
1596
1597
/// Try to sanitize a string to be used as a Redis key.
1598
/// We don't actually modify the string, just check if it's valid.
1599
31
const fn try_sanitize(s: &str) -> bool {
1600
    // Note: We cannot use for loops or iterators here because they are not const.
1601
    // Allowing us to use a const function here gives the compiler the ability to
1602
    // optimize this function away entirely in the case where the input is constant.
1603
31
    let chars = s.as_bytes();
1604
31
    let mut i: usize = 0;
1605
31
    let len = s.len();
1606
    loop {
1607
737
        if i >= len {
1608
31
            break;
1609
706
        }
1610
706
        let c = chars[i];
1611
706
        if !c.is_ascii_alphanumeric() && 
c != b'_'30
{
1612
0
            return false;
1613
706
        }
1614
706
        i += 1;
1615
    }
1616
31
    true
1617
31
}
1618
1619
/// An individual subscription to a key in Redis.
1620
#[derive(Debug)]
1621
pub struct RedisSubscription {
1622
    receiver: Option<tokio::sync::watch::Receiver<String>>,
1623
    weak_subscribed_keys: Weak<RwLock<StringPatriciaMap<RedisSubscriptionPublisher>>>,
1624
}
1625
1626
impl SchedulerSubscription for RedisSubscription {
1627
    /// Wait for the subscription key to change.
1628
23
    async fn changed(&mut self) -> Result<(), Error> {
1629
23
        let receiver = self
1630
23
            .receiver
1631
23
            .as_mut()
1632
23
            .ok_or_else(|| 
make_err!0
(
Code::Internal0
, "In RedisSubscription::changed::as_mut"))
?0
;
1633
23
        receiver.changed().await.
map_err19
(|err|
{0
1634
0
            Error::from_std_err(Code::Internal, &err)
1635
0
                .append("In RedisSubscription::changed::changed")
1636
0
        })
1637
19
    }
1638
}
1639
1640
// If the subscription is dropped, we need to possibly remove the key from the
1641
// subscribed keys map.
1642
impl Drop for RedisSubscription {
1643
411
    fn drop(&mut self) {
1644
411
        let Some(receiver) = self.receiver.take() else {
1645
0
            warn!("RedisSubscription has already been dropped, nothing to do.");
1646
0
            return; // Already dropped, nothing to do.
1647
        };
1648
411
        let key = receiver.borrow().clone();
1649
411
        let Some(
subscribed_keys408
) = self.weak_subscribed_keys.upgrade() else {
1650
3
            return; // Parent dropped — nothing to do.
1651
        };
1652
408
        let mut subscribed_keys = subscribed_keys.write();
1653
408
        let Some(publisher) = subscribed_keys.get(&key) else {
1654
0
            warn!(
1655
                %key,
1656
                "RedisSubscription::drop: key absent from subscribed_keys under write lock — \
1657
                 indicates an unexpected removal path",
1658
            );
1659
0
            return;
1660
        };
1661
        // Count includes our own (still-alive) receiver. If we are the
1662
        // sole subscriber, remove the publisher entry.
1663
408
        if publisher.receiver_count() == 1 {
1664
206
            subscribed_keys.remove(&key);
1665
206
        
}202
1666
408
        drop(receiver);
1667
411
    }
1668
}
1669
1670
/// A publisher for a key in Redis.
1671
#[derive(Debug)]
1672
struct RedisSubscriptionPublisher {
1673
    sender: Mutex<tokio::sync::watch::Sender<String>>,
1674
}
1675
1676
impl RedisSubscriptionPublisher {
1677
209
    fn new(
1678
209
        key: String,
1679
209
        weak_subscribed_keys: Weak<RwLock<StringPatriciaMap<Self>>>,
1680
209
    ) -> (Self, RedisSubscription) {
1681
209
        let (sender, receiver) = tokio::sync::watch::channel(key);
1682
209
        let publisher = Self {
1683
209
            sender: Mutex::new(sender),
1684
209
        };
1685
209
        let subscription = RedisSubscription {
1686
209
            receiver: Some(receiver),
1687
209
            weak_subscribed_keys,
1688
209
        };
1689
209
        (publisher, subscription)
1690
209
    }
1691
1692
202
    fn subscribe(
1693
202
        &self,
1694
202
        weak_subscribed_keys: Weak<RwLock<StringPatriciaMap<Self>>>,
1695
202
    ) -> RedisSubscription {
1696
202
        let receiver = self.sender.lock().subscribe();
1697
202
        RedisSubscription {
1698
202
            receiver: Some(receiver),
1699
202
            weak_subscribed_keys,
1700
202
        }
1701
202
    }
1702
1703
408
    fn receiver_count(&self) -> usize {
1704
408
        self.sender.lock().receiver_count()
1705
408
    }
1706
1707
19
    fn notify(&self) {
1708
        // TODO(https://github.com/sile/patricia_tree/issues/40) When this is addressed
1709
        // we can remove the `Mutex` and use the mutable iterator directly.
1710
19
        self.sender.lock().send_modify(|_| {});
1711
19
    }
1712
}
1713
1714
#[derive(Debug, Clone)]
1715
pub struct RedisSubscriptionManager {
1716
    subscribed_keys: Arc<RwLock<StringPatriciaMap<RedisSubscriptionPublisher>>>,
1717
    tx_for_test: UnboundedSender<String>,
1718
    _subscription_spawn: Arc<Mutex<JoinHandleDropGuard<()>>>,
1719
}
1720
1721
impl RedisSubscriptionManager {
1722
53
    pub fn new(
1723
53
        subscriber_channel: UnboundedReceiver<PushInfo>,
1724
53
        remove_callbacks: Arc<async_lock::Mutex<Vec<RemoveCallback>>>,
1725
53
        key_prefix: String,
1726
53
    ) -> Self {
1727
53
        let subscribed_keys = Arc::new(RwLock::new(StringPatriciaMap::new()));
1728
53
        let subscribed_keys_weak = Arc::downgrade(&subscribed_keys);
1729
53
        let (tx_for_test, mut rx_for_test) = unbounded_channel();
1730
53
        let mut local_subscriber_channel = UnboundedReceiverStream::new(subscriber_channel);
1731
        Self {
1732
53
            subscribed_keys,
1733
53
            tx_for_test,
1734
53
            _subscription_spawn: Arc::new(Mutex::new(spawn!(
1735
                "redis_subscribe_spawn",
1736
31
                async move {
1737
31
                    debug!("running subscribe loop");
1738
                    loop {
1739
                        loop {
1740
52
                            let 
key17
= select! {
1741
52
                                
value16
= rx_for_test.recv() => {
1742
16
                                    let Some(value) = value else {
1743
0
                                        unreachable!("Channel should never close");
1744
                                    };
1745
16
                                    value
1746
                                },
1747
52
                                
maybe_push_info14
= local_subscriber_channel.next() => {
1748
14
                                    if let Some(
push_info5
) = maybe_push_info {
1749
5
                                        match push_info.kind {
1750
4
                                            redis::PushKind::PMessage => {},
1751
                                            redis::PushKind::PSubscribe => {
1752
1
                                                trace!(?push_info, "PSubscribe, ignore");
1753
1
                                                continue;
1754
                                            }
1755
                                            _ => {
1756
0
                                                warn!(?push_info, "Other push_info message, discarded");
1757
0
                                                continue;
1758
                                            },
1759
                                        }
1760
4
                                        if push_info.data.len() != 3 {
1761
0
                                            error!(?push_info, "Expected exactly 3 values on subscriber channel (pattern, channel, value)");
1762
0
                                            continue;
1763
4
                                        }
1764
4
                                        let value = match push_info.data.last().unwrap() {
1765
0
                                            Value::SimpleString(s) => {
1766
0
                                                s.clone()
1767
                                            }
1768
4
                                            Value::BulkString(v) => {
1769
4
                                                String::from_utf8(v.clone()).expect("String message")
1770
                                            }
1771
0
                                            other => {
1772
0
                                                error!(?other, "Received non-string message in RedisSubscriptionManager");
1773
0
                                                continue;
1774
                                            }
1775
                                        };
1776
4
                                        if value == "evicted" {
1777
3
                                            trace!(?push_info, "Eviction event");
1778
3
                                            let eviction_key = if let Some(key) = push_info.data.get(1) {
1779
3
                                                if let Value::BulkString(s) = key {
1780
3
                                                    String::from_utf8(s.clone()).expect("String message")
1781
                                                } else {
1782
0
                                                    error!(?push_info, "Eviction key wasn't bulk-string");
1783
0
                                                    continue;
1784
                                                }
1785
                                            } else {
1786
0
                                                error!(?push_info, "No key in eviction event");
1787
0
                                                continue;
1788
                                            };
1789
3
                                            trace!(?eviction_key, "Eviction key");
1790
3
                                            let Some((_prefix, internal_key)) = eviction_key.split_once(':') else {
1791
0
                                                error!(?eviction_key, "Eviction key doesn't contain a colon");
1792
0
                                                continue;
1793
                                            };
1794
1795
3
                                            let store_key = match decode_key(&key_prefix, Cow::from(internal_key)) {
1796
3
                                                Ok(k) => k.into_owned(),
1797
0
                                                Err(err) => {
1798
0
                                                    error!(%err, internal_key, "Bad redis key");
1799
0
                                                    continue;
1800
                                                }
1801
                                            };
1802
3
                                            let locked_remove_callbacks = remove_callbacks.lock().await;
1803
3
                                            let mut callbacks: FuturesUnordered<_> =
1804
3
                                                locked_remove_callbacks.iter()
1805
3
                                                .map(|callback| callback.callback(store_key.borrow()))
1806
3
                                                .collect();
1807
6
                                            while callbacks.next().await.is_some() 
{}3
1808
3
                                            continue
1809
1
                                        }
1810
1
                                        value
1811
                                    } else {
1812
9
                                        error!("Error receiving message in RedisSubscriptionManager from subscriber_channel");
1813
9
                                        break;
1814
                                    }
1815
                                }
1816
                            };
1817
17
                            trace!(key, "New subscription manager key");
1818
17
                            let Some(subscribed_keys) = subscribed_keys_weak.upgrade() else {
1819
0
                                warn!(
1820
                                    "It appears our parent has been dropped, exiting RedisSubscriptionManager spawn"
1821
                                );
1822
0
                                return;
1823
                            };
1824
17
                            let subscribed_keys_mux = subscribed_keys.read();
1825
17
                            subscribed_keys_mux
1826
17
                                .common_prefix_values(&*key)
1827
17
                                .for_each(RedisSubscriptionPublisher::notify);
1828
                        }
1829
                        // Sleep for a small amount of time to ensure we don't reconnect too quickly.
1830
9
                        sleep(Duration::from_secs(1)).await;
1831
                        // If we reconnect or lag behind we might have had dirty keys, so we need to
1832
                        // flag all of them as changed.
1833
0
                        let Some(subscribed_keys) = subscribed_keys_weak.upgrade() else {
1834
0
                            warn!(
1835
                                "It appears our parent has been dropped, exiting RedisSubscriptionManager spawn"
1836
                            );
1837
0
                            return;
1838
                        };
1839
0
                        let subscribed_keys_mux = subscribed_keys.read();
1840
                        // Just in case also get a new receiver.
1841
0
                        for publisher in subscribed_keys_mux.values() {
1842
0
                            publisher.notify();
1843
0
                        }
1844
                    }
1845
0
                }
1846
            ))),
1847
        }
1848
53
    }
1849
}
1850
1851
impl SubscriptionManagerNotify for RedisSubscriptionManager {
1852
18
    fn notify_for_test(&self, value: String) {
1853
18
        self.tx_for_test.send(value).unwrap();
1854
18
    }
1855
}
1856
1857
impl SchedulerSubscriptionManager for RedisSubscriptionManager {
1858
    type Subscription = RedisSubscription;
1859
1860
411
    fn subscribe<K>(&self, key: K) -> Result<Self::Subscription, Error>
1861
411
    where
1862
411
        K: SchedulerStoreKeyProvider,
1863
    {
1864
411
        let weak_subscribed_keys = Arc::downgrade(&self.subscribed_keys);
1865
411
        let mut subscribed_keys = self.subscribed_keys.write();
1866
411
        let key = key.get_key();
1867
411
        let key_str = key.as_str();
1868
411
        let mut subscription = if let Some(
publisher202
) = subscribed_keys.get(&key_str) {
1869
202
            publisher.subscribe(weak_subscribed_keys)
1870
        } else {
1871
209
            let (publisher, subscription) =
1872
209
                RedisSubscriptionPublisher::new(key_str.to_string(), weak_subscribed_keys);
1873
209
            subscribed_keys.insert(key_str, publisher);
1874
209
            subscription
1875
        };
1876
411
        subscription
1877
411
            .receiver
1878
411
            .as_mut()
1879
411
            .ok_or_else(|| 
{0
1880
0
                make_err!(
1881
0
                    Code::Internal,
1882
                    "Receiver should be set in RedisSubscriptionManager::subscribe"
1883
                )
1884
0
            })?
1885
411
            .mark_changed();
1886
1887
411
        Ok(subscription)
1888
411
    }
1889
1890
1
    fn is_reliable() -> bool {
1891
1
        false
1892
1
    }
1893
}
1894
1895
impl<C, M> SchedulerStore for RedisStore<C, M>
1896
where
1897
    C: Clone + ConnectionLike + Sync + Send + 'static,
1898
    M: RedisManager<C> + Sync + Send + 'static,
1899
{
1900
    type SubscriptionManager = RedisSubscriptionManager;
1901
1902
8
    async fn subscription_manager(&self) -> Result<Arc<RedisSubscriptionManager>, Error> {
1903
8
        if self.pub_sub_channel.is_none() {
1904
0
            return Err(make_input_err!(
1905
0
                "RedisStore must have a pubsub for Redis Scheduler if using subscriptions"
1906
0
            ));
1907
8
        }
1908
8
        Ok(self.subscription_manager.clone())
1909
8
    }
1910
1911
26
    async fn update_data<T>(&self, data: T, expiry: Option<Duration>) -> Result<Option<i64>, Error>
1912
26
    where
1913
26
        T: SchedulerStoreDataProvider
1914
26
            + SchedulerStoreKeyProvider
1915
26
            + SchedulerCurrentVersionProvider
1916
26
            + Send,
1917
26
    {
1918
26
        let key = data.get_key();
1919
26
        let redis_key = self.encode_key(&key);
1920
26
        let mut client = self.get_client().await
?0
;
1921
26
        let maybe_index = data.get_indexes().err_tip(|| 
{0
1922
0
            format!("Err getting index in RedisStore::update_data::versioned for {redis_key}")
1923
0
        })?;
1924
26
        if <T as SchedulerStoreKeyProvider>::Versioned::VALUE {
1925
19
            let current_version = data.current_version();
1926
19
            let data = data.try_into_bytes().err_tip(|| 
{0
1927
0
                format!("Could not convert value to bytes in RedisStore::update_data::versioned for {redis_key}")
1928
0
            })?;
1929
19
            let mut script = self.connection_manager.update_script(redis_key.as_ref());
1930
19
            let mut script_invocation = script
1931
19
                .arg(format!("{current_version}"))
1932
19
                .arg(expiry.unwrap_or(Duration::ZERO).as_secs())
1933
19
                .arg(data.to_vec());
1934
55
            for (name, value) in 
maybe_index19
{
1935
55
                script_invocation = script_invocation.arg(name).arg(value.to_vec());
1936
55
            }
1937
19
            let start = Instant::now();
1938
19
            let (success, new_version): (bool, i64) = match script_invocation
1939
19
                .invoke_async(&mut client.connection_manager)
1940
19
                .await
1941
            {
1942
19
                Ok(v) => v,
1943
0
                Err(err) if is_retryable_redis_error(&err) => {
1944
0
                    client.reconnect(&self.connection_manager).await?;
1945
0
                    script_invocation
1946
0
                        .invoke_async(&mut client.connection_manager)
1947
0
                        .await
1948
0
                        .err_tip(|| format!("(after reconnect) In RedisStore::update_data::versioned for {key:?}"))?
1949
                }
1950
0
                Err(err) => {
1951
0
                    let mut error: Error = err.into();
1952
0
                    error
1953
0
                        .messages
1954
0
                        .push(format!("In RedisStore::update_data::versioned for {key:?}"));
1955
0
                    return Err(error);
1956
                }
1957
            };
1958
1959
19
            let elapsed = start.elapsed();
1960
1961
19
            if elapsed > Duration::from_millis(100) {
1962
0
                warn!(
1963
                    %redis_key,
1964
                    ?elapsed,
1965
                    "Slow Redis version-set operation"
1966
                );
1967
19
            }
1968
19
            if !success {
1969
4
                warn!(
1970
                    %redis_key,
1971
                    %key,
1972
                    %current_version,
1973
                    %new_version,
1974
4
                    caller = core::any::type_name::<T>(),
1975
                    "Redis version conflict - optimistic lock failed"
1976
                );
1977
4
                return Ok(None);
1978
15
            }
1979
15
            trace!(
1980
                %redis_key,
1981
                %key,
1982
                old_version = %current_version,
1983
                %new_version,
1984
                "Updated redis key to new version"
1985
            );
1986
            // If we have a publish channel configured, send a notice that the key has been set.
1987
15
            if let Some(
pub_sub_channel13
) = &self.pub_sub_channel {
1988
13
                return Ok(client
1989
13
                    .connection_manager
1990
13
                    .publish(pub_sub_channel, redis_key.as_ref())
1991
13
                    .await
?0
);
1992
2
            }
1993
2
            Ok(Some(new_version))
1994
        } else {
1995
7
            let data = data.try_into_bytes().err_tip(|| 
{0
1996
0
                format!("Could not convert value to bytes in RedisStore::update_data::noversion for {redis_key}")
1997
0
            })?;
1998
7
            let mut fields: Vec<(String, _)> = vec![];
1999
7
            fields.push((DATA_FIELD_NAME.into(), data.to_vec()));
2000
7
            for (
name6
,
value6
) in maybe_index {
2001
6
                fields.push((name.into(), value.to_vec()));
2002
6
            }
2003
7
            match client
2004
7
                .connection_manager
2005
7
                .hset_multiple::<_, _, _, ()>(redis_key.as_ref(), &fields)
2006
7
                .await
2007
            {
2008
6
                Ok(_v) => {
2009
6
                    if let Some(expiry_v) = expiry {
2010
6
                        let seconds =
2011
6
                            TryInto::<i64>::try_into(expiry_v.as_secs()).err_tip(|| 
{0
2012
0
                                format!("Expiry seconds doesn't map to i64: {expiry_v:#?}")
2013
0
                            })?;
2014
6
                        let expiry_result: u8 = client
2015
6
                            .connection_manager
2016
6
                            .expire(redis_key.as_ref(), seconds)
2017
6
                            .await
2018
6
                            .err_tip(|| 
{0
2019
0
                                format!(
2020
                                    "In RedisStore::update_data::noversion (expiry) for {redis_key}"
2021
                                )
2022
0
                            })?;
2023
6
                        if expiry_result != 1 {
2024
1
                            warn!(%redis_key, seconds, "Wasn't able to set expiry for Redis key");
2025
5
                        }
2026
0
                    }
2027
                }
2028
1
                Err(err) if is_retryable_redis_error(&err) => {
2029
1
                    client.reconnect(&self.connection_manager).await
?0
;
2030
1
                    client
2031
1
                        .connection_manager
2032
1
                        .hset_multiple::<_, _, _, ()>(redis_key.as_ref(), &fields)
2033
1
                        .await
2034
1
                        .err_tip(|| 
format!0
("(after reconnect) In RedisStore::update_data::noversion (hset) for {redis_key}"))
?0
;
2035
1
                    if let Some(expiry_v) = expiry {
2036
1
                        let seconds =
2037
1
                            TryInto::<i64>::try_into(expiry_v.as_secs()).err_tip(|| 
{0
2038
0
                                format!("Expiry seconds doesn't map to i64: {expiry_v:#?}")
2039
0
                            })?;
2040
1
                        let expiry_result: u8 = client.connection_manager.expire(redis_key.as_ref(), seconds).await
2041
1
                        .err_tip(|| 
format!0
("(after reconnect) In RedisStore::update_data::noversion (expiry) for {redis_key}"))
?0
;
2042
1
                        if expiry_result != 1 {
2043
0
                            warn!(%redis_key, seconds, "Wasn't able to set expiry for Redis key");
2044
1
                        }
2045
0
                    }
2046
                }
2047
0
                Err(err) => {
2048
0
                    let mut error: Error = err.into();
2049
0
                    error.messages.push(format!(
2050
                        "In RedisStore::update_data::noversion for {redis_key}"
2051
                    ));
2052
0
                    return Err(error);
2053
                }
2054
            }
2055
            // If we have a publish channel configured, send a notice that the key has been set.
2056
7
            if let Some(
pub_sub_channel4
) = &self.pub_sub_channel {
2057
4
                return Ok(client
2058
4
                    .connection_manager
2059
4
                    .publish(pub_sub_channel, redis_key.as_ref())
2060
4
                    .await
?0
);
2061
3
            }
2062
3
            Ok(Some(0)) // Always use "0" version since this is not a versioned request.
2063
        }
2064
26
    }
2065
2066
31
    async fn search_by_index_prefix<K>(
2067
31
        &self,
2068
31
        index: K,
2069
31
    ) -> Result<
2070
31
        impl Stream<Item = Result<<K as SchedulerStoreDecodeTo>::DecodeOutput, Error>> + Send,
2071
31
        Error,
2072
31
    >
2073
31
    where
2074
31
        K: SchedulerIndexProvider + SchedulerStoreDecodeTo + Send,
2075
31
    {
2076
31
        let index_value = index.index_value();
2077
31
        try_sanitize(index_value.as_ref())
2078
31
            .then_some(())
2079
31
            .err_tip(|| 
{0
2080
0
                format!("In RedisStore::search_by_index_prefix::try_sanitize - {index_value:?}")
2081
0
            })?;
2082
35
        let 
run_ft_aggregate31
= |connection_manager: C| async {
2083
35
            ft_aggregate(
2084
35
                connection_manager,
2085
35
                format!(
2086
                    "{}",
2087
35
                    get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY)
2088
                ),
2089
35
                if index_value.is_empty() {
2090
2
                    "*".to_string()
2091
                } else {
2092
33
                    format!("@{}:{{ {} }}", K::INDEX_NAME, index_value)
2093
                },
2094
                FtAggregateOptions {
2095
35
                    load: vec![DATA_FIELD_NAME.into(), VERSION_FIELD_NAME.into()],
2096
35
                    cursor: FtAggregateCursor {
2097
35
                        count: self.max_count_per_cursor,
2098
35
                        max_idle: CURSOR_IDLE_MS,
2099
35
                    },
2100
35
                    sort_by: K::MAYBE_SORT_KEY.map_or_else(Vec::new, |v| 
vec!29
[
format!29
("@{v}")]),
2101
                },
2102
            )
2103
35
            .await
2104
70
        };
2105
31
        let run_ft_create = |connection_manager: C| async 
{3
2106
3
            let mut schema = vec![SearchSchema {
2107
3
                field_name: K::INDEX_NAME.into(),
2108
3
                sortable: false,
2109
3
            }];
2110
3
            if let Some(sort_key) = K::MAYBE_SORT_KEY {
2111
3
                schema.push(SearchSchema {
2112
3
                    field_name: sort_key.into(),
2113
3
                    sortable: true,
2114
3
                });
2115
3
            
}0
2116
3
            let create_options = FtCreateOptions {
2117
3
                prefixes: vec![K::KEY_PREFIX.into()],
2118
3
                nohl: true,
2119
3
                nofields: true,
2120
3
                nofreqs: true,
2121
3
                nooffsets: true,
2122
3
                temporary: Some(INDEX_TTL_S),
2123
3
            };
2124
3
            let index = format!(
2125
                "{}",
2126
3
                get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY)
2127
            );
2128
3
            ft_create(connection_manager, index, create_options, schema).await
2129
6
        };
2130
2131
31
        let (connection_manager, connect_id) = self.connection_manager.get_connection().await
?0
;
2132
31
        let 
stream28
= match run_ft_aggregate(connection_manager.clone()).await {
2133
            // A demoted master answers READONLY and a dead/old master drops the
2134
            // connection or times the command out. Both mean the master moved
2135
            // (Sentinel failover) — re-resolve it and retry rather than letting
2136
            // the scheduler's matching loop spin on a stale handle. (A missing
2137
            // index is not retryable here; it falls through to the create path
2138
            // below, which re-runs on the next matching cycle if needed.)
2139
4
            Err(
err1
) if is_retryable_redis_error(&err
)1
=> {
2140
1
                let (connection_manager, _connect_id) =
2141
1
                    self.connection_manager.reconnect(connect_id).await
?0
;
2142
1
                run_ft_aggregate(connection_manager).await.err_tip(|| 
{0
2143
0
                    format!(
2144
                        "Error with reconnected ft_aggregate in RedisStore::search_by_index_prefix({})",
2145
0
                        get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY),
2146
                    )
2147
0
                })
2148
            }
2149
            Err(_) => {
2150
3
                let (connection_manager, result) =
2151
3
                    match run_ft_create(connection_manager.clone()).await {
2152
3
                        Err(
err0
) if is_retryable_redis_error(&err
)0
=> {
2153
0
                            let (connection_manager, _connect_id) =
2154
0
                                self.connection_manager.reconnect(connect_id).await?;
2155
                            (
2156
0
                                connection_manager.clone(),
2157
0
                                run_ft_create(connection_manager).await,
2158
                            )
2159
                        }
2160
3
                        result => (connection_manager, result),
2161
                    };
2162
2163
                // RediSearch returns ErrorKind::Extension with code "Index"
2164
                // and detail along the lines of "Index already exists" when
2165
                // FT.CREATE races with another node.
2166
3
                let create_result = result.or_else(|e| {
2167
3
                    let is_already_exists = e.kind() == redis::ErrorKind::Extension
2168
2
                        && e.code() == Some("Index")
2169
1
                        && e.detail()
2170
1
                            .is_some_and(|d| d.to_ascii_lowercase().contains("already exists"));
2171
3
                    if is_already_exists {
2172
1
                        Ok(())
2173
                    } else {
2174
2
                        Err(e).err_tip(|| {
2175
2
                            format!(
2176
                                "Error with ft_create in RedisStore::search_by_index_prefix({})",
2177
2
                                get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY),
2178
                            )
2179
2
                        })
2180
                    }
2181
3
                });
2182
2183
3
                let run_result = run_ft_aggregate(connection_manager).await.err_tip(|| {
2184
3
                    format!(
2185
                        "Error with second ft_aggregate in RedisStore::search_by_index_prefix({})",
2186
3
                        get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY),
2187
                    )
2188
3
                });
2189
2190
                // Creating the index will race which is ok. If it fails to create, we only
2191
                // error if the second ft_aggregate call fails and fails to create.
2192
3
                run_result.or_else(move |e| create_result.merge(Err(e)))
2193
            }
2194
27
            Ok(stream) => Ok(stream),
2195
3
        }?;
2196
2197
28
        Ok(stream.filter_map(|result| async move 
{17
2198
17
            let raw_redis_map = match result {
2199
17
                Ok(v) => v,
2200
0
                Err(e) => {
2201
                    return Some(
2202
0
                        Err(Error::from(e))
2203
0
                            .err_tip(|| "Error in stream of in RedisStore::search_by_index_prefix"),
2204
                    );
2205
                }
2206
            };
2207
2208
17
            if 
matches!16
(raw_redis_map, Value::Int(_)) {
2209
1
                return None;
2210
16
            }
2211
2212
16
            let Some(redis_map) = raw_redis_map.as_sequence() else {
2213
0
                return Some(Err(Error::new(
2214
0
                    Code::Internal,
2215
0
                    format!("Non-array from ft_aggregate: {raw_redis_map:?}"),
2216
0
                )));
2217
            };
2218
16
            let mut redis_map_iter = redis_map.iter();
2219
16
            let mut bytes_data: Option<Bytes> = None;
2220
16
            let mut version: Option<i64> = None;
2221
49
            while let Some(
key33
) = redis_map_iter.next() {
2222
33
                let value = redis_map_iter.next().unwrap();
2223
33
                let Value::BulkString(k) = key else {
2224
0
                    return Some(Err(Error::new(
2225
0
                        Code::Internal,
2226
0
                        format!("Non-BulkString key from ft_aggregate: {key:?}"),
2227
0
                    )));
2228
                };
2229
33
                let Ok(str_key) = str::from_utf8(k) else {
2230
0
                    return Some(Err(Error::new(
2231
0
                        Code::Internal,
2232
0
                        format!("Non-utf8 key from ft_aggregate: {key:?}"),
2233
0
                    )));
2234
                };
2235
33
                let Value::BulkString(v) = value else {
2236
0
                    return Some(Err(Error::new(
2237
0
                        Code::Internal,
2238
0
                        format!("Non-BulkString value from ft_aggregate: {key:?}"),
2239
0
                    )));
2240
                };
2241
33
                match str_key {
2242
33
                    DATA_FIELD_NAME => {
2243
16
                        bytes_data = Some(v.clone().into());
2244
16
                    }
2245
17
                    VERSION_FIELD_NAME => {
2246
16
                        let Ok(str_v) = str::from_utf8(v) else {
2247
0
                            return Some(Err(Error::new(
2248
0
                                Code::Internal,
2249
0
                                format!("Non-utf8 version value from ft_aggregate: {v:?}"),
2250
0
                            )));
2251
                        };
2252
16
                        let Ok(raw_version) = str_v.parse::<i64>() else {
2253
0
                            return Some(Err(Error::new(
2254
0
                                Code::Internal,
2255
0
                                format!("Non-integer version value from ft_aggregate: {str_v:?}"),
2256
0
                            )));
2257
                        };
2258
16
                        version = Some(raw_version);
2259
                    }
2260
1
                    other => {
2261
1
                        if K::MAYBE_SORT_KEY == Some(other) {
2262
1
                            // ignore sort keys
2263
1
                        } else {
2264
0
                            return Some(Err(Error::new(
2265
0
                                Code::Internal,
2266
0
                                format!("Extra keys from ft_aggregate: {other}"),
2267
0
                            )));
2268
                        }
2269
                    }
2270
                }
2271
            }
2272
16
            let Some(found_bytes_data) = bytes_data else {
2273
0
                return Some(Err(Error::new(
2274
0
                    Code::Internal,
2275
0
                    format!("Missing '{DATA_FIELD_NAME}' in ft_aggregate, got: {raw_redis_map:?}"),
2276
0
                )));
2277
            };
2278
            Some(
2279
16
                K::decode(version.unwrap_or(0), found_bytes_data)
2280
16
                    .err_tip(|| "In RedisStore::search_by_index_prefix::decode"),
2281
            )
2282
34
        }))
2283
31
    }
2284
2285
94
    async fn get_and_decode<K>(
2286
94
        &self,
2287
94
        key: K,
2288
94
    ) -> Result<Option<<K as SchedulerStoreDecodeTo>::DecodeOutput>, Error>
2289
94
    where
2290
94
        K: SchedulerStoreKeyProvider + SchedulerStoreDecodeTo + Send,
2291
94
    {
2292
94
        let key = key.get_key();
2293
94
        let key = self.encode_key(&key);
2294
94
        let mut client = self.get_client().await
?0
;
2295
        // hmget is idempotent, so re-resolve the master and retry on a transient
2296
        // failover (matching the read paths in get_part/list) instead of failing
2297
        // a scheduler-state read while the master is moving.
2298
94
        let results: Vec<Value> = {
2299
94
            let mut attempt: u32 = 0;
2300
            loop {
2301
94
                attempt += 1;
2302
94
                match client
2303
94
                    .connection_manager
2304
94
                    .hmget::<_, Vec<String>, Vec<Value>>(
2305
94
                        key.as_ref(),
2306
94
                        vec![VERSION_FIELD_NAME.into(), DATA_FIELD_NAME.into()],
2307
                    )
2308
94
                    .await
2309
                {
2310
94
                    Ok(v) => break v,
2311
0
                    Err(err)
2312
0
                        if attempt < MAX_REDIS_RETRY_ATTEMPTS && is_retryable_redis_error(&err) =>
2313
                    {
2314
0
                        client.reconnect(&self.connection_manager).await?;
2315
0
                        sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;
2316
                    }
2317
0
                    Err(err) => {
2318
0
                        return Err(Error::from(err).append(format!(
2319
0
                            "In RedisStore::get_without_version::notversioned {key}"
2320
0
                        )));
2321
                    }
2322
                }
2323
            }
2324
        };
2325
94
        let Some(Value::BulkString(
data48
)) = results.get(1) else {
2326
46
            return Ok(None);
2327
        };
2328
        #[allow(clippy::get_first)]
2329
48
        let version = if let Some(raw_v) = results.get(0) {
2330
48
            match raw_v {
2331
0
                Value::Int(v) => *v,
2332
46
                Value::BulkString(v) => i64::from_str(str::from_utf8(v).expect("utf-8 bulkstring"))
2333
46
                    .expect("integer bulkstring"),
2334
2
                Value::Nil => 0,
2335
                _ => {
2336
0
                    warn!(?raw_v, "Non-integer version!");
2337
0
                    0
2338
                }
2339
            }
2340
        } else {
2341
0
            0
2342
        };
2343
        Ok(Some(
2344
48
            K::decode(version, Bytes::from(data.clone())).err_tip(|| 
{0
2345
0
                format!("In RedisStore::get_with_version::notversioned::decode {key}")
2346
0
            })?,
2347
        ))
2348
94
    }
2349
}