/build/source/nativelink-util/src/store_trait.rs

Source
// 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::borrow::{Borrow, BorrowMut};
use core::convert::Into;
use core::fmt::{self, Debug, Display};
use core::future;
use core::hash::{Hash, Hasher};
use core::ops::{Bound, RangeBounds};
use core::pin::Pin;
use core::ptr::addr_eq;
use core::time::Duration;
use std::borrow::Cow;
use std::collections::hash_map::DefaultHasher as StdHasher;
use std::ffi::OsString;
use std::sync::{Arc, OnceLock};

use async_trait::async_trait;
use bytes::{Bytes, BytesMut};
use futures::{Future, FutureExt, Stream, join, try_join};
use nativelink_error::{Code, Error, ResultExt, error_if, make_err};
use nativelink_metric::MetricsComponent;
use rand::rngs::StdRng;
use rand::{RngCore, SeedableRng};
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncReadExt, AsyncSeekExt};
use tracing::warn;

use crate::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair};
use crate::common::DigestInfo;
use crate::digest_hasher::{DigestHasher, DigestHasherFunc, default_digest_hasher_func};
use crate::fs;
use crate::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};

static DEFAULT_DIGEST_SIZE_HEALTH_CHECK: OnceLock<usize> = OnceLock::new();
/// Default digest size for health check data. Any change in this value
/// changes the default contract. `GlobalConfig` should be updated to reflect
/// changes in this value.
pub const DEFAULT_DIGEST_SIZE_HEALTH_CHECK_CFG: usize = 1024 * 1024;

// Get the default digest size for health check data, if value is unset a system wide default is used.
pub fn default_digest_size_health_check() -> usize {
    *DEFAULT_DIGEST_SIZE_HEALTH_CHECK.get_or_init(|| DEFAULT_DIGEST_SIZE_HEALTH_CHECK_CFG)
}

/// Set the default digest size for health check data, this should be called once.
pub fn set_default_digest_size_health_check(size: usize) -> Result<(), Error> {
    DEFAULT_DIGEST_SIZE_HEALTH_CHECK.set(size).map_err(|_| {
        make_err!(
            Code::Internal,
            "set_default_digest_size_health_check already set"
        )
    })
}

#[derive(
    Debug,
    PartialEq,
    Eq,
    Copy,
    Clone,
    Serialize,
    Deserialize,
    wincode::SchemaWrite,
    wincode::SchemaRead,
)]
pub enum UploadSizeInfo {
    /// When the data transfer amount is known to be exact size, this enum should be used.
    /// The receiver store can use this to better optimize the way the data is sent or stored.
    ExactSize(u64),

    /// When the data transfer amount is not known to be exact, the caller should use this enum
    /// to provide the maximum size that could possibly be sent. This will bypass the exact size
    /// checks, but still provide useful information to the underlying store about the data being
    /// sent that it can then use to optimize the upload process.
    MaxSize(u64),
}

/// Utility to send all the data to the store from a file.
// Note: This is not inlined because some code may want to bypass any underlying
// optimizations that may be present in the inner store.
pub async fn slow_update_store_with_file<S: StoreDriver + ?Sized>(
    store: Pin<&S>,
    digest: impl Into<StoreKey<'_>>,
    file: &mut fs::FileSlot,
    upload_size: UploadSizeInfo,
) -> Result<u64, Error> {
    file.rewind()
        .await
        .err_tip(|| "Failed to rewind in upload_file_to_store")?0;
    let (mut tx, rx) = make_buf_channel_pair();

    let update_fut = store
        .update(digest.into(), rx, upload_size)
        .map(|r| r.err_tip(|| "Could not upload data to store in upload_file_to_store"));
    let read_data_fut = async move {
        loop {
            let mut buf = BytesMut::with_capacity(fs::DEFAULT_READ_BUFF_SIZE);
            let read = file
                .read_buf(&mut buf)
                .await
                .err_tip(|| "Failed to read in upload_file_to_store")?0;
            if read == 0 {
                break;
            }
            tx.send(buf.freeze())
                .await
                .err_tip(|| "Failed to send in upload_file_to_store")?0;
        }
        tx.send_eof()
            .err_tip(|| "Could not send EOF to store in upload_file_to_store")
    };
    tokio::pin!(read_data_fut);
    let (update_res, read_res) = tokio::join!(update_fut, read_data_fut);
    read_res.merge(update_res)
}

