// Copyright 2024 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::ops::{Bound, RangeBounds};

use core::pin::Pin;

use core::time::Duration;

use core::{cmp, iter};

use std::borrow::Cow;

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

use async_trait::async_trait;

use bytes::Bytes;

use const_format::formatcp;

use fred::clients::SubscriberClient;

use fred::interfaces::{ClientLike, KeysInterface, PubsubInterface};

use fred::prelude::{Client, EventInterface, HashesInterface, RediSearchInterface};

use fred::types::config::{

    Config as RedisConfig, ConnectionConfig, PerformanceConfig, ReconnectPolicy, UnresponsiveConfig,

};

use fred::types::redisearch::{

    AggregateOperation, FtAggregateOptions, FtCreateOptions, IndexKind, Load, SearchField,

    SearchSchema, SearchSchemaKind, WithCursor,

};

use fred::types::scan::Scanner;

use fred::types::scripts::Script;

use fred::types::{Builder, Key as RedisKey, Map as RedisMap, SortOrder, Value as RedisValue};

use futures::stream::FuturesUnordered;

use futures::{FutureExt, Stream, StreamExt, 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_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};

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

use nativelink_util::spawn;

use nativelink_util::store_trait::{

    BoolValue, RemoveItemCallback, SchedulerCurrentVersionProvider, SchedulerIndexProvider,

    SchedulerStore, SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider,

    SchedulerSubscription, SchedulerSubscriptionManager, StoreDriver, StoreKey, UploadSizeInfo,

};

use nativelink_util::task::JoinHandleDropGuard;

use parking_lot::{Mutex, RwLock};

use patricia_tree::StringPatriciaMap;

use tokio::select;

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

use tokio::time::sleep;

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

use uuid::Uuid;

use crate::cas_utils::is_zero_digest;

use crate::redis_utils::ft_aggregate;

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

/// The amount of jitter to add to the retry delay if not specified.

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

const DEFAULT_RETRY_JITTER: f32 = 0.5;

/// The default maximum capacity of the broadcast channel if not specified.

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

const DEFAULT_BROADCAST_CHANNEL_CAPACITY: usize = 4096;

/// 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 maximum number of chunk uploads per update.

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

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: u32 = 10_000;

const DEFAULT_CLIENT_PERMITS: usize = 500;

#[derive(Clone, Debug)]

pub struct RecoverablePool {

    clients: Arc<RwLock<Vec<Client>>>,

    builder: Builder,

    counter: Arc<core::sync::atomic::AtomicUsize>,

}

impl RecoverablePool {

    pub fn new(builder: Builder, size: usize) -> Result<Self, Error> {

        let mut clients = Vec::with_capacity(size);

        for _ in 0..size {

            let client = builder

                .build()

                .err_tip(|| "Failed to build client in RecoverablePool::new")
?0
;

            clients.push(client);

        }

        Ok(Self {

            clients: Arc::new(RwLock::new(clients)),

            builder,

            counter: Arc::new(core::sync::atomic::AtomicUsize::new(0)),

        })

    }

    fn connect(&self) {

        let clients = self.clients.read();

        for client in 
clients14
.
iter14
() {

            client.connect();

        }

    }

    fn next(&self) -> Client {

        let clients = self.clients.read();

        let index = self

            .counter

            .fetch_add(1, core::sync::atomic::Ordering::Relaxed);

        clients[index % clients.len()].clone()

    }

    async fn replace_client(&self, old_client: &Client) -> Result<Client, Error> {

        {

            let clients = self.clients.read();

            if !clients.iter().any(|c| c.id() == old_client.id()) {

  Branch (142:16): [True: 0, False: 0]

  Branch (142:16): [True: 0, False: 2]

  Branch (142:16): [Folded - Ignored]

  Branch (142:16): [True: 0, False: 0]

                // Someone else swapped this client already; just hand out the next pooled one.

                return Ok(self.next());

            }

        }

        let new_client = self

            .builder

            .build()

            .err_tip(|| "Failed to build new client in RecoverablePool::replace_client")
?0
;

        new_client.connect();

        new_client.wait_for_connect().await.err_tip(|| {

            format!(

                "Failed to connect new client while replacing Redis client {}",

                old_client.id()

            )

        })?;

        let replaced_client = {

            let mut clients = self.clients.write();

            clients

                .iter()

                .position(|c| c.id() == old_client.id())

                .map(|index| core::mem::replace(&mut clients[index], new_client.clone()))

        };

        if let Some(old_client) = replaced_client {

  Branch (168:16): [True: 0, False: 0]

  Branch (168:16): [True: 0, False: 0]

  Branch (168:16): [Folded - Ignored]

  Branch (168:16): [True: 0, False: 0]

            let _unused = old_client.quit().await;

            info!("Replaced Redis client {}", old_client.id());

            Ok(new_client)

        } else {

            // Second race: pool entry changed after we connected the new client.

            let _unused = new_client.quit().await;

            Ok(self.next())

        }

    }

}

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

