// Copyright 2024-2026 The NativeLink Authors. All rights reserved.

//

// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//    See LICENSE file for details

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

use core::cmp;

use core::fmt::Debug;

use core::marker::PhantomData;

use core::ops::{Bound, RangeBounds};

use core::pin::Pin;

use core::str::FromStr;

use core::time::Duration;

use std::borrow::Cow;

use std::collections::HashSet;

use std::sync::{Arc, Weak};

use std::time::Instant;

use async_trait::async_trait;

use bytes::Bytes;

use const_format::formatcp;

use futures::stream::FuturesUnordered;

use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt, future};

use itertools::izip;

use nativelink_config::stores::{RedisMode, RedisSpec};

use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err};

use nativelink_metric::MetricsComponent;

use nativelink_redis_tester::SubscriptionManagerNotify;

use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};

use nativelink_util::common::DigestInfo;

use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};

use nativelink_util::metrics::{record_connection_acquired, record_connection_reconnect};

use nativelink_util::store_trait::{

    BoolValue, RemoveCallback, SchedulerCurrentVersionProvider, SchedulerIndexProvider,

    SchedulerStore, SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider,

    SchedulerSubscription, SchedulerSubscriptionManager, StoreDriver, StoreKey, UploadSizeInfo,

};

use nativelink_util::task::JoinHandleDropGuard;

use nativelink_util::{background_spawn, spawn};

use parking_lot::{Mutex, RwLock};

use patricia_tree::StringPatriciaMap;

use redis::aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig};

use redis::cluster::ClusterClient;

use redis::cluster_async::ClusterConnection;

use redis::sentinel::{SentinelClient, SentinelNodeConnectionInfo, SentinelServerType};

use redis::{

    AsyncCommands, AsyncIter, Client, IntoConnectionInfo, PushInfo, ScanOptions, Script, Value,

    pipe,

};

use serde::Deserialize;

use serde::de::IntoDeserializer;

use tokio::select;

use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};

use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore};

use tokio::time::{sleep, timeout};

use tokio_stream::wrappers::UnboundedReceiverStream;

use tracing::{debug, error, info, trace, warn};

use url::Url;

use uuid::Uuid;

use crate::cas_utils::is_zero_digest;

use crate::redis_utils::{

    FtAggregateCursor, FtAggregateOptions, FtCreateOptions, SearchSchema, ft_aggregate, ft_create,

    ft_search_count,

};

/// The default size of the read chunk when reading data from Redis.

/// Note: If this changes it should be updated in the config documentation.

const DEFAULT_READ_CHUNK_SIZE: usize = 64 * 1024;

/// The default size of the connection pool if not specified.

/// Note: If this changes it should be updated in the config documentation.

const DEFAULT_CONNECTION_POOL_SIZE: usize = 3;

/// The default delay between retries if not specified.

/// Note: If this changes it should be updated in the config documentation.

const DEFAULT_RETRY_DELAY: f32 = 0.1;

/// Maximum attempts (re-resolving the master each time) for an operation to

/// ride out a transient Redis topology change — a failover, a dropped

/// connection, or brief replica lag. Used by `update`'s verify and by the read

/// paths so a transient error doesn't fail an otherwise healthy request.

const MAX_REDIS_RETRY_ATTEMPTS: u32 = 5;

/// Whether a Redis error is worth re-resolving the master and retrying: a

/// dropped/refused connection, an IO error, or a write that hit a freshly

/// demoted replica. Anything else (e.g. a real protocol/logic error) is

/// returned as-is.

fn is_retryable_redis_error(err: &redis::RedisError) -> bool {

    err.is_connection_dropped()

        || err.is_connection_refusal()

        || err.is_io_error()

        || err.kind() == redis::ErrorKind::Server(redis::ServerErrorKind::ReadOnly)

}

/// The default connection timeout in milliseconds if not specified.

/// Note: If this changes it should be updated in the config documentation.

const DEFAULT_CONNECTION_TIMEOUT_MS: u64 = 3000;

/// The default command timeout in milliseconds if not specified.

/// Note: If this changes it should be updated in the config documentation.

const DEFAULT_COMMAND_TIMEOUT_MS: u64 = 10_000;

/// The default `check_health` PING ceiling in milliseconds if not specified.

/// Note: If this changes it should be updated in the config documentation.

const DEFAULT_HEALTH_CHECK_TIMEOUT_MS: u64 = 4000;

/// The default maximum number of chunk uploads per update.

/// Note: If this changes it should be updated in the config documentation.

pub const DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE: usize = 10;

/// The default COUNT value passed when scanning keys in Redis.

/// Note: If this changes it should be updated in the config documentation.

const DEFAULT_SCAN_COUNT: usize = 10_000;

/// The default COUNT value passed when scanning search indexes

/// Note: If this changes it should be updated in the config documentation.