/// Optimizations that stores may want to expose to the callers.
/// This is useful for specific cases when the store can optimize the processing
/// of the data being processed.
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub enum StoreOptimizations {
    /// The store can optimize the upload process when it knows the data is coming from a file.
    FileUpdates,

    /// If the store will ignore the data uploads.
    NoopUpdates,

    /// If the store will never serve downloads.
    NoopDownloads,

    /// If the store will determine whether a key has associated data once a read has been
    /// attempted instead of calling `.has()` first.
    LazyExistenceOnSync,

    /// The store provides an optimized `update_oneshot` implementation that bypasses
    /// channel overhead for direct Bytes writes. Stores with this optimization can
    /// accept complete data directly without going through the MPSC channel.
    SubscribesToUpdateOneshot,
}

/// A wrapper struct for [`StoreKey`] to work around
/// lifetime limitations in `HashMap::get()` as described in
/// <https://github.com/rust-lang/rust/issues/80389>
///
/// As such this is a wrapper type that is stored in the
/// maps using the workaround as described in
/// <https://blinsay.com/blog/compound-keys/>
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct StoreKeyBorrow(StoreKey<'static>);

impl From<StoreKey<'static>> for StoreKeyBorrow {
    fn from(key: StoreKey<'static>) -> Self {
        Self(key)
    }
}

impl From<StoreKeyBorrow> for StoreKey<'static> {
    fn from(key_borrow: StoreKeyBorrow) -> Self {
        key_borrow.0
    }
}

impl<'a> Borrow<StoreKey<'a>> for StoreKeyBorrow {
    fn borrow(&self) -> &StoreKey<'a> {
        &self.0
    }
}

impl<'a> Borrow<StoreKey<'a>> for &StoreKeyBorrow {
    fn borrow(&self) -> &StoreKey<'a> {
        &self.0
    }
}

/// Holds something that can be converted into a key the
/// store API can understand. Generally this is a digest
/// but it can also be a string if the caller wishes to
/// store the data directly and reference it by a string
/// directly.
#[derive(Debug, Eq)]
pub enum StoreKey<'a> {
    /// A string key.
    Str(Cow<'a, str>),

    /// A key that is a digest.
    Digest(DigestInfo),
}

impl<'a> StoreKey<'a> {
    /// Creates a new store key from a string.
    pub const fn new_str(s: &'a str) -> Self {
        StoreKey::Str(Cow::Borrowed(s))
    }

    /// Returns a shallow clone of the key.
    /// This is extremely cheap and should be used when clone
    /// is needed but the key is not going to be modified.
    #[must_use]
    #[allow(
        clippy::missing_const_for_fn,
        reason = "False positive on stable, but not on nightly"
    )]
    pub fn borrow(&'a self) -> Self {
        match self {
            StoreKey::Str(Cow::Owned(s)) => StoreKey::Str(Cow::Borrowed(s)),
            StoreKey::Str(Cow::Borrowed(s)) => StoreKey::Str(Cow::Borrowed(s)),
            StoreKey::Digest(d) => StoreKey::Digest(*d),
        }
    }

    /// Converts the key into an owned version. This is useful
    /// when the caller needs an owned version of the key.
    pub fn into_owned(self) -> StoreKey<'static> {
        match self {
            StoreKey::Str(Cow::Owned(s)) => StoreKey::Str(Cow::Owned(s)),
            StoreKey::Str(Cow::Borrowed(s)) => StoreKey::Str(Cow::Owned(s.to_owned())),
            StoreKey::Digest(d) => StoreKey::Digest(d),
        }
    }

    /// Converts the key into a digest. This is useful when the caller
    /// must have a digest key. If the data is not a digest, it may
    /// hash the underlying key and return a digest of the hash of the key
    pub fn into_digest(self) -> DigestInfo {
        match self {
            StoreKey::Digest(digest) => digest,
            StoreKey::Str(s) => {
                let mut hasher = DigestHasherFunc::Blake3.hasher();
                hasher.update(s.as_bytes());
                hasher.finalize_digest()
            }
        }
    }

    /// Returns the key as a string. If the key is a digest, it will
    /// return a string representation of the digest. If the key is a string,
    /// it will return the string itself.
    pub fn as_str(&'a self) -> Cow<'a, str> {
        match self {
            StoreKey::Str(Cow::Owned(s)) => Cow::Borrowed(s),
            StoreKey::Str(Cow::Borrowed(s)) => Cow::Borrowed(s),
            StoreKey::Digest(d) => Cow::Owned(format!("{d}")),
        }
    }
}

