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

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

use core::pin::Pin;

use core::time::Duration;

use std::borrow::Cow;

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

use async_trait::async_trait;

use bytes::Bytes;

use const_format::formatcp;

use fred::clients::{Pool as RedisPool, 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, TryFutureExt, TryStreamExt, future};

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::time::sleep;

use tracing::{error, info, 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;

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

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

}

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 (148:12): [True: 0, False: 2]

  Branch (148: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 (154:13): [True: 2, False: 0]

  Branch (154: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 (173:16): [True: 2, False: 0]

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

                spec.retry.delay = DEFAULT_RETRY_DELAY;

            
}0

            if spec.retry.jitter == 0.0 {

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

  Branch (176: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 (201:16): [True: 2, False: 0]

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

                spec.broadcast_channel_capacity = DEFAULT_BROADCAST_CHANNEL_CAPACITY;

            
}0

            if spec.connection_timeout_ms == 0 {

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

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

                spec.connection_timeout_ms = DEFAULT_CONNECTION_TIMEOUT_MS;

            
}0

            if spec.command_timeout_ms == 0 {

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

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

                spec.command_timeout_ms = DEFAULT_COMMAND_TIMEOUT_MS;

            
}0

            if spec.connection_pool_size == 0 {

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

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

                spec.connection_pool_size = DEFAULT_CONNECTION_POOL_SIZE;

            
}0

            if spec.read_chunk_size == 0 {

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

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

                spec.read_chunk_size = DEFAULT_READ_CHUNK_SIZE;

            
}0

            if spec.max_chunk_uploads_per_update == 0 {

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

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

                spec.max_chunk_uploads_per_update = DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE;

            
}0

            if spec.scan_count == 0 {

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

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

                spec.scan_count = DEFAULT_SCAN_COUNT;

            
}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 = builder

            .build_pool(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,

        )

        .map(Arc::new)

    }

    /// Used for testing when determinism is required.

    #[expect(clippy::too_many_arguments)]

    pub fn new_from_builder_and_parts(

        client_pool: RedisPool,

        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,

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

        })

    }

    async fn get_client(&'_ self) -> Result<&'_ Client, Error> {

        let client = self.client_pool.next();

        let config = client.client_config();

        if config.mocks.is_none() {

  Branch (303:12): [True: 0, False: 0]

  Branch (303:12): [True: 2, False: 24]

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

  Branch (303:12): [True: 0, False: 57]

            client.wait_for_connect().await.err_tip(||

                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.unwrap_or_else(|| "None".to_string()),

                    config.database.unwrap_or_default()

                )

            )?;

        }

        Ok(client)

    }

    /// 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 (319:12): [True: 70, False: 4]

  Branch (319: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?;

        keys.iter()

            .zip(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()) {

  Branch (356:20): [True: 2, False: 4]

  Branch (356:20): [Folded - Ignored]

                    *result = Some(0);

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

                }

                let encoded_key = self.encode_key(key);

                let pipeline = client.pipeline();

                pipeline

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

                    .await

                    .err_tip(|| 
{0

                        format!("In RedisStore::has_with_results::strlen for {encoded_key}")

                    })?;

                // 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(|| 
{0

                        format!("In RedisStore::has_with_results::exists for {encoded_key}")

                    })?;

                let (blob_len, exists) = pipeline

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

                    .await

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

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

  Branch (383:30): [True: 4, False: 0]

  Branch (383:30): [Folded - Ignored]

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

                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

            .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

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

        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,

        );

        loop {

            let chunk: Bytes = 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

                .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 (757:15): [True: 101, False: 1]

            let byte = DATA[i];

            crc ^= byte as u32;

            let mut j = 0;

            while j < 8 {

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

                crc = if crc & 1 != 0 {

  Branch (763: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 (810:12): [True: 22, False: 547]

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

            break;

        }

        let c = chars[i];

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

  Branch (814:12): [True: 25, False: 522]
  Branch (814:42): [True: 0, False: 25]

  Branch (814:12): [Folded - Ignored]
  Branch (814: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 (847:13): [True: 6, False: 0]

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

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

            return; // Already dropped, nothing to do.

        };

        let key = receiver.borrow().clone();

        // IMPORTANT: This must be dropped before receiver_count() is called.

        drop(receiver);

        let Some(
subscribed_keys4
) = self.weak_subscribed_keys.upgrade() else {

  Branch (854:13): [True: 4, False: 2]

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

            return; // Already dropped, nothing to do.

        };

        let mut subscribed_keys = subscribed_keys.write();

        let Some(value) = subscribed_keys.get(&key) else {

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

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

            error!(

                "Key {key} was not found in subscribed keys when checking if it should be removed."

            );

            return;

        };

        // If we have no receivers, cleanup the entry from our map.

        if value.receiver_count() == 0 {

  Branch (865:12): [True: 4, False: 0]

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

            subscribed_keys.remove(key);

        
}0

    }

}

/// A publisher for a key in Redis.

#[derive(Debug)]

struct RedisSubscriptionPublisher {

    sender: Mutex<tokio::sync::watch::Sender<String>>,

}

impl RedisSubscriptionPublisher {

    fn new(

        key: String,

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

    ) -> (Self, RedisSubscription) {

        let (sender, receiver) = tokio::sync::watch::channel(key);

        let publisher = Self {

            sender: Mutex::new(sender),

        };

        let subscription = RedisSubscription {

            receiver: Some(receiver),

            weak_subscribed_keys,

        };

        (publisher, subscription)

    }

    fn subscribe(

        &self,

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

    ) -> RedisSubscription {

        let receiver = self.sender.lock().subscribe();

        RedisSubscription {

            receiver: Some(receiver),

            weak_subscribed_keys,

        }

    }

    fn receiver_count(&self) -> usize {

        self.sender.lock().receiver_count()

    }

    fn notify(&self) {

        // TODO(https://github.com/sile/patricia_tree/issues/40) When this is addressed

        // we can remove the `Mutex` and use the mutable iterator directly.

        self.sender.lock().send_modify(|_| {});

    }

}

#[derive(Debug)]

pub struct RedisSubscriptionManager {

    subscribed_keys: Arc<RwLock<StringPatriciaMap<RedisSubscriptionPublisher>>>,

    tx_for_test: tokio::sync::mpsc::UnboundedSender<String>,

    _subscription_spawn: JoinHandleDropGuard<()>,

}

impl RedisSubscriptionManager {

    pub fn new(subscribe_client: SubscriberClient, pub_sub_channel: String) -> Self {

        let subscribed_keys = Arc::new(RwLock::new(StringPatriciaMap::new()));

        let subscribed_keys_weak = Arc::downgrade(&subscribed_keys);

        let (tx_for_test, mut rx_for_test) = tokio::sync::mpsc::unbounded_channel();

        Self {

            subscribed_keys,

            tx_for_test,

            _subscription_spawn: spawn!("redis_subscribe_spawn", async move {

                let mut rx = subscribe_client.message_rx();

                loop {

                    if let Err(
e0
) = subscribe_client.subscribe(&pub_sub_channel).await {

  Branch (933:28): [True: 0, False: 4]

  Branch (933:28): [Folded - Ignored]

                        error!("Error subscribing to pattern - {e}");

                        return;

                    }

                    let mut reconnect_rx = subscribe_client.reconnect_rx();

                    let reconnect_fut = reconnect_rx.recv().fuse();

                    tokio::pin!(reconnect_fut);

                    loop {

                        let 
key17
 = select! {

                            
value17
 = rx_for_test.recv() => {

                                let Some(value) = value else {

  Branch (943:37): [True: 17, False: 0]

  Branch (943:37): [Folded - Ignored]

                                    unreachable!("Channel should never close");

                                };

                                value.into()

                            },

                            
msg0
 = rx.recv() => {

                                match msg {

                                    Ok(msg) => {

                                        if let RedisValue::String(s) = msg.value {

  Branch (951:48): [True: 0, False: 0]

  Branch (951:48): [Folded - Ignored]

                                            s

                                        } else {

                                            error!("Received non-string message in RedisSubscriptionManager");

                                            continue;

                                        }

                                    },

                                    Err(e) => {

                                        // Check to see if our parent has been dropped and if so kill spawn.

                                        if subscribed_keys_weak.upgrade().is_none() {

  Branch (960:44): [True: 0, False: 0]

  Branch (960:44): [Folded - Ignored]

                                            warn!("It appears our parent has been dropped, exiting RedisSubscriptionManager spawn");

                                            return;

                                        }

                                        error!("Error receiving message in RedisSubscriptionManager reconnecting and flagging everything changed - {e}");

                                        break;

                                    }

                                }

                            },

                            _ = &mut reconnect_fut => {

                                warn!("Redis reconnected flagging all subscriptions as changed and resuming");

                                break;

                            }

                        };

                        let Some(subscribed_keys) = subscribed_keys_weak.upgrade() else {

  Branch (974:29): [True: 17, False: 0]

  Branch (974:29): [Folded - Ignored]

                            warn!(

                                "It appears our parent has been dropped, exiting RedisSubscriptionManager spawn"

                            );

                            return;

                        };

                        let subscribed_keys_mux = subscribed_keys.read();

                        subscribed_keys_mux

                            .common_prefix_values(&*key)

                            .for_each(RedisSubscriptionPublisher::notify);

                    }

                    // Sleep for a small amount of time to ensure we don't reconnect too quickly.

                    sleep(Duration::from_secs(1)).await;

                    // If we reconnect or lag behind we might have had dirty keys, so we need to

                    // flag all of them as changed.

                    let Some(subscribed_keys) = subscribed_keys_weak.upgrade() else {

  Branch (989:25): [True: 0, False: 0]

  Branch (989:25): [Folded - Ignored]

                        warn!(

                            "It appears our parent has been dropped, exiting RedisSubscriptionManager spawn"

                        );

                        return;