pub const DEFAULT_MAX_COUNT_PER_CURSOR: u64 = 1_500;

const DEFAULT_CLIENT_PERMITS: usize = 500;

/// Converts a TTL to the seconds `EXPIRE` takes.

///

/// Clamped to at least one second. Redis treats `EXPIRE` with a value of zero

/// or less as "delete now", so a sub-second duration reaching here would

/// destroy the value it was meant to protect. Config already rejects zero, so

/// this only guards a future caller.

fn ttl_seconds(key_ttl: Duration) -> i64 {

    i64::try_from(key_ttl.as_secs()).unwrap_or(i64::MAX).max(1)

}

/// A wrapper around Redis to allow it to be reconnected.

pub trait RedisManager<C>

where

    C: ConnectionLike + Clone,

{

    /// Get a connection manager and a unique identifier for this connection

    /// which may be used to issue a reconnect later.

    fn get_connection(&self) -> impl Future<Output = Result<(C, Uuid), Error>> + Send;

    /// Reconnect if the uuid matches the uuid returned from `get_connection()`.

    fn reconnect(&self, uuid: Uuid) -> impl Future<Output = Result<(C, Uuid), Error>> + Send;

    /// Get an invocation of the update version script for a given `key`.

    fn update_script(&self, key: &str) -> redis::ScriptInvocation<'_>;

    /// Configure the connection to have a psubscribe on it and perform the

    /// subscription on reconnect.

    fn psubscribe(&self, pattern: &str) -> impl Future<Output = Result<(), Error>> + Send;

}

#[derive(Debug)]

pub struct ClusterRedisManager<C>

where

    C: ConnectionLike + Clone,

{

    /// A constant Uuid, we never reconnect.

    uuid: Uuid,

    /// Redis script used to update a value in redis if the version matches.

    /// This is done by incrementing the version number and then setting the new

    /// data only if the version number matches the existing version number.

    update_if_version_matches_script: Script,

    /// The client pool connecting to the backing Redis instance(s).

    connection_manager: C,

}

impl<C> ClusterRedisManager<C>

where

    C: ConnectionLike + Clone,

{

    pub async fn new(mut connection_manager: C) -> Result<Self, Error> {

        let update_if_version_matches_script = Script::new(LUA_VERSION_SET_SCRIPT);

        update_if_version_matches_script

            .load_async(&mut connection_manager)

            .await
?0
;

        Ok(Self {

            uuid: Uuid::new_v4(),

            update_if_version_matches_script,

            connection_manager,

        })

    }

}

impl<C> RedisManager<C> for ClusterRedisManager<C>

where

    C: ConnectionLike + Clone + Send + Sync,

{

    fn get_connection(&self) -> impl Future<Output = Result<(C, Uuid), Error>> + Send {

        future::ready(Ok((self.connection_manager.clone(), self.uuid)))

    }

    fn reconnect(&self, _uuid: Uuid) -> impl Future<Output = Result<(C, Uuid), Error>> + Send {

        self.get_connection()

    }

    fn update_script(&self, key: &str) -> redis::ScriptInvocation<'_> {

        self.update_if_version_matches_script.key(key)

    }

    fn psubscribe(&self, _pattern: &str) -> impl Future<Output = Result<(), Error>> + Send {

        // This is a no-op for cluster connections.

        future::ready(Ok(()))

    }

}

type RedisConnectFuture<C> = dyn Future<Output = Result<C, Error>> + Send;

type RedisConnectFn<C> = dyn Fn() -> Pin<Box<RedisConnectFuture<C>>> + Send + Sync;

pub struct StandardRedisManager<C>

where

    C: ConnectionLike + Clone,

{

    /// Function used to re-connect to Redis.

    connect_func: Box<RedisConnectFn<C>>,

    /// Redis script used to update a value in redis if the version matches.

    /// This is done by incrementing the version number and then setting the new

    /// data only if the version number matches the existing version number.

    update_if_version_matches_script: Script,

    /// The client pool connecting to the backing Redis instance(s) and a Uuid

    /// for this connection in order to avoid multiple reconnection attempts.

    connection_manager: tokio::sync::RwLock<(C, Uuid)>,

    /// Serializes reconnect attempts so a Sentinel master failover triggers a

    /// single re-resolution instead of a thundering herd. Kept separate from

    /// `connection_manager` on purpose: the (potentially multi-second) connect

    /// runs while holding only this lock, so in-flight `get_connection` readers

    /// keep using the existing handle and are never frozen behind a slow

    /// reconnect. See [`Self::reconnect`].

    reconnect_lock: tokio::sync::Mutex<()>,

    /// A list of subscription that should be performed on reconnect.

    subscriptions: Mutex<HashSet<String>>,

}

impl<C> Debug for StandardRedisManager<C>