impl Clone for StoreKey<'static> {
    fn clone(&self) -> Self {
        match self {
            StoreKey::Str(s) => StoreKey::Str(s.clone()),
            StoreKey::Digest(d) => StoreKey::Digest(*d),
        }
    }
}

impl PartialOrd for StoreKey<'_> {
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for StoreKey<'_> {
    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
        match (self, other) {
            (StoreKey::Str(a), StoreKey::Str(b)) => a.cmp(b),
            (StoreKey::Digest(a), StoreKey::Digest(b)) => a.cmp(b),
            (StoreKey::Str(_), StoreKey::Digest(_)) => core::cmp::Ordering::Less,
            (StoreKey::Digest(_), StoreKey::Str(_)) => core::cmp::Ordering::Greater,
        }
    }
}

impl PartialEq for StoreKey<'_> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (StoreKey::Str(a), StoreKey::Str(b)) => a == b,
            (StoreKey::Digest(a), StoreKey::Digest(b)) => a == b,
            _ => false,
        }
    }
}

impl Hash for StoreKey<'_> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        /// Salts the hash with the enum value that represents
        /// the type of the key.
        #[repr(u8)]
        enum HashId {
            Str = 0,
            Digest = 1,
        }
        match self {
            StoreKey::Str(s) => {
                (HashId::Str as u8).hash(state);
                s.hash(state);
            }
            StoreKey::Digest(d) => {
                (HashId::Digest as u8).hash(state);
                d.hash(state);
            }
        }
    }
}

impl<'a> From<&'a str> for StoreKey<'a> {
    fn from(s: &'a str) -> Self {
        StoreKey::Str(Cow::Borrowed(s))
    }
}

impl From<String> for StoreKey<'static> {
    fn from(s: String) -> Self {
        StoreKey::Str(Cow::Owned(s))
    }
}

impl From<DigestInfo> for StoreKey<'_> {
    fn from(d: DigestInfo) -> Self {
        StoreKey::Digest(d)
    }
}

impl From<&DigestInfo> for StoreKey<'_> {
    fn from(d: &DigestInfo) -> Self {
        StoreKey::Digest(*d)
    }
}

// mostly for use with tracing::Value
impl Display for StoreKey<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StoreKey::Str(s) => {
                write!(f, "{s}")
            }
            StoreKey::Digest(d) => {
                write!(f, "Digest: {d}")
            }
        }
    }
}

#[derive(Clone, MetricsComponent)]
#[repr(transparent)]
pub struct Store {
    #[metric]
    inner: Arc<dyn StoreDriver>,
}

impl Debug for Store {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Store").finish_non_exhaustive()
    }
}

impl Store {
    pub fn new(inner: Arc<dyn StoreDriver>) -> Self {
        Self { inner }
    }

    /// Returns the immediate inner store driver.
    /// Note: This does not recursively try to resolve underlying store drivers
    /// like `.inner_store()` does.
    #[inline]
    pub fn into_inner(self) -> Arc<dyn StoreDriver> {
        self.inner
    }

    /// Gets the underlying store for the given digest.
    /// A caller might want to use this to obtain a reference to the "real" underlying store
    /// (if applicable) and check if it implements some special traits that allow optimizations.
    /// Note: If the store performs complex operations on the data, it should return itself.
    #[inline]
    pub fn inner_store<'a, K: Into<StoreKey<'a>>>(&self, digest: Option<K>) -> &dyn StoreDriver {
        self.inner.inner_store(digest.map(Into::into))
    }

    /// Tries to cast the underlying store to the given type.
    #[inline]
    pub fn downcast_ref<U: StoreDriver>(&self, maybe_digest: Option<StoreKey<'_>>) -> Option<&U> {
        self.inner.inner_store(maybe_digest).as_any().downcast_ref()
    }