#[derive(Debug, MetricsComponent)]

pub struct RedisStore {

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

    client_pool: RecoverablePool,

    /// 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 redis client for managing subscriptions.

    /// TODO: This should be moved into the store in followups once a standard use pattern has been determined.

    subscriber_client: SubscriberClient,

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

    #[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: u32,

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

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

    subscription_manager: Mutex<Option<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>,

}

struct ClientWithPermit {

    client: Client,

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

    #[allow(dead_code)]

    semaphore_permit: OwnedSemaphorePermit,

}

impl Drop for ClientWithPermit {

    fn drop(&mut self) {

        trace!(

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

            "Dropping a client permit"

        );

    }

}

impl RedisStore {

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

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

        if spec.addresses.is_empty() {

  Branch (253:12): [True: 0, False: 2]

  Branch (253:12): [Folded - Ignored]

            return Err(make_err!(

                Code::InvalidArgument,

                "No addresses were specified in redis store configuration."

            ));

        }

        let [addr] = spec.addresses.as_slice() else {

  Branch (259:13): [True: 2, False: 0]

  Branch (259:13): [Folded - Ignored]

            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 redis_config = match spec.mode {

            RedisMode::Cluster => RedisConfig::from_url_clustered(addr),

            RedisMode::Sentinel => RedisConfig::from_url_sentinel(addr),

            RedisMode::Standard => RedisConfig::from_url_centralized(addr),

        }

        .err_tip_with_code(|e| 
{0

            (

                Code::InvalidArgument,

                format!("while parsing redis node address: {e}"),

            )

        })?;

        let reconnect_policy = {

            if spec.retry.delay == 0.0 {

  Branch (278:16): [True: 2, False: 0]

  Branch (278:16): [Folded - Ignored]

                spec.retry.delay = DEFAULT_RETRY_DELAY;

            
}0

            if spec.retry.jitter == 0.0 {

  Branch (281:16): [True: 2, False: 0]

  Branch (281:16): [Folded - Ignored]

                spec.retry.jitter = DEFAULT_RETRY_JITTER;

            
}0

            let 
to_ms2
 = |secs: f32| -> u32 {

                Duration::from_secs_f32(secs)

                    .as_millis()

                    .try_into()

                    .unwrap_or(u32::MAX)

            };

            let max_retries = u32::try_from(spec.retry.max_retries)

                .err_tip(|| "max_retries could not be converted to u32 in RedisStore::new")
?0
;

            let min_delay_ms = to_ms(spec.retry.delay);

            let max_delay_ms = 8000;

            let jitter = to_ms(spec.retry.jitter * spec.retry.delay);

            let mut reconnect_policy =

                ReconnectPolicy::new_exponential(max_retries, min_delay_ms, max_delay_ms, 2);

            reconnect_policy.set_jitter(jitter);

            reconnect_policy

        };

        {

            if spec.broadcast_channel_capacity == 0 {

  Branch (306:16): [True: 2, False: 0]

  Branch (306:16): [Folded - Ignored]

                spec.broadcast_channel_capacity = DEFAULT_BROADCAST_CHANNEL_CAPACITY;

            
}0

            if spec.connection_timeout_ms == 0 {

  Branch (309:16): [True: 2, False: 0]

  Branch (309:16): [Folded - Ignored]

                spec.connection_timeout_ms = DEFAULT_CONNECTION_TIMEOUT_MS;

            
}0

            if spec.command_timeout_ms == 0 {

  Branch (312:16): [True: 2, False: 0]

  Branch (312:16): [Folded - Ignored]

                spec.command_timeout_ms = DEFAULT_COMMAND_TIMEOUT_MS;

            
}0

            if spec.connection_pool_size == 0 {

  Branch (315:16): [True: 2, False: 0]

  Branch (315:16): [Folded - Ignored]

                spec.connection_pool_size = DEFAULT_CONNECTION_POOL_SIZE;

            
}0

            if spec.read_chunk_size == 0 {

  Branch (318:16): [True: 2, False: 0]

  Branch (318:16): [Folded - Ignored]

                spec.read_chunk_size = DEFAULT_READ_CHUNK_SIZE;

            
}0

            if spec.max_chunk_uploads_per_update == 0 {

  Branch (321:16): [True: 2, False: 0]

  Branch (321:16): [Folded - Ignored]

                spec.max_chunk_uploads_per_update = DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE;

            
}0

            if spec.scan_count == 0 {

  Branch (324:16): [True: 2, False: 0]

  Branch (324:16): [Folded - Ignored]

                spec.scan_count = DEFAULT_SCAN_COUNT;

            
}0

            if spec.max_client_permits == 0 {

  Branch (327:16): [True: 2, False: 0]

  Branch (327:16): [Folded - Ignored]

                spec.max_client_permits = DEFAULT_CLIENT_PERMITS;

            
}0

        }

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

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

        let mut builder = Builder::from_config(redis_config);

        builder

            .set_performance_config(PerformanceConfig {

                default_command_timeout: command_timeout,

                broadcast_channel_capacity: spec.broadcast_channel_capacity,

                ..Default::default()

            })

            .set_connection_config(ConnectionConfig {

                connection_timeout,

                internal_command_timeout: command_timeout,

                unresponsive: UnresponsiveConfig {

                    max_timeout: Some(connection_timeout),

                    // This number needs to be less than the connection timeout.

                    // We use 4 as it is a good balance between not spamming the server

                    // and not waiting too long.

                    interval: connection_timeout / 4,

                },

                ..Default::default()

            })

            .set_policy(reconnect_policy);

        let client_pool = RecoverablePool::new(builder.clone(), spec.connection_pool_size)

            .err_tip(|| "while creating redis connection pool")
?0
;

        let subscriber_client = builder

            .build_subscriber_client()

            .err_tip(|| "while creating redis subscriber client")
?0
;

        Self::new_from_builder_and_parts(

            client_pool,

            subscriber_client,

            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,

        )

        .map(Arc::new)

    }