where

    C: ConnectionLike + Clone,

{

    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {

        f.debug_struct("StandardRedisManager")

            .field(

                "update_if_version_matches_script",

                &self.update_if_version_matches_script,

            )

            .field("subscriptions", &self.subscriptions)

            .finish()

    }

}

impl<C> StandardRedisManager<C>

where

    C: ConnectionLike + Clone + Send + Sync,

{

    async fn configure(&self, connection_manager: &mut C) -> Result<(), Error> {

        self.update_if_version_matches_script

            .load_async(connection_manager)

            .await
?0
;

        Ok(())

    }

    async fn new(connect_func: Box<RedisConnectFn<C>>) -> Result<Self, Error> {

        let 
connection_manager21
 = connect_func().await
?3
;

        let update_if_version_matches_script = Script::new(LUA_VERSION_SET_SCRIPT);

        let connection = Self {

            connect_func,

            update_if_version_matches_script,

            connection_manager: tokio::sync::RwLock::new((connection_manager, Uuid::new_v4())),

            reconnect_lock: tokio::sync::Mutex::new(()),

            subscriptions: Mutex::new(HashSet::new()),

        };

        {

            let mut connection_manager = connection.connection_manager.write().await;

            connection.configure(&mut connection_manager.0).await
?0
;

        }

        Ok(connection)

    }

}

impl RedisManager<ConnectionManager> for StandardRedisManager<ConnectionManager> {

    async fn get_connection(&self) -> Result<(ConnectionManager, Uuid), Error> {

        Ok(self.connection_manager.read().await.clone())

    }

    async fn reconnect(&self, uuid: Uuid) -> Result<(ConnectionManager, Uuid), Error> {

        // Fast path: another caller already reconnected past this generation,

        // so the handle is fresh — hand it back without re-resolving.

        {

            let guard = self.connection_manager.read().await;

            if guard.1 != uuid {

                return Ok(guard.clone());

            }

        }

        // Serialize reconnect attempts on a dedicated lock so a failover storm

        // resolves the new master once. Critically this is NOT the

        // `connection_manager` lock, so the slow connect below does not block

        // the `get_connection` readers that every Redis op needs — they keep

        // using the old (now-failing) handle and fail fast on their own

        // command timeout instead of all freezing behind one reconnect.

        let _reconnect_guard = self.reconnect_lock.lock().await;

        // Re-check: a reconnect may have completed while we waited for the lock.

        {

            let guard = self.connection_manager.read().await;

            if guard.1 != uuid {

                return Ok(guard.clone());

            }

        }

        record_connection_reconnect("redis");

        let mut connection_manager = (self.connect_func)().await
?0
;

        let new_uuid = Uuid::new_v4();

        self.configure(&mut connection_manager).await
?0
;

        let subscriptions = {

            let guard = self.subscriptions.lock();

            guard.iter().map(Clone::clone).collect::<Vec<_>>()

        };

        for 
subscription0
 in subscriptions {

            connection_manager.psubscribe(&subscription).await?;

        }

        // Publish the new handle under a brief exclusive lock.

        {

            let mut guard = self.connection_manager.write().await;

            *guard = (connection_manager.clone(), new_uuid);

        }

        info!(old = %uuid, new = %new_uuid, "StandardRedisManager re-resolved the Redis master");

        Ok((connection_manager, new_uuid))

    }

    fn update_script(&self, key: &str) -> redis::ScriptInvocation<'_> {

        self.update_if_version_matches_script.key(key)

    }

    async fn psubscribe(&self, pattern: &str) -> Result<(), Error> {

        debug!(pattern, "new psubscribe");

        let mut connection = self.get_connection().await
?0
.0;

        let new_subscription = self.subscriptions.lock().insert(String::from(pattern));

        if new_subscription {

            let result = connection.psubscribe(pattern).await;

            if result.is_err() {

                self.subscriptions.lock().remove(pattern);

            }

            result
?0
;

        }

        debug!(pattern, new_subscription, "new psubscribe complete");

        Ok(())

    }

}

/// A [`StoreDriver`] implementation that uses Redis as a backing store.

#[derive(MetricsComponent)]

pub struct RedisStore<C, M>

where

    C: ConnectionLike + Clone,

    M: RedisManager<C>,