    /// Register health checks used to monitor the store.
    #[inline]
    pub fn register_health(&self, registry: &mut HealthRegistryBuilder) {
        self.inner.clone().register_health(registry);
    }

    #[inline]
    pub fn register_remove_callback(
        &self,
        callback: Arc<dyn RemoveItemCallback>,
    ) -> Result<(), Error> {
        self.inner.clone().register_remove_callback(callback)
    }
}

impl StoreLike for Store {
    #[inline]
    fn as_store_driver(&self) -> &'_ dyn StoreDriver {
        self.inner.as_ref()
    }

    fn as_pin(&self) -> Pin<&Self> {
        Pin::new(self)
    }
}

impl<T> StoreLike for T
where
    T: StoreDriver + Sized,
{
    #[inline]
    fn as_store_driver(&self) -> &'_ dyn StoreDriver {
        self
    }

    fn as_pin(&self) -> Pin<&Self> {
        Pin::new(self)
    }
}

pub trait StoreLike: Send + Sync + Sized + Unpin + 'static {
    /// Returns the immediate inner store driver.
    fn as_store_driver(&self) -> &'_ dyn StoreDriver;

    /// Utility function to return a pinned reference to self.
    fn as_pin(&self) -> Pin<&Self>;

    /// Utility function to return a pinned reference to the store driver.
    #[inline]
    fn as_store_driver_pin(&self) -> Pin<&'_ dyn StoreDriver> {
        Pin::new(self.as_store_driver())
    }

    /// Look up a digest in the store and return None if it does not exist in
    /// the store, or Some(size) if it does.
    /// Note: On an AC store the size will be incorrect and should not be used!
    #[inline]
    fn has<'a>(
        &'a self,
        digest: impl Into<StoreKey<'a>>,
    ) -> impl Future<Output = Result<Option<u64>, Error>> + 'a {
        self.as_store_driver_pin().has(digest.into())
    }

    /// Look up a list of digests in the store and return a result for each in
    /// the same order as input.  The result will either be None if it does not
    /// exist in the store, or Some(size) if it does.
    /// Note: On an AC store the size will be incorrect and should not be used!
    #[inline]
    fn has_many<'a>(
        &'a self,
        digests: &'a [StoreKey<'a>],
    ) -> impl Future<Output = Result<Vec<Option<u64>>, Error>> + Send + 'a {
        if digests.is_empty() {
            return future::ready(Ok(vec![])).boxed();
        }
        self.as_store_driver_pin().has_many(digests)
    }

    /// The implementation of the above has and `has_many` functions.  See their
    /// documentation for details.
    #[inline]
    fn has_with_results<'a>(
        &'a self,
        digests: &'a [StoreKey<'a>],
        results: &'a mut [Option<u64>],
    ) -> impl Future<Output = Result<(), Error>> + Send + 'a {
        if digests.is_empty() {
            return future::ready(Ok(())).boxed();
        }
        self.as_store_driver_pin()
            .has_with_results(digests, results)
    }

    /// List all the keys in the store that are within the given range.
    /// `handler` is called for each key in the range. If `handler` returns
    /// false, the listing is stopped.
    ///
    /// The number of keys passed through the handler is the return value.
    #[inline]
    fn list<'a, 'b>(
        &'a self,
        range: impl RangeBounds<StoreKey<'b>> + Send + 'b,
        mut handler: impl for<'c> FnMut(&'c StoreKey) -> bool + Send + Sync + 'a,
    ) -> impl Future<Output = Result<u64, Error>> + Send + 'a
    where
        'b: 'a,
    {
        // Note: We use a manual async move, so the future can own the `range` and `handler`,
        // otherwise we'd require the caller to pass them in by reference making more borrow
        // checker noise.
        async move {
            self.as_store_driver_pin()
                .list(
                    (
                        range.start_bound().map(StoreKey::borrow),
                        range.end_bound().map(StoreKey::borrow),
                    ),
                    &mut handler,
                )
                .await
        }
    }

    /// Sends the data to the store.
    #[inline]
    fn update<'a>(
        &'a self,
        digest: impl Into<StoreKey<'a>>,
        reader: DropCloserReadHalf,
        upload_size: UploadSizeInfo,
    ) -> impl Future<Output = Result<u64, Error>> + Send + 'a {
        self.as_store_driver_pin()
            .update(digest.into(), reader, upload_size)
    }

    /// Any optimizations the store might want to expose to the callers.
    /// By default, no optimizations are exposed.
    #[inline]
    fn optimized_for(&self, optimization: StoreOptimizations) -> bool {
        self.as_store_driver_pin().optimized_for(optimization)
    }

    /// Specialized version of `.update()` which takes a `FileSlot`.
    /// This is useful if the underlying store can optimize the upload process
    /// when it knows the data is coming from a file.
    #[inline]
    fn update_with_whole_file<'a>(
        &'a self,
        digest: impl Into<StoreKey<'a>>,
        path: OsString,
        file: fs::FileSlot,
        upload_size: UploadSizeInfo,
    ) -> impl Future<Output = Result<(u64, Option<fs::FileSlot>), Error>> + Send + 'a {
        self.as_store_driver_pin()
            .update_with_whole_file(digest.into(), path, file, upload_size)
    }

    /// Utility to send all the data to the store when you have all the bytes.
    #[inline]
    fn update_oneshot<'a>(
        &'a self,
        digest: impl Into<StoreKey<'a>>,
        data: Bytes,
    ) -> impl Future<Output = Result<(), Error>> + Send + 'a {
        self.as_store_driver_pin()
            .update_oneshot(digest.into(), data)
    }

    /// Retrieves part of the data from the store and writes it to the given writer.
    #[inline]
    fn get_part<'a>(
        &'a self,
        digest: impl Into<StoreKey<'a>>,
        mut writer: impl BorrowMut<DropCloserWriteHalf> + Send + 'a,
        offset: u64,
        length: Option<u64>,
    ) -> impl Future<Output = Result<(), Error>> + Send + 'a {
        let key = digest.into();
        // Note: We need to capture `writer` just in case the caller
        // expects the drop() method to be called on it when the future
        // is done due to the complex interaction between the DropCloserWriteHalf
        // and the DropCloserReadHalf during drop().
        async move {
            self.as_store_driver_pin()
                .get_part(key, writer.borrow_mut(), offset, length)
                .await
        }
    }

    /// Utility that works the same as `.get_part()`, but writes all the data.
    #[inline]
    fn get<'a>(
        &'a self,
        key: impl Into<StoreKey<'a>>,
        writer: DropCloserWriteHalf,
    ) -> impl Future<Output = Result<(), Error>> + Send + 'a {
        self.as_store_driver_pin().get(key.into(), writer)
    }

    /// Utility that will return all the bytes at once instead of in a streaming manner.
    #[inline]
    fn get_part_unchunked<'a>(
        &'a self,
        key: impl Into<StoreKey<'a>>,
        offset: u64,
        length: Option<u64>,
    ) -> impl Future<Output = Result<Bytes, Error>> + Send + 'a {
        self.as_store_driver_pin()
            .get_part_unchunked(key.into(), offset, length)
    }

    /// Default implementation of the health check. Some stores may want to override this
    /// in situations where the default implementation is not sufficient.
    #[inline]
    fn check_health(
        &self,
        namespace: Cow<'static, str>,
    ) -> impl Future<Output = HealthStatus> + Send {
        self.as_store_driver_pin().check_health(namespace)
    }
}