    /// Used for testing when determinism is required.

    #[expect(clippy::too_many_arguments)]

    pub fn new_from_builder_and_parts(

        client_pool: RecoverablePool,

        subscriber_client: SubscriberClient,

        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: u32,

        max_client_permits: usize,

    ) -> Result<Self, Error> {

        // Start connection pool (this will retry forever by default).

        client_pool.connect();

        subscriber_client.connect();

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

        Ok(Self {

            client_pool,

            pub_sub_channel,

            subscriber_client,

            temp_name_generator_fn,

            key_prefix,

            read_chunk_size,

            max_chunk_uploads_per_update,

            scan_count,

            update_if_version_matches_script: Script::from_lua(LUA_VERSION_SET_SCRIPT),

            subscription_manager: Mutex::new(None),

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

        })

    }

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

        let mut client = self.client_pool.next();

        loop {

            let config = client.client_config();

            if config.mocks.is_some() {

  Branch (414:16): [True: 0, False: 0]

  Branch (414:16): [True: 24, False: 2]

  Branch (414:16): [Folded - Ignored]

  Branch (414:16): [True: 96, False: 0]

                break;

            }

            let connection_info = format!(

                "Connection issue connecting to redis server with hosts: {:?}, username: {}, database: {}",

                config

                    .server

                    .hosts()

                    .iter()

                    .map(|s| format!("{}:{}", s.host, s.port))

                    .collect::<Vec<String>>(),

                config

                    .username

                    .clone()

                    .unwrap_or_else(|| "None".to_string()),

                config.database.unwrap_or_default()

            );

            match client.wait_for_connect().await {

                Ok(()) => break,

                Err(e) => {

                    warn!("{connection_info}: {e:?}. Replacing client.");

                    
client0
 = self

                        .client_pool

                        .replace_client(&client)

                        .await

                        .err_tip(|| connection_info.clone())?;

                }

            }

        }