{

    /// The client pool connecting to the backing Redis instance(s).

    connection_manager: M,

    /// The underlying connection type in the connection manager.

    _connection_type: PhantomData<C>,

    /// A channel to publish updates to when a key is added, removed, or modified.

    #[metric(

        help = "The pubsub channel to publish updates to when a key is added, removed, or modified"

    )]

    pub_sub_channel: Option<String>,

    /// A function used to generate names for temporary keys.

    temp_name_generator_fn: fn() -> String,

    /// A common prefix to append to all keys before they are sent to Redis.

    ///

    /// See [`RedisStore::key_prefix`](`nativelink_config::stores::RedisStore::key_prefix`).

    #[metric(help = "Prefix to append to all keys before sending to Redis")]

    key_prefix: String,

    /// The amount of data to read from Redis at a time.

    #[metric(help = "The amount of data to read from Redis at a time")]

    read_chunk_size: usize,

    /// The maximum number of chunk uploads per update.

    /// This is used to limit the number of chunk uploads per update to prevent

    /// overloading when uploading large blocks of data

    #[metric(help = "The maximum number of chunk uploads per update")]

    max_chunk_uploads_per_update: usize,

    /// The COUNT value passed when scanning keys in Redis.

    /// This is used to hint the amount of work that should be done per response.

    #[metric(help = "The COUNT value passed when scanning keys in Redis")]

    scan_count: usize,

    /// The COUNT value used with search indexes

    #[metric(help = "The maximum number of results to return per cursor")]

    max_count_per_cursor: u64,

    /// A manager for subscriptions to keys in Redis.

    subscription_manager: Arc<RedisSubscriptionManager>,

    /// Permits to limit inflight Redis requests. Technically only

    /// limits the calls to `get_client()`, but the requests per client

    /// are small enough that it works well enough.

    client_permits: Arc<Semaphore>,

    /// Per-call ceiling for `check_health` PING.

    health_check_timeout: Duration,

    /// Expire keys this store writes after this long. `None` keeps them

    /// forever, which is the default and what every store other than a BEP

    /// store wants.

    key_ttl: Option<Duration>,

    /// Have we done a subscribe for messages for `remove_callback` subscribes?

    has_remove_callback_subscribe: OnceCell<()>,

    remove_callbacks: Arc<async_lock::Mutex<Vec<RemoveCallback>>>,

}

impl<C, M> Debug for RedisStore<C, M>

where

    C: ConnectionLike + Clone,

    M: RedisManager<C>,

{

    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {

        f.debug_struct("RedisStore")

            .field("temp_name_generator_fn", &self.temp_name_generator_fn)

            .field("key_prefix", &self.key_prefix)

            .field("read_chunk_size", &self.read_chunk_size)

            .field(

                "max_chunk_uploads_per_update",

                &self.max_chunk_uploads_per_update,

            )

            .field("scan_count", &self.scan_count)

            .field("subscription_manager", &self.subscription_manager)

            .field("client_permits", &self.client_permits)

            .finish()

    }

}

struct ClientWithPermit<C: ConnectionLike> {

    connection_manager: C,

    uuid: Uuid,

    // here so it sticks around with the client and doesn't get dropped until that does

    #[allow(dead_code)]

    semaphore_permit: OwnedSemaphorePermit,

}

impl<C: ConnectionLike + Clone> ClientWithPermit<C> {

    async fn reconnect<M: RedisManager<C> + Sync>(&mut self, manager: &M) -> Result<(), Error> {

        (self.connection_manager, self.uuid) = manager.reconnect(self.uuid).await
?0
;

        Ok(())

    }

}

impl<C: ConnectionLike> Drop for ClientWithPermit<C> {

    fn drop(&mut self) {

        trace!(

            remaining = self.semaphore_permit.semaphore().available_permits(),

            "Dropping a client permit"

        );

    }

}

/// Decode a `StoreKey` coming from Redis

pub fn decode_key<'a>(

    key_prefix: &String,

    encoded_key: Cow<'a, str>,

) -> Result<StoreKey<'a>, Error> {

    let key_no_prefix = if key_prefix.is_empty() {

        encoded_key

    } else {

        match encoded_key.strip_prefix(key_prefix) {

            Some(r) => Cow::from(r.to_string()),

            None => {

                return Err(make_err!(

                    Code::InvalidArgument,

                    "Redis key ({}) is missing prefix ({})",

                    encoded_key,

                    key_prefix

                ));

            }

        }

    };

    let maybe_digest_info: Result<_, serde::de::value::Error> =

        DigestInfo::deserialize(key_no_prefix.clone().into_deserializer());

    if let Ok(
digest_info8
) = maybe_digest_info {

        Ok(StoreKey::Digest(digest_info))

    } else {

        Ok(StoreKey::Str(key_no_prefix))

    }

}

impl<C, M> RedisStore<C, M>

where

    C: ConnectionLike + Clone + Sync,

    M: RedisManager<C> + Sync,

