/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 | | } |