#[async_trait]
pub trait StoreDriver:
    Sync + Send + Unpin + MetricsComponent + HealthStatusIndicator + 'static
{
    // Do "all the stores are setup" init e.g. if we need access to the store manager
    // for ref stores
    async fn post_init(self: Arc<Self>) -> Result<(), Error>;

    /// See: [`StoreLike::has`] for details.
    #[inline]
    async fn has(self: Pin<&Self>, key: StoreKey<'_>) -> Result<Option<u64>, Error> {
        let mut result = [None];
        self.has_with_results(&[key], &mut result).await?;
        Ok(result[0])
    }

    /// See: [`StoreLike::has_many`] for details.
    #[inline]
    async fn has_many(
        self: Pin<&Self>,
        digests: &[StoreKey<'_>],
    ) -> Result<Vec<Option<u64>>, Error> {
        let mut results = vec![None; digests.len()];
        self.has_with_results(digests, &mut results).await?;
        Ok(results)
    }

    /// See: [`StoreLike::has_with_results`] for details.
    async fn has_with_results(
        self: Pin<&Self>,
        digests: &[StoreKey<'_>],
        results: &mut [Option<u64>],
    ) -> Result<(), Error>;

    /// See: [`StoreLike::list`] for details.
    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> {
        // TODO(palfrey) We should force all stores to implement this function instead of
        // providing a default implementation.
        Err(make_err!(
            Code::Unimplemented,
            "Store::list() not implemented for this store"
        ))
    }

    /// See: [`StoreLike::update`] for details.
    async fn update(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        reader: DropCloserReadHalf,
        upload_size: UploadSizeInfo,
    ) -> Result<u64, Error>;

    /// See: [`StoreLike::optimized_for`] for details.
    fn optimized_for(&self, _optimization: StoreOptimizations) -> bool {
        false
    }

    /// See: [`StoreLike::update_with_whole_file`] for details.
    async fn update_with_whole_file(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        path: OsString,
        mut file: fs::FileSlot,
        upload_size: UploadSizeInfo,
    ) -> Result<(u64, Option<fs::FileSlot>), Error> {
        let inner_store = self.inner_store(Some(key.borrow()));
        if inner_store.optimized_for(StoreOptimizations::FileUpdates) {
            error_if!(
                addr_eq(inner_store, &raw const *self),
                "Store::inner_store() returned self when optimization present"
            );
            return Pin::new(inner_store)
                .update_with_whole_file(key, path, file, upload_size)
                .await;
        }
        let size = slow_update_store_with_file(self, key, &mut file, upload_size).await?;
        Ok((size, Some(file)))
    }

    /// See: [`StoreLike::update_oneshot`] for details.
    async fn update_oneshot(self: Pin<&Self>, key: StoreKey<'_>, data: Bytes) -> Result<(), Error> {
        // TODO(palfrey) This is extremely inefficient, since we have exactly
        // what we need here. Maybe we could instead make a version of the stream
        // that can take objects already fully in memory instead?
        let (mut tx, rx) = make_buf_channel_pair();

        let data_len =
            u64::try_from(data.len()).err_tip(|| "Could not convert data.len() to u64")?;
        let send_fut = async move {
            // Only send if we are not EOF.
            if !data.is_empty() {
                tx.send(data)
                    .await
                    .err_tip(|| "Failed to write data in update_oneshot")?0;
            }
            tx.send_eof()
                .err_tip(|| "Failed to write EOF in update_oneshot")?0;
            Ok(())
        };
        try_join!(
            send_fut,
            self.update(key, rx, UploadSizeInfo::ExactSize(data_len))
        )?;
        Ok(())
    }

    /// See: [`StoreLike::get_part`] for details.
    async fn get_part(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        writer: &mut DropCloserWriteHalf,
        offset: u64,
        length: Option<u64>,
    ) -> Result<(), Error>;

    /// See: [`StoreLike::get`] for details.
    #[inline]
    async fn get(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        mut writer: DropCloserWriteHalf,
    ) -> Result<(), Error> {
        self.get_part(key, &mut writer, 0, None).await
    }

    /// See: [`StoreLike::get_part_unchunked`] for details.
    async fn get_part_unchunked(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        offset: u64,
        length: Option<u64>,
    ) -> Result<Bytes, Error> {
        let length_usize = length
            .map(|v| usize::try_from(v).err_tip(|| "Could not convert length to usize"))
            .transpose()?;

        // TODO(palfrey) This is extremely inefficient, since we have exactly
        // what we need here. Maybe we could instead make a version of the stream
        // that can take objects already fully in memory instead?
        let (mut tx, mut rx) = make_buf_channel_pair();

        let (data_res, get_part_res) = join!(
            rx.consume(length_usize),
            // We use a closure here to ensure that the `tx` is dropped when the
            // future is done.
            async move { self.get_part(key, &mut tx, offset, length).await }5.50k,
        );
        get_part_res
            .err_tip(|| "Failed to get_part in get_part_unchunked")
            .merge(data_res.err_tip(|| "Failed to read stream to completion in get_part_unchunked"))
    }

    /// See: [`StoreLike::check_health`] for details.
    async fn check_health(self: Pin<&Self>, namespace: Cow<'static, str>) -> HealthStatus {
        let digest_data_size = default_digest_size_health_check();
        let mut digest_data = vec![0u8; digest_data_size];

        let mut namespace_hasher = StdHasher::new();
        namespace.hash(&mut namespace_hasher);
        self.get_name().hash(&mut namespace_hasher);
        let hash_seed = namespace_hasher.finish();

        // Fill the digest data with random data based on a stable
        // hash of the namespace and store name. Intention is to
        // have randomly filled data that is unique per store and
        // does not change between health checks. This is to ensure
        // we are not adding more data to store on each health check.
        let mut rng: StdRng = StdRng::seed_from_u64(hash_seed);
        rng.fill_bytes(&mut digest_data);

        let mut digest_hasher = default_digest_hasher_func().hasher();
        digest_hasher.update(&digest_data);
        let digest_data_len = digest_data.len() as u64;
        let digest_info = StoreKey::from(digest_hasher.finalize_digest());

        let digest_bytes = Bytes::copy_from_slice(&digest_data);

        if let Err(e) = self
            .update_oneshot(digest_info.borrow(), digest_bytes.clone())
            .await
        {
            warn!(?e, "check_health Store.update_oneshot() failed");
            return HealthStatus::new_failed(
                self.get_ref(),
                format!("Store.update_oneshot() failed: {e}").into(),
            );
        }

        match self.has(digest_info.borrow()).await {
            Ok(Some(s)) => {
                if s != digest_data_len {
                    return HealthStatus::new_failed(
                        self.get_ref(),
                        format!("Store.has() size mismatch {s} != {digest_data_len}").into(),
                    );
                }
            }
            Ok(None) => {
                return HealthStatus::new_failed(
                    self.get_ref(),
                    "Store.has() size not found".into(),
                );
            }
            Err(e) => {
                return HealthStatus::new_failed(
                    self.get_ref(),
                    format!("Store.has() failed: {e}").into(),
                );
            }
        }

        match self
            .get_part_unchunked(digest_info, 0, Some(digest_data_len))
            .await
        {
            Ok(b) => {
                if b != digest_bytes {
                    return HealthStatus::new_failed(
                        self.get_ref(),
                        "Store.get_part_unchunked() data mismatch".into(),
                    );
                }
            }
            Err(e) => {
                return HealthStatus::new_failed(
                    self.get_ref(),
                    format!("Store.get_part_unchunked() failed: {e}").into(),
                );
            }
        }

        HealthStatus::new_ok(self.get_ref(), "Successfully store health check".into())
    }

    /// See: [`Store::inner_store`] for details.
    fn inner_store(&self, _digest: Option<StoreKey<'_>>) -> &dyn StoreDriver;

    /// Returns an Any variation of whatever Self is.
    fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static);
    fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static>;

    // Register health checks used to monitor the store.
    fn register_health(self: Arc<Self>, _registry: &mut HealthRegistryBuilder) {}

    fn register_remove_callback(
        self: Arc<Self>,
        callback: Arc<dyn RemoveItemCallback>,
    ) -> Result<(), Error>;
}

// Callback to be called when a store deletes an item. This is used so
// compound stores can remove items from their internal state when their
// underlying stores remove items e.g. caches
pub trait RemoveItemCallback: Debug + Send + Sync {
    fn callback<'a>(
        &'a self,
        store_key: StoreKey<'a>,
    ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
}

/// The instructions on how to decode a value from a Bytes & version into
/// the underlying type.
pub trait SchedulerStoreDecodeTo {
    type DecodeOutput;
    fn decode(version: i64, data: Bytes) -> Result<Self::DecodeOutput, Error>;
}

pub trait SchedulerSubscription: Send + Sync {
    fn changed(&mut self) -> impl Future<Output = Result<(), Error>> + Send;
}

pub trait SchedulerSubscriptionManager: Send + Sync {
    type Subscription: SchedulerSubscription;

    fn subscribe<K>(&self, key: K) -> Result<Self::Subscription, Error>
    where
        K: SchedulerStoreKeyProvider;

    fn is_reliable() -> bool;
}

/// The API surface for a scheduler store.
pub trait SchedulerStore: Send + Sync + 'static {
    type SubscriptionManager: SchedulerSubscriptionManager;

    /// Returns the subscription manager for the scheduler store.
    fn subscription_manager(
        &self,
    ) -> impl Future<Output = Result<Arc<Self::SubscriptionManager>, Error>> + Send;

    /// Updates or inserts an entry into the underlying store.
    /// Metadata about the key is attached to the compile-time type.
    /// If `StoreKeyProvider::Versioned` is `TrueValue`, the data will not
    /// be updated if the current version in the database does not match
    /// the version in the passed in data.
    /// No guarantees are made about when `Version` is `FalseValue`.
    /// Indexes are guaranteed to be updated atomically with the data.
    fn update_data<T>(
        &self,
        data: T,
        expiry: Option<Duration>,
    ) -> impl Future<Output = Result<Option<i64>, Error>> + Send
    where
        T: SchedulerStoreDataProvider
            + SchedulerStoreKeyProvider
            + SchedulerCurrentVersionProvider
            + Send;

    /// Searches for all keys in the store that match the given index prefix.
    fn search_by_index_prefix<K>(
        &self,
        index: K,
    ) -> impl Future<
        Output = Result<
            impl Stream<Item = Result<<K as SchedulerStoreDecodeTo>::DecodeOutput, Error>> + Send,
            Error,
        >,
    > + Send
    where
        K: SchedulerIndexProvider + SchedulerStoreDecodeTo + Send,
        <K as SchedulerStoreDecodeTo>::DecodeOutput: Send;

    /// Returns data for the provided key with the given version if
    /// `StoreKeyProvider::Versioned` is `TrueValue`.
    fn get_and_decode<K>(
        &self,
        key: K,
    ) -> impl Future<Output = Result<Option<<K as SchedulerStoreDecodeTo>::DecodeOutput>, Error>> + Send
    where
        K: SchedulerStoreKeyProvider + SchedulerStoreDecodeTo + Send;
}

/// A type that is used to let the scheduler store know what
/// index is being requested.
pub trait SchedulerIndexProvider {
    /// Only keys inserted with this prefix will be indexed.
    const KEY_PREFIX: &'static str;

    /// The name of the index.
    const INDEX_NAME: &'static str;

    /// The sort key for the index (if any).
    const MAYBE_SORT_KEY: Option<&'static str> = None;

    /// If the data is versioned.
    type Versioned: BoolValue;

    /// The value of the index.
    fn index_value(&self) -> Cow<'_, str>;
}

/// Provides a key to lookup data in the store.
pub trait SchedulerStoreKeyProvider {
    /// If the data is versioned.
    type Versioned: BoolValue;

    /// Returns the key for the data.
    fn get_key(&self) -> StoreKey<'static>;
}

/// Provides data to be stored in the scheduler store.
pub trait SchedulerStoreDataProvider {
    /// Converts the data into bytes to be stored in the store.
    fn try_into_bytes(self) -> Result<Bytes, Error>;

    /// Returns the indexes for the data if any.
    fn get_indexes(&self) -> Result<Vec<(&'static str, Bytes)>, Error> {
        Ok(Vec::new())
    }
}

/// Provides the current version of the data in the store.
pub trait SchedulerCurrentVersionProvider {
    /// Returns the current version of the data in the store.
    fn current_version(&self) -> i64;
}

/// Default implementation for when we are not providing a version
/// for the data.
impl<T> SchedulerCurrentVersionProvider for T
where
    T: SchedulerStoreKeyProvider<Versioned = FalseValue>,
{
    fn current_version(&self) -> i64 {
        0
    }
}

/// Compile time types for booleans.
pub trait BoolValue {
    const VALUE: bool;
}
/// Compile time check if something is false.
pub trait IsFalse {}
/// Compile time check if something is true.
pub trait IsTrue {}

/// Compile time true value.
#[derive(Debug, Clone, Copy)]
pub struct TrueValue;
impl BoolValue for TrueValue {
    const VALUE: bool = true;
}
impl IsTrue for TrueValue {}

/// Compile time false value.
#[derive(Debug, Clone, Copy)]
pub struct FalseValue;
impl BoolValue for FalseValue {
    const VALUE: bool = false;
}
impl IsFalse for FalseValue {}

Coverage Report

Created: 2026-06-29 16:10