{

    /// Used for testing when determinism is required.

    #[expect(clippy::too_many_arguments)]

    pub async fn new_from_builder_and_parts(

        pub_sub_channel: Option<String>,

        temp_name_generator_fn: fn() -> String,

        key_prefix: String,

        read_chunk_size: usize,

        max_chunk_uploads_per_update: usize,

        scan_count: usize,

        max_client_permits: usize,

        max_count_per_cursor: u64,

        health_check_timeout: Duration,

        key_ttl: Option<Duration>,

        subscriber_channel: UnboundedReceiver<PushInfo>,

        connection_manager: M,

    ) -> Result<Self, Error> {

        info!("Redis index fingerprint: {FINGERPRINT_CREATE_INDEX_HEX}");

        let remove_callbacks = Arc::new(async_lock::Mutex::new(Vec::new()));

        let subscription_manager = Arc::new(RedisSubscriptionManager::new(

            subscriber_channel,

            remove_callbacks.clone(),

            key_prefix.clone(),

        ));

        if let Some(
channel5
) = &pub_sub_channel {

            connection_manager.psubscribe(channel).await
?0
;

        }

        Ok(Self {

            connection_manager,

            _connection_type: PhantomData,

            pub_sub_channel,

            temp_name_generator_fn,

            key_prefix,

            read_chunk_size,

            max_chunk_uploads_per_update,

            scan_count,

            subscription_manager,

            client_permits: Arc::new(Semaphore::new(max_client_permits)),

            max_count_per_cursor,

            health_check_timeout,

            key_ttl,

            has_remove_callback_subscribe: OnceCell::const_new(),

            remove_callbacks,

        })

    }

    async fn get_client(&self) -> Result<ClientWithPermit<C>, Error> {

        let local_client_permits = self.client_permits.clone();

        let remaining = local_client_permits.available_permits();

        // Zero here means every client is busy and this call is about to wait,

        // which is the saturation signal worth alerting on.

        record_connection_acquired("redis", Some(remaining), remaining == 0);

        let semaphore_permit = local_client_permits.acquire_owned().await
?0
;

        trace!(remaining, "Got a client permit");

        let (connection_manager, uuid) = self.connection_manager.get_connection().await
?0
;

        Ok(ClientWithPermit {

            connection_manager,

            uuid,

            semaphore_permit,

        })

    }

    /// Encode a `StoreKey` so it can be sent to Redis.

    pub fn encode_key<'a>(&self, key: &'a StoreKey<'a>) -> Cow<'a, str> {

        let key_body = key.as_str();

        if self.key_prefix.is_empty() {

            key_body

        } else {

            // This is in the hot path for all redis operations, so we try to reuse the allocation

            // from `key.as_str()` if possible.

            match key_body {

                Cow::Owned(mut encoded_key) => {

                    encoded_key.insert_str(0, &self.key_prefix);

                    Cow::Owned(encoded_key)

                }

                Cow::Borrowed(body) => {

                    let mut encoded_key = String::with_capacity(self.key_prefix.len() + body.len());

                    encoded_key.push_str(&self.key_prefix);

                    encoded_key.push_str(body);

                    Cow::Owned(encoded_key)

                }

            }

        }

    }

    fn set_spec_defaults(spec: &mut RedisSpec) -> Result<(), Error> {

        if spec.addresses.is_empty() {

            return Err(make_err!(

                Code::InvalidArgument,

                "No addresses were specified in redis store configuration."

            ));

        }

        if spec.broadcast_channel_capacity != 0 {

            warn!("broadcast_channel_capacity in Redis spec is deprecated and ignored");

        }

        if spec.response_timeout_s != 0 {

            warn!(

                "response_timeout_s in Redis spec is deprecated and ignored, use command_timeout_ms"

            );

        }

        if spec.connection_timeout_s != 0 {

            if spec.connection_timeout_ms != 0 {

                return Err(make_err!(

                    Code::InvalidArgument,

                    "Both connection_timeout_s and connection_timeout_ms were set, can only have one!"

                ));

            }

            warn!("connection_timeout_s in Redis spec is deprecated, use connection_timeout_ms");

            spec.connection_timeout_ms = spec.connection_timeout_s * 1000;

        }

        if spec.connection_timeout_ms == 0 {

            spec.connection_timeout_ms = DEFAULT_CONNECTION_TIMEOUT_MS;

        
}5

        if spec.command_timeout_ms == 0 {

            spec.command_timeout_ms = DEFAULT_COMMAND_TIMEOUT_MS;

        
}2

        if spec.health_check_timeout_ms == 0 {

            spec.health_check_timeout_ms = DEFAULT_HEALTH_CHECK_TIMEOUT_MS;

        
}0

        if spec.connection_pool_size == 0 {

            spec.connection_pool_size = DEFAULT_CONNECTION_POOL_SIZE;

        
}0

        if spec.read_chunk_size == 0 {

            spec.read_chunk_size = DEFAULT_READ_CHUNK_SIZE;

        
}0

        if spec.max_count_per_cursor == 0 {

            spec.max_count_per_cursor = DEFAULT_MAX_COUNT_PER_CURSOR;

        
}0

        if spec.max_chunk_uploads_per_update == 0 {

            spec.max_chunk_uploads_per_update = DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE;

        
}0

        if spec.scan_count == 0 {

            spec.scan_count = DEFAULT_SCAN_COUNT;

        
}0

        if spec.max_client_permits == 0 {

            spec.max_client_permits = DEFAULT_CLIENT_PERMITS;

        
}0

        if spec.retry.delay == 0.0 {

            spec.retry.delay = DEFAULT_RETRY_DELAY;

        
}0

        if spec.retry.max_retries == 0 {

            spec.retry.max_retries = 1;

        
}0

        trace!(?spec, "redis spec is after setting defaults");

        Ok(())

    }

    // Only used by tests, because we need to make a real redis connection, then fix this to get fixed values

    pub fn replace_temp_name_generator(&mut self, replacement: fn() -> String) {

        self.temp_name_generator_fn = replacement;

    }

}