        let local_client_permits = self.client_permits.clone();

        let remaining = local_client_permits.available_permits();

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

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

        Ok(ClientWithPermit {

            client,

            semaphore_permit,

        })

    }

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

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

        let key_body = key.as_str();

        if self.key_prefix.is_empty() {

  Branch (456:12): [True: 89, False: 4]

  Branch (456:12): [Folded - Ignored]

            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)

                }

            }

        }

    }

}

#[async_trait]

impl StoreDriver for RedisStore {

    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.

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

        // If we ask for many keys in one go, this can timeout, so limit that

        let max_in_one_go = Arc::new(Semaphore::const_new(5));

        izip!(

            keys.iter(),

            results.iter_mut(),

            iter::repeat(&max_in_one_go),

            iter::repeat(&client)

        )

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

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

            if is_zero_digest(key.borrow()) {

  Branch (502:16): [True: 2, False: 4]

  Branch (502:16): [Folded - Ignored]

                *result = Some(0);

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

            }

            let encoded_key = self.encode_key(key);

            let guard = local_semaphore.acquire().await
?0
;

            let pipeline = client.client.pipeline();

            pipeline

                .strlen::<(), _>(encoded_key.as_ref())

                .await

                .err_tip(|| format!(
"In RedisStore::has_with_results::strlen for {encoded_key}"0
))
?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.

            pipeline

                .exists::<(), _>(encoded_key.as_ref())

                .await

                .err_tip(|| format!(
"In RedisStore::has_with_results::exists for {encoded_key}"0
))
?0
;

            let (blob_len, exists) = pipeline

                .all::<(u64, bool)>()

                .await

                .err_tip(|| "In RedisStore::has_with_results::all")
?0
;

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

  Branch (528:26): [True: 4, False: 0]

  Branch (528:26): [Folded - Ignored]

            drop(guard);

            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 client = self.get_client().await?;

        let mut scan_stream = client.client.scan(pattern, Some(self.scan_count), None);

        let mut iterations = 0;

        'outer: while let Some(mut page) = scan_stream.try_next().await? {

            if let Some(keys) = page.take_results() {

                for key in keys {

                    // TODO: Notification of conversion errors

                    // Any results that do not conform to expectations are ignored.

                    if let Some(key) = key.as_str() {

                        if let Some(key) = key.strip_prefix(&self.key_prefix) {

                            let key = StoreKey::new_str(key);

                            if range.contains(&key) {

                                iterations += 1;

                                if !handler(&key) {

                                    break 'outer;

                                }

                            }

                        }

                    }

                }

            }