impl RedisStore<ClusterConnection, ClusterRedisManager<ClusterConnection>> {

    pub async fn new_cluster(mut spec: RedisSpec) -> Result<Arc<Self>, Error> {

        if spec.mode != RedisMode::Cluster {

            return Err(Error::new(

                Code::InvalidArgument,

                "new_cluster only works for Cluster mode".to_string(),

            ));

        }

        Self::set_spec_defaults(&mut spec)?;

        let parsed_addrs: Vec<_> = spec

            .addresses

            .iter_mut()

            .map(|addr| {

                addr.clone().into_connection_info().map(|connection_info| {

                    let redis_settings = connection_info

                        .redis_settings()

                        .clone()

                        // We need RESP3 here because the cluster mode doesn't support RESP2 pubsub

                        // See also https://docs.rs/redis/latest/redis/cluster_async/index.html#pubsub

                        .set_protocol(redis::ProtocolVersion::RESP3);

                    connection_info.set_redis_settings(redis_settings)

                })

            })

            .collect::<Result<Vec<_>, _>>()?;

        let connection_timeout = Duration::from_millis(spec.connection_timeout_ms);

        let command_timeout = Duration::from_millis(spec.command_timeout_ms);

        let (tx, subscriber_channel) = unbounded_channel();

        let builder = ClusterClient::builder(parsed_addrs)

            .connection_timeout(connection_timeout)

            .response_timeout(command_timeout)

            .push_sender(tx)

            .retries(u32::try_from(spec.retry.max_retries)?);

        let client = builder.build()?;

        Self::new_from_builder_and_parts(

            spec.experimental_pub_sub_channel,

            || Uuid::new_v4().to_string(),

            spec.key_prefix.clone(),

            spec.read_chunk_size,

            spec.max_chunk_uploads_per_update,

            spec.scan_count,

            spec.max_client_permits,

            spec.max_count_per_cursor,

            Duration::from_millis(spec.health_check_timeout_ms),

            (spec.key_ttl_s > 0).then(|| Duration::from_secs(spec.key_ttl_s)),

            subscriber_channel,

            ClusterRedisManager::new(client.get_async_connection().await?).await?,

        )

        .await

        .map(Arc::new)

    }

}

impl RedisStore<ConnectionManager, StandardRedisManager<ConnectionManager>> {

    async fn connect(

        spec: RedisSpec,

        tx: UnboundedSender<PushInfo>,

    ) -> Result<ConnectionManager, Error> {

        let connection_timeout = Duration::from_millis(spec.connection_timeout_ms);

        let command_timeout = Duration::from_millis(spec.command_timeout_ms);

        let addr = &spec.addresses[0];

        let local_addr = addr.clone();

        let mut parsed_addr = local_addr

            .replace("redis+sentinel://", "redis://")

            .into_connection_info()
?0
;

        let redis_settings = parsed_addr

            .redis_settings()

            .clone()

            // We need RESP3 here because we want to do set_push_sender

            .set_protocol(redis::ProtocolVersion::RESP3);

        parsed_addr = parsed_addr.set_redis_settings(redis_settings);

        debug!(?parsed_addr, "Parsed redis addr");

        let 
client26
 = timeout(

            connection_timeout,

            spawn!("connect", async move {

                match spec.mode {

                    RedisMode::Standard => Client::open(parsed_addr).map_err(Into::<Error>::into),

                    RedisMode::Cluster => {

                        return Err(Error::new(

                            Code::Internal,

                            "Use RedisStore::new_cluster for cluster connections".to_owned(),

                        ));

                    }

                    RedisMode::Sentinel => async {

                        let url_parsing = Url::parse(&local_addr)
?0
;

                        let master_name = url_parsing

                            .query_pairs()

                            .find(|(key, _)| 
key1
 == 
"sentinelServiceName"1
)

                            .map_or_else(|| 
"master"8
.
into8
(), |(_, value)| 
value1
.
to_string1
());

                        let redis_connection_info = parsed_addr.redis_settings().clone();

                        let sentinel_connection_info = SentinelNodeConnectionInfo::default()

                            .set_redis_connection_info(redis_connection_info);

                        // We fish this out because sentinels don't support db, we need to set it

                        // on the client only. See also https://github.com/redis-rs/redis-rs/issues/1950

                        let original_db = parsed_addr.redis_settings().db();

                        if original_db != 0 {

                            // sentinel_connection_info has the actual DB set

                            let revised_settings = parsed_addr.redis_settings().clone().set_db(0);

                            parsed_addr = parsed_addr.set_redis_settings(revised_settings);

                        }

                        SentinelClient::build(

                            vec![parsed_addr],

                            master_name,

                            Some(sentinel_connection_info),

                            SentinelServerType::Master,

                        )

                        .map_err(Into::<Error>::into)

                    }

                    .
and_then9
(|mut s| async move 
{9
 Ok(
s9
.async_get_client().await) })

                    .await
?0

                    .map_err(Into::<Error>::into),

                }

                // Keep the underlying error's code rather than forcing

                // InvalidArgument. Not every connect failure is a bad URL: a

                // sentinel failover in progress reports the master as missing,

                // which is transient and retryable, and overriding it to a

                // permanent status made clients give up on a recoverable blip.

                .err_tip(|| 
format!1
("While connecting to redis with url: {local_addr}"))

            }),

        )

        .await

        .err_tip(|| 
format!0
("Timeout while connecting to redis with url: {addr}"))
?0
?0
?1
;

        let connection_manager_config = {

            ConnectionManagerConfig::new()

                .set_number_of_retries(spec.retry.max_retries)

                .set_connection_timeout(Some(connection_timeout))

                .set_response_timeout(Some(command_timeout))

                .set_push_sender(tx)

        };

        let mut connection_manager =

            ConnectionManager::new_with_config(client, connection_manager_config)

                .await

                .err_tip(|| 
format!2
("While connecting to redis with url: {addr}"))
?2
;

        if let Some(
pub_sub_channel5
) = spec.experimental_pub_sub_channel {

            connection_manager.psubscribe(pub_sub_channel).await
?0
;

        }

        Ok(connection_manager)

    }

    /// Create a new `RedisStore` from the given configuration.

    pub async fn new_standard(mut spec: RedisSpec) -> Result<Arc<Self>, Error> {

        Self::set_spec_defaults(&mut spec)
?0
;

        if spec.addresses.len() != 1 {

            return Err(make_err!(

                Code::Unimplemented,

                "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."

            ));

        }

        let (tx, subscriber_channel) = unbounded_channel();

        Self::new_from_builder_and_parts(

            spec.experimental_pub_sub_channel.clone(),

            || Uuid::new_v4().to_string(),

            spec.key_prefix.clone(),

            spec.read_chunk_size,

            spec.max_chunk_uploads_per_update,

            spec.scan_count,

            spec.max_client_permits,

            spec.max_count_per_cursor,

            Duration::from_millis(spec.health_check_timeout_ms),

            (spec.key_ttl_s > 0).then(|| 
Duration::from_secs0
(
spec.key_ttl_s0
)),

            subscriber_channel,

            
StandardRedisManager::new24
(
Box::new24
(move || {

                Box::pin(Self::connect(spec.clone(), tx.clone()))

            }))

            .await
?3
,

        )

        .await

        .map(Arc::new)

    }

}

#[async_trait]

impl<C, M> StoreDriver for RedisStore<C, M>

where

    C: ConnectionLike + Clone + Send + Sync + Unpin + 'static,

    M: RedisManager<C> + Unpin + Send + Sync + 'static,