            page.next();

        }

        Ok(iterations)

    }

    async fn update(

        self: Pin<&Self>,

        key: StoreKey<'_>,

        mut reader: DropCloserReadHalf,

        _upload_size: UploadSizeInfo,

    ) -> Result<(), Error> {

        let final_key = self.encode_key(&key);

        // While the name generation function can be supplied by the user, we need to have the curly

        // braces in place in order to manage redis' hashing behavior and make sure that the temporary

        // key name and the final key name are directed to the same cluster node. See

        // https://redis.io/blog/redis-clustering-best-practices-with-keys/

        //

        // The TL;DR is that if we're in cluster mode and the names hash differently, we can't use request

        // pipelining. By using these braces, we tell redis to only hash the part of the temporary key that's

        // identical to the final key -- so they will always hash to the same node.

        let temp_key = format!(

            "temp-{}-{{{}}}",

            (self.temp_name_generator_fn)(),

            &final_key

        );

        if is_zero_digest(key.borrow()) {

            let chunk = reader

                .peek()

                .await

                .err_tip(|| "Failed to peek in RedisStore::update")?;

            if chunk.is_empty() {

                reader

                    .drain()

                    .await

                    .err_tip(|| "Failed to drain in RedisStore::update")?;

                // Zero-digest keys are special -- we don't need to do anything with it.

                return Ok(());

            }

        }

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

        let mut read_stream = reader

            .scan(0u32, |bytes_read, chunk_res| {

                future::ready(Some(

                    chunk_res

                        .err_tip(|| "Failed to read chunk in update in redis store")

                        .and_then(|chunk| 
{5

                            let offset = *bytes_read;

                            let chunk_len = u32::try_from(chunk.len()).err_tip(

                                || "Could not convert chunk length to u32 in RedisStore::update",

                            )?;

                            let new_bytes_read = bytes_read

                                .checked_add(chunk_len)

                                .err_tip(|| "Overflow protection in RedisStore::update")
?0
;

                            *bytes_read = new_bytes_read;

                            Ok::<_, Error>((offset, *bytes_read, chunk))

                        }),

                ))

            })

            .map(|res| {

                let (
offset5
, 
end_pos5
, 
chunk5
) = res
?1
;

                let temp_key_ref = &temp_key;

                let client = client.client.clone();

                Ok(async move {

                    client

                        .setrange::<(), _, _>(temp_key_ref, offset, chunk)

                        .await

                        .err_tip(

                            || format!("While appending to temp key ({temp_key_ref}) in RedisStore::update. offset = {offset}. end_pos = {end_pos}"),

                        )?;

                    Ok::<u32, Error>(end_pos)

                })

            })

            .try_buffer_unordered(self.max_chunk_uploads_per_update);

        let mut total_len: u32 = 0;

        while let Some(last_pos) = read_stream.try_next().await? {

            if last_pos > total_len {

                total_len = last_pos;

            }

        }

        let blob_len = client

            .client

            .strlen::<u64, _>(&temp_key)

            .await

            .err_tip(|| format!("In RedisStore::update strlen check for {temp_key}"))?;

        // This is a safety check to ensure that in the event some kind of retry was to happen

        // and the data was appended to the key twice, we reject the data.

        if blob_len != u64::from(total_len) {

            return Err(make_input_err!(

                "Data length mismatch in RedisStore::update for {}({}) - expected {} bytes, got {} bytes",

                key.borrow().as_str(),

                temp_key,

                total_len,

                blob_len,

            ));

        }

        // Rename the temp key so that the data appears under the real key. Any data already present in the real key is lost.

        client

            .client

            .rename::<(), _, _>(&temp_key, final_key.as_ref())

            .await

            .err_tip(|| "While queueing key rename in RedisStore::update()")?;

        // If we have a publish channel configured, send a notice that the key has been set.

        if let Some(pub_sub_channel) = &self.pub_sub_channel {

            return Ok(client

                .client

                .publish(pub_sub_channel, final_key.as_ref())

                .await?);

        }

        Ok(())

    }

    async fn get_part(

        self: Pin<&Self>,

        key: StoreKey<'_>,

        writer: &mut DropCloserWriteHalf,

        offset: u64,

        length: Option<u64>,

    ) -> Result<(), Error> {

        let offset = usize::try_from(offset).err_tip(|| "Could not convert offset to usize")?;

        let length = length

            .map(|v| usize::try_from(v).err_tip(|| "Could not convert length to usize"))

            .transpose()?;

        // To follow RBE spec we need to consider any digest's with

        // zero size to be existing.

        if is_zero_digest(key.borrow()) {

            return writer

                .send_eof()

                .err_tip(|| "Failed to send zero EOF in redis store get_part");

        }

        let encoded_key = self.encode_key(&key);

        let encoded_key = encoded_key.as_ref();

        // N.B. the `-1`'s you see here are because redis GETRANGE is inclusive at both the start and end, so when we

        // do math with indices we change them to be exclusive at the end.

        // We want to read the data at the key from `offset` to `offset + length`.

        let data_start = offset;

        let data_end = data_start

            .saturating_add(length.unwrap_or(isize::MAX as usize))

            .saturating_sub(1);

        // And we don't ever want to read more than `read_chunk_size` bytes at a time, so we'll need to iterate.

        let mut chunk_start = data_start;

        let mut chunk_end = cmp::min(

            data_start.saturating_add(self.read_chunk_size) - 1,

            data_end,

        );

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

        loop {

            let chunk: Bytes = client

                .client

                .getrange(encoded_key, chunk_start, chunk_end)

                .await

                .err_tip(|| "In RedisStore::get_part::getrange")?;

            let didnt_receive_full_chunk = chunk.len() < self.read_chunk_size;

            let reached_end_of_data = chunk_end == data_end;

            if didnt_receive_full_chunk || reached_end_of_data {

                if !chunk.is_empty() {

                    writer

                        .send(chunk)

                        .await

                        .err_tip(|| "Failed to write data in RedisStore::get_part")?;

                }

                break; // No more data to read.

            }

            // We received a full chunk's worth of data, so write it...

            writer

                .send(chunk)

                .await

                .err_tip(|| "Failed to write data in RedisStore::get_part")?;

            // ...and go grab the next chunk.

            chunk_start = chunk_end + 1;

            chunk_end = cmp::min(

                chunk_start.saturating_add(self.read_chunk_size) - 1,

                data_end,

            );

        }

        // If we didn't write any data, check if the key exists, if not return a NotFound error.

        // This is required by spec.

        if writer.get_bytes_written() == 0 {

            // We're supposed to read 0 bytes, so just check if the key exists.

            let exists = client

                .client

                .exists::<bool, _>(encoded_key)

                .await

                .err_tip(|| "In RedisStore::get_part::zero_exists")?;

            if !exists {

                return Err(make_err!(

                    Code::NotFound,

                    "Data not found in Redis store for digest: {key:?}"

                ));

            }

        }

        writer

            .send_eof()

            .err_tip(|| "Failed to write EOF in redis store get_part")

    }

    fn inner_store(&self, _digest: Option<StoreKey>) -> &dyn StoreDriver {

        self

    }

    fn as_any(&self) -> &(dyn core::any::Any + Sync + Send) {

        self

    }

    fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send> {

        self

    }

    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {

        registry.register_indicator(self);

    }

    fn register_remove_callback(

        self: Arc<Self>,

        _callback: Arc<dyn RemoveItemCallback>,

    ) -> Result<(), Error> {

        // As redis doesn't drop stuff, we can just ignore this

        Ok(())

    }

}