{

    async fn post_init(self: Arc<Self>) -> Result<(), Error> {

        Ok(())

    }

    async fn has_with_results(

        self: Pin<&Self>,

        keys: &[StoreKey<'_>],

        results: &mut [Option<u64>],

    ) -> Result<(), Error> {

        // TODO(palfrey) We could use pipeline here, but it makes retry more

        // difficult and it doesn't work very well in cluster mode.

        // If we wanted to optimize this with pipeline be careful to

        // implement retry and to support cluster mode.

        izip!(keys.iter(), results.iter_mut(),)

            .map(|(key, result)| async move {

                // We need to do a special pass to ensure our zero key exist.

                if is_zero_digest(key.borrow()) {

                    *result = Some(0);

                    return Ok::<_, Error>(());

                }

                let encoded_key = self.encode_key(key);

                let mut client = self.get_client().await
?0
;

                // Redis returns 0 when the key doesn't exist

                // AND when the key exists with value of length 0.

                // Therefore, we need to check both length and existence

                // and do it in a pipeline for efficiency

                // Re-resolve the master and retry on a transient failover so a

                // topology change doesn't fail the existence check.

                let (blob_len, exists) = {

                    let mut attempt: u32 = 0;

                    loop {

                        attempt += 1;

                        match pipe()

                            .strlen(encoded_key.as_ref())

                            .exists(encoded_key.as_ref())

                            .query_async::<(u64, bool)>(&mut client.connection_manager)

                            .await

                        {

                            Ok(v) => break v,

                            Err(err)

                                if attempt < MAX_REDIS_RETRY_ATTEMPTS

                                    && is_retryable_redis_error(&err) =>

                            {

                                client.reconnect(&self.connection_manager).await
?0
;

                                sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;

                            }

                            Err(err) => {

                                return Err(

                                    Error::from(err).append("In RedisStore::has_with_results::all")

                                );

                            }

                        }

                    }

                };

                *result = if exists { Some(blob_len) } else { 
None0
 };

                Ok::<_, Error>(())

            })

            .collect::<FuturesUnordered<_>>()

            .try_collect()

            .await

    }

    async fn list(

        self: Pin<&Self>,

        range: (Bound<StoreKey<'_>>, Bound<StoreKey<'_>>),

        handler: &mut (dyn for<'a> FnMut(&'a StoreKey) -> bool + Send + Sync + '_),

    ) -> Result<u64, Error> {

        let range = (

            range.0.map(StoreKey::into_owned),

            range.1.map(StoreKey::into_owned),

        );

        let pattern = match range.0 {

            Bound::Included(ref start) | Bound::Excluded(ref start) => match range.1 {

                Bound::Included(ref end) | Bound::Excluded(ref end) => {

                    let start = start.as_str();

                    let end = end.as_str();

                    let max_length = start.len().min(end.len());

                    let length = start

                        .chars()

                        .zip(end.chars())

                        .position(|(a, b)| a != b)

                        .unwrap_or(max_length);

                    format!("{}{}*", self.key_prefix, &start[..length])

                }

                Bound::Unbounded => format!("{}*", self.key_prefix),

            },

            Bound::Unbounded => format!("{}*", self.key_prefix),

        };

        let mut client = self.get_client().await?;

        trace!(%pattern, count=self.scan_count, "Running SCAN");

        // Restart the scan on a transient failover. Redis SCAN may re-emit keys

        // even without failures, so callers already tolerate duplicates;

        // re-resolving the master and rescanning is within that contract and

        // avoids failing on a topology change mid-iteration.

        let mut attempt: u32 = 0;

        loop {

            attempt += 1;

            let opts = ScanOptions::default()

                .with_pattern(pattern.clone())

                .with_count(self.scan_count);

            // Scan via a cloned connection handle so the iterator doesn't hold a

            // borrow on `client` across a reconnect.

            let mut conn = client.connection_manager.clone();

            let mut scan_stream: AsyncIter<Value> = match conn.scan_options(opts).await {

                Ok(s) => s,

                Err(err)

                    if attempt < MAX_REDIS_RETRY_ATTEMPTS && is_retryable_redis_error(&err) =>

                {

                    client.reconnect(&self.connection_manager).await?;

                    sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;

                    continue;

                }

                Err(err) => return Err(Error::from(err).append("During scan_options")),

            };

            let mut iterations = 0;

            let mut errors = vec![];

            let mut transient_err = false;

            while let Some(key) = scan_stream.next_item().await {

                match key {

                    Ok(Value::BulkString(raw_key)) => {

                        let Ok(str_key) = str::from_utf8(&raw_key) else {

                            error!(?raw_key, "Non-utf8 key");

                            errors.push(format!("Non-utf8 key {raw_key:?}"));

                            continue;

                        };

                        match decode_key(&self.key_prefix, Cow::from(str_key)) {

                            Ok(key) => {

                                if range.contains(&key) {

                                    iterations += 1;

                                    if !handler(&key) {

                                        error!("Issue in handler");

                                        errors.push("Issue in handler".to_string());

                                    }

                                } else {

                                    trace!(%key, ?range, "Key not in range");

                                }

                            }

                            Err(e) => {

                                errors.push(e.to_string());

                            }

                        }

                    }

                    Err(err)

                        if attempt < MAX_REDIS_RETRY_ATTEMPTS && is_retryable_redis_error(&err) =>

                    {

                        // Connection dropped mid-scan; restart from scratch.

                        transient_err = true;

                        break;

                    }

                    other => {

                        error!(?other, "Non-string in key");

                        errors.push("Non-string in key".to_string());

                    }

                }

            }

            if transient_err {

                drop(scan_stream);

                client.reconnect(&self.connection_manager).await?;

                sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await;

                continue;

            }

            return if errors.is_empty() {

                Ok(iterations)

            } else {

                error!(?errors, "Errors in scan stream");

                Err(Error::new(Code::Internal, format!("Errors: {errors:?}")))