#[async_trait]

impl HealthStatusIndicator for RedisStore {

    fn get_name(&self) -> &'static str {

        "RedisStore"

    }

    async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus {

        StoreDriver::check_health(Pin::new(self), namespace).await

    }

}

// -------------------------------------------------------------------

// Below this line are specific to the redis scheduler implementation.

// -------------------------------------------------------------------

/// The maximum number of results to return per cursor.

const MAX_COUNT_PER_CURSOR: u64 = 256;

/// The time in milliseconds that a redis cursor can be idle before it is closed.

const CURSOR_IDLE_MS: u64 = 2_000;

/// The name of the field in the Redis hash that stores the data.

const DATA_FIELD_NAME: &str = "data";

/// The name of the field in the Redis hash that stores the version.

const VERSION_FIELD_NAME: &str = "version";

/// The time to live of indexes in seconds. After this time redis may delete the index.

const INDEX_TTL_S: u64 = 60 * 60 * 24; // 24 hours.

/// Lua script to set a key if the version matches.

/// Args:

///   KEYS[1]: The key where the version is stored.

///   ARGV[1]: The expected version.

///   ARGV[2]: The new data.

///   ARGV[3*]: Key-value pairs of additional data to include.

/// Returns:

///   The new version if the version matches. nil is returned if the

///   value was not set.

const LUA_VERSION_SET_SCRIPT: &str = formatcp!(

    r"

local key = KEYS[1]

local expected_version = tonumber(ARGV[1])

local new_data = ARGV[2]

local new_version = redis.call('HINCRBY', key, '{VERSION_FIELD_NAME}', 1)

local i

local indexes = {{}}

if new_version-1 ~= expected_version then

    redis.call('HINCRBY', key, '{VERSION_FIELD_NAME}', -1)

    return {{ 0, new_version-1 }}

end

-- Skip first 2 argvs, as they are known inputs.

-- Remember: Lua is 1-indexed.

for i=3, #ARGV do

    indexes[i-2] = ARGV[i]

end

-- In testing we witnessed redis sometimes not update our FT indexes

-- resulting in stale data. It appears if we delete our keys then insert

-- them again it works and reduces risk significantly.

redis.call('DEL', key)

redis.call('HSET', key, '{DATA_FIELD_NAME}', new_data, '{VERSION_FIELD_NAME}', new_version, unpack(indexes))

return {{ 1, new_version }}

"

);

/// This is the output of the calculations below hardcoded into the executable.

const FINGERPRINT_CREATE_INDEX_HEX: &str = "3e762c15";

#[cfg(test)]

mod test {

    use super::FINGERPRINT_CREATE_INDEX_HEX;

    /// String of the `FT.CREATE` command used to create the index template.

    const CREATE_INDEX_TEMPLATE: &str = "FT.CREATE {} ON HASH PREFIX 1 {} NOOFFSETS NOHL NOFIELDS NOFREQS SCHEMA {} TAG CASESENSITIVE SORTABLE";

    /// Compile-time fingerprint of the `FT.CREATE` command used to create the

    /// index template. This is a simple CRC32 checksum of the command string.

    /// We don't care about it actually being a valid CRC32 checksum, just that

    /// it's a unique identifier with a low chance of collision.

    const fn fingerprint_create_index_template() -> u32 {

        const POLY: u32 = 0xEDB8_8320;

        const DATA: &[u8] = CREATE_INDEX_TEMPLATE.as_bytes();

        let mut crc = 0xFFFF_FFFF;

        let mut i = 0;

        while i < DATA.len() {

  Branch (912:15): [True: 101, False: 1]

            let byte = DATA[i];

            crc ^= byte as u32;

            let mut j = 0;

            while j < 8 {

  Branch (917:19): [True: 808, False: 101]

                crc = if crc & 1 != 0 {

  Branch (918:26): [True: 386, False: 422]

                    (crc >> 1) ^ POLY

                } else {

                    crc >> 1

                };

                j += 1;

            }

            i += 1;

        }

        crc

    }

    /// Verify that our calculation always evaluates to this fixed value.

    #[test]

    fn test_fingerprint_value() {

        assert_eq!(

            format!("{:08x}", &fingerprint_create_index_template()),

            FINGERPRINT_CREATE_INDEX_HEX,

        );

    }

}

/// Get the name of the index to create for the given field.

/// This will add some prefix data to the name to try and ensure

/// if the index definition changes, the name will get a new name.

macro_rules! get_index_name {

    ($prefix:expr, $field:expr, $maybe_sort:expr) => {

        format_args!(

            "{}_{}_{}_{}",

            $prefix,

            $field,

            $maybe_sort.unwrap_or(""),

            FINGERPRINT_CREATE_INDEX_HEX

        )

    };

}

/// Try to sanitize a string to be used as a Redis key.

/// We don't actually modify the string, just check if it's valid.

const fn try_sanitize(s: &str) -> Option<&str> {

    // Note: We cannot use for loops or iterators here because they are not const.

    // Allowing us to use a const function here gives the compiler the ability to

    // optimize this function away entirely in the case where the input is constant.

    let chars = s.as_bytes();

    let mut i: usize = 0;

    let len = s.len();

    loop {

        if i >= len {

  Branch (965:12): [True: 20, False: 369]

  Branch (965:12): [Folded - Ignored]

            break;

        }

        let c = chars[i];

        if !c.is_ascii_alphanumeric() && 
c != b'_'15
 {

  Branch (969:12): [True: 15, False: 354]
  Branch (969:42): [True: 0, False: 15]

  Branch (969:12): [Folded - Ignored]
  Branch (969:42): [Folded - Ignored]

            return None;

        }

        i += 1;

    }

    Some(s)

}

/// An individual subscription to a key in Redis.

#[derive(Debug)]

pub struct RedisSubscription {

    receiver: Option<tokio::sync::watch::Receiver<String>>,

    weak_subscribed_keys: Weak<RwLock<StringPatriciaMap<RedisSubscriptionPublisher>>>,

}

impl SchedulerSubscription for RedisSubscription {

    /// Wait for the subscription key to change.

    async fn changed(&mut self) -> Result<(), Error> {

        let receiver = self

            .receiver

            .as_mut()

            .ok_or_else(|| make_err!(
Code::Internal0
, "In RedisSubscription::changed::as_mut"))
?0
;

        receiver

            .changed()

            .await

            .map_err(|_| make_err!(
Code::Internal0
, "In RedisSubscription::changed::changed"))

    }

}

// If the subscription is dropped, we need to possibly remove the key from the

// subscribed keys map.

impl Drop for RedisSubscription {

    fn drop(&mut self) {

        let Some(receiver) = self.receiver.take() else {

  Branch (1002:13): [True: 4, False: 0]

  Branch (1002:13): [Folded - Ignored]

            warn!("RedisSubscription has already been dropped, nothing to do.");