// 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::convert::Into;

use core::fmt::{Debug, Formatter};

use core::pin::Pin;

use core::sync::atomic::{AtomicU64, Ordering};

use core::time::Duration;

use std::collections::hash_map::Entry;

use std::collections::{HashMap, HashSet};

use std::sync::Arc;

use std::time::{Instant, SystemTime, UNIX_EPOCH};

use bytes::BytesMut;

use futures::future::pending;

use futures::stream::unfold;

use futures::{Future, Stream, TryFutureExt, try_join};

use nativelink_config::cas_server::{ByteStreamConfig, InstanceName, WithInstanceName};

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

use nativelink_metric::{

    MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, group, publish,

};

use nativelink_proto::build::bazel::remote::execution::v2::compressor;

use nativelink_proto::google::bytestream::byte_stream_server::{

    ByteStream, ByteStreamServer as Server,

};

use nativelink_proto::google::bytestream::{

    QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, ReadResponse, WriteRequest,

    WriteResponse,

};

use nativelink_store::grpc_store::GrpcStore;

use nativelink_store::store_manager::StoreManager;

use nativelink_util::buf_channel::{

    DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair,

};

use nativelink_util::common::DigestInfo;

use nativelink_util::digest_hasher::{

    DigestHasherFunc, default_digest_hasher_func, make_ctx_for_hash_func,

};

use nativelink_util::proto_stream_utils::WriteRequestStreamWrapper;

use nativelink_util::resource_info::ResourceInfo;

use nativelink_util::spawn;

use nativelink_util::store_trait::{Store, StoreLike, StoreOptimizations, UploadSizeInfo};

use nativelink_util::task::JoinHandleDropGuard;

use opentelemetry::context::FutureExt;

use parking_lot::Mutex;

use tokio::time::sleep;

use tonic::{Request, Response, Status, Streaming};

use tracing::{Instrument, Level, debug, error, error_span, info, instrument, trace, warn};

use crate::wire_compression::RemoteCacheCompressionInstances;

/// If this value changes update the documentation in the config definition.

const DEFAULT_PERSIST_STREAM_ON_DISCONNECT_TIMEOUT: Duration = Duration::from_mins(1);

/// If this value changes update the documentation in the config definition.

const DEFAULT_MAX_BYTES_PER_STREAM: usize = 64 * 1024;

/// Metrics for `ByteStream` server operations.

/// Tracks upload/download activity, throughput, and latency.

#[derive(Debug, Default)]

pub struct ByteStreamMetrics {

    /// Number of currently active uploads (includes idle streams waiting for resume)

    pub active_uploads: AtomicU64,

    /// Total number of write requests received

    pub write_requests_total: AtomicU64,

    /// Total number of successful write requests

    pub write_requests_success: AtomicU64,

    /// Total number of failed write requests

    pub write_requests_failure: AtomicU64,

    /// Total number of read requests received

    pub read_requests_total: AtomicU64,

    /// Total number of successful read requests

    pub read_requests_success: AtomicU64,

    /// Total number of failed read requests

    pub read_requests_failure: AtomicU64,

    /// Total number of `query_write_status` requests

    pub query_write_status_total: AtomicU64,

    /// Total bytes written via `ByteStream`

    pub bytes_written_total: AtomicU64,

    /// Total bytes read via `ByteStream`

    pub bytes_read_total: AtomicU64,

    /// Sum of write durations in nanoseconds (for average latency calculation)

    pub write_duration_ns: AtomicU64,

    /// Sum of read durations in nanoseconds (for average latency calculation)

    pub read_duration_ns: AtomicU64,

    /// Number of UUID collisions detected

    pub uuid_collisions: AtomicU64,

    /// Number of resumed uploads (client reconnected to existing stream)

    pub resumed_uploads: AtomicU64,

    /// Number of idle streams that timed out

    pub idle_stream_timeouts: AtomicU64,

}

impl MetricsComponent for ByteStreamMetrics {

    fn publish(

        &self,

        _kind: MetricKind,

        field_metadata: MetricFieldData,

    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {

        let _enter = group!(field_metadata.name).entered();

        publish!(

            "active_uploads",

            &self.active_uploads,

            MetricKind::Counter,

            "Number of currently active uploads"

        );

        publish!(

            "write_requests_total",

            &self.write_requests_total,

            MetricKind::Counter,

            "Total write requests received"

        );

        publish!(

            "write_requests_success",

            &self.write_requests_success,

            MetricKind::Counter,

            "Total successful write requests"

        );

        publish!(

            "write_requests_failure",

            &self.write_requests_failure,

            MetricKind::Counter,

            "Total failed write requests"

        );

        publish!(

            "read_requests_total",

            &self.read_requests_total,

            MetricKind::Counter,

            "Total read requests received"

        );

        publish!(

            "read_requests_success",

            &self.read_requests_success,

            MetricKind::Counter,

            "Total successful read requests"

        );

        publish!(

            "read_requests_failure",

            &self.read_requests_failure,

            MetricKind::Counter,

            "Total failed read requests"

        );

        publish!(

            "query_write_status_total",

            &self.query_write_status_total,

            MetricKind::Counter,

            "Total query_write_status requests"

        );

        publish!(

            "bytes_written_total",

            &self.bytes_written_total,

            MetricKind::Counter,

            "Total bytes written via ByteStream"

        );

        publish!(

            "bytes_read_total",

            &self.bytes_read_total,

            MetricKind::Counter,

            "Total bytes read via ByteStream"

        );

        publish!(

            "write_duration_ns",

            &self.write_duration_ns,

            MetricKind::Counter,

            "Sum of write durations in nanoseconds"

        );

        publish!(

            "read_duration_ns",

            &self.read_duration_ns,

            MetricKind::Counter,

            "Sum of read durations in nanoseconds"

        );

        publish!(

            "uuid_collisions",

            &self.uuid_collisions,

            MetricKind::Counter,

            "Number of UUID collisions detected"

        );

        publish!(

            "resumed_uploads",

            &self.resumed_uploads,

            MetricKind::Counter,

            "Number of resumed uploads"

        );

        publish!(

            "idle_stream_timeouts",

            &self.idle_stream_timeouts,

            MetricKind::Counter,

            "Number of idle streams that timed out"

        );

        Ok(MetricPublishKnownKindData::Component)

    }

}

type BytesWrittenAndIdleStream = (Arc<AtomicU64>, Option<IdleStream>);

/// Type alias for the UUID key used in `active_uploads` `HashMap`.

/// Using u128 instead of String reduces memory allocations and improves

/// cache locality for `HashMap` operations.

type UuidKey = u128;

/// Parse a UUID string to a u128 for use as a `HashMap` key.

/// This avoids heap allocation for String keys and improves `HashMap` performance.

/// Falls back to hashing the string if it's not a valid hex UUID.

#[inline]

fn parse_uuid_to_key(uuid_str: &str) -> UuidKey {

    // UUIDs are typically 32 hex chars (128 bits) or 36 chars with dashes.

    // We'll try to parse as hex first, then fall back to hashing.

    let clean: String = uuid_str.chars().filter(char::is_ascii_hexdigit).collect();

    if clean.len() >= 16 {

        // Take up to 32 hex chars (128 bits)

        let hex_str = if clean.len() > 32 {

            &clean[..32]

        } else {

            &clean

        };

        u128::from_str_radix(hex_str, 16).unwrap_or_else(|_| 
{0

            // Hash fallback for non-hex strings

            use core::hash::{Hash, Hasher};

            let mut hasher = std::collections::hash_map::DefaultHasher::new();

            uuid_str.hash(&mut hasher);

            u128::from(hasher.finish())

        })

    } else {

        // Short strings: use hash

        use core::hash::{Hash, Hasher};

        let mut hasher = std::collections::hash_map::DefaultHasher::new();

        uuid_str.hash(&mut hasher);

        u128::from(hasher.finish())

    }

}

pub struct InstanceInfo {

    store: Store,

    // Max number of bytes to send on each grpc stream chunk.

    max_bytes_per_stream: usize,

    /// Active uploads keyed by UUID as u128 for better performance.

    /// Using u128 keys instead of String reduces heap allocations

    /// and improves `HashMap` lookup performance.

    active_uploads: Arc<Mutex<HashMap<UuidKey, BytesWrittenAndIdleStream>>>,

    /// How long to keep idle streams before timing them out.

    idle_stream_timeout: Duration,

    metrics: Arc<ByteStreamMetrics>,

    /// Handle to the global sweeper task. Kept alive for the lifetime of the instance.

    _sweeper_handle: Arc<JoinHandleDropGuard<()>>,

    /// Whether this instance supports Bazel remote cache compression.

    remote_cache_compression_enabled: bool,

}

impl Debug for InstanceInfo {

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

        f.debug_struct("InstanceInfo")

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

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

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

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

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

            .field(

                "remote_cache_compression_enabled",

                &self.remote_cache_compression_enabled,

            )

            .finish()

    }

}

struct CompressedUploadGuard {

    uuid_key: u128,

    active_uploads: Arc<Mutex<HashMap<UuidKey, BytesWrittenAndIdleStream>>>,

    metrics: Arc<ByteStreamMetrics>,

}

impl Drop for CompressedUploadGuard {

    fn drop(&mut self) {

        self.active_uploads.lock().remove(&self.uuid_key);

        self.metrics.active_uploads.fetch_sub(1, Ordering::Relaxed);

    }

}

impl InstanceInfo {

    fn track_compressed_upload(

        &self,

        uuid_key: UuidKey,

    ) -> (Arc<AtomicU64>, CompressedUploadGuard) {

        let bytes_received = Arc::new(AtomicU64::new(0));

        let uuid_key = {

            let mut active_uploads = self.active_uploads.lock();

            match active_uploads.entry(uuid_key) {

                Entry::Occupied(entry) => {

                    // Another upload already owns this UUID; rekey instead of

                    // clobbering its entry (which would break its QueryWriteStatus

                    // visibility and double-decrement the metric on guard drop).

                    // Mirrors create_or_join_upload_stream's collision handling.

                    let original_key = *entry.key();

                    let unique_key = ByteStreamServer::generate_unique_uuid_key(original_key);

                    warn!(

                        msg = "UUID collision detected on compressed upload, generating unique UUID to prevent conflict",

                        original_uuid = format!("{original_key:032x}"),

                        unique_uuid = format!("{unique_key:032x}")

                    );

                    self.metrics.uuid_collisions.fetch_add(1, Ordering::Relaxed);

                    // Release the Occupied entry's borrow so we can insert on the same guard.

                    let _ = entry;

                    active_uploads.insert(unique_key, (bytes_received.clone(), None));

                    unique_key

                }

                Entry::Vacant(entry) => {

                    let key = *entry.key();

                    entry.insert((bytes_received.clone(), None));

                    key

                }

            }

        };

        self.metrics.active_uploads.fetch_add(1, Ordering::Relaxed);

        (

            bytes_received,

            CompressedUploadGuard {

                uuid_key,

                active_uploads: self.active_uploads.clone(),

                metrics: self.metrics.clone(),

            },

        )

    }

}

/// Pump compressed `ByteStream` upload chunks into the decoder.

///

/// Compressed uploads intentionally do not support the identity upload resume

/// protocol; offsets are validated against compressed wire bytes.

async fn process_compressed_client_stream(

    mut stream: WriteRequestStreamWrapper<impl Stream<Item = Result<WriteRequest, Status>> + Unpin>,

    mut tx: DropCloserWriteHalf,

    bytes_received: &Arc<AtomicU64>,

) -> Result<(), Error> {

    loop {

        match stream.next().await {

            Some(Ok(write_request)) => {

                if write_request.write_offset < 0 {

                    return Err(make_input_err!(

                        "Invalid negative compressed write offset in write request: {}",

                        write_request.write_offset

                    ));

                }

                let write_offset = u64::try_from(write_request.write_offset)

                    .err_tip(|| "Compressed write offset was not convertible to u64")
?0
;

                let compressed_offset = tx.get_bytes_written();

                if write_offset != compressed_offset {

                    return Err(make_input_err!(

                        "Received out of order compressed data. Got {}, expected {}",

                        write_offset,

                        compressed_offset

                    ));

                }

                if !write_request.data.is_empty() {

                    tx.send(write_request.data)

                        .await

                        .err_tip(|| "Failed to forward compressed upload data")
?0
;

                    bytes_received.store(tx.get_bytes_written(), Ordering::Release);

                }

                if write_request.finish_write {

                    tx.send_eof()

                        .err_tip(|| "Failed to send compressed upload EOF")
?0
;

                    return Ok(());

                }

            }

            Some(Err(e)) => {

                return Err(e);

            }

            None => {

                return Err(make_err!(

                    Code::InvalidArgument,

                    "Compressed write stream ended without finish_write"

                ));

            }

        }

    }

}

type ReadStream = Pin<Box<dyn Stream<Item = Result<ReadResponse, Status>> + Send + 'static>>;

type StoreUpdateFuture = Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'static>>;

struct StreamState {

    uuid: UuidKey,

    tx: DropCloserWriteHalf,

    store_update_fut: StoreUpdateFuture,

    digest_function: DigestHasherFunc,

}

impl Debug for StreamState {

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

        f.debug_struct("StreamState")

            .field("uuid", &format!("{:032x}", self.uuid))

            .finish()

    }

}

/// If a stream is in this state, it will automatically be put back into an `IdleStream` and

/// placed back into the `active_uploads` map as an `IdleStream` after it is dropped.

/// To prevent it from being put back into an `IdleStream` you must call `.graceful_finish()`.

struct ActiveStreamGuard {

    stream_state: Option<StreamState>,

    bytes_received: Arc<AtomicU64>,

    active_uploads: Arc<Mutex<HashMap<UuidKey, BytesWrittenAndIdleStream>>>,

    metrics: Arc<ByteStreamMetrics>,

}

impl ActiveStreamGuard {

    /// Consumes the guard. The stream will be considered "finished", will

    /// remove it from the `active_uploads`.

    fn graceful_finish(mut self) {

        let stream_state = self.stream_state.take().unwrap();

        self.active_uploads.lock().remove(&stream_state.uuid);

        // Decrement active uploads counter on successful completion

        self.metrics.active_uploads.fetch_sub(1, Ordering::Relaxed);

    }

}

impl Drop for ActiveStreamGuard {

    fn drop(&mut self) {

        let Some(
stream_state7
) = self.stream_state.take() else {

            return; // If None it means we don't want it put back into an IdleStream.

        };

        let mut active_uploads = self.active_uploads.lock();

        let uuid = stream_state.uuid; // u128 is Copy, no clone needed

        let Some(active_uploads_slot) = active_uploads.get_mut(&uuid) else {

            error!(

                err = "Failed to find active upload. This should never happen.",

                uuid = format!("{:032x}", uuid),

            );

            return;

        };

        // Mark stream as idle with current timestamp.

        // The global sweeper will clean it up after idle_stream_timeout.

        // This avoids spawning a task per stream, reducing overhead from O(n) to O(1).

        active_uploads_slot.1 = Some(IdleStream {

            stream_state,

            idle_since: Instant::now(),

        });

    }

}

/// Represents a stream that is in the "idle" state. this means it is not currently being used

/// by a client. If it is not used within a certain amount of time it will be removed from the

/// `active_uploads` map automatically by the global sweeper task.

#[derive(Debug)]

struct IdleStream {

    stream_state: StreamState,

    /// When this stream became idle. Used by the global sweeper to determine expiration.

    idle_since: Instant,

}

impl IdleStream {

    fn into_active_stream(

        self,

        bytes_received: Arc<AtomicU64>,

        instance_info: &InstanceInfo,

    ) -> ActiveStreamGuard {

        ActiveStreamGuard {

            stream_state: Some(self.stream_state),

            bytes_received,

            active_uploads: instance_info.active_uploads.clone(),

            metrics: instance_info.metrics.clone(),

        }

    }

}

#[derive(Debug)]

pub struct ByteStreamServer {

    instance_infos: HashMap<InstanceName, InstanceInfo>,

    reported_digest_function_mismatches: Mutex<HashSet<(DigestHasherFunc, DigestHasherFunc)>>,

}

impl ByteStreamServer {

    /// Generate a unique UUID key by `XOR`ing the base key with a nanosecond timestamp.

    /// This ensures virtually zero collision probability while being O(1).

    fn generate_unique_uuid_key(base_key: UuidKey) -> UuidKey {

        let timestamp = SystemTime::now()

            .duration_since(UNIX_EPOCH)

            .unwrap_or_default()

            .as_nanos();

        // XOR with timestamp to create unique key

        base_key ^ timestamp

    }

    pub fn new(

        configs: &[WithInstanceName<ByteStreamConfig>],

        store_manager: &StoreManager,

        remote_cache_compression_instances: &RemoteCacheCompressionInstances,

    ) -> Result<Self, Error> {

        let mut instance_infos: HashMap<String, InstanceInfo> = HashMap::new();

        for config in configs {

            let idle_stream_timeout = if config.persist_stream_on_disconnect_timeout_s == 0 {

                DEFAULT_PERSIST_STREAM_ON_DISCONNECT_TIMEOUT

            } else {

                Duration::from_secs(config.persist_stream_on_disconnect_timeout_s as u64)

            };

            let remote_cache_compression_enabled =

                remote_cache_compression_instances.enabled_for(&config.instance_name);

            let _old_value = instance_infos.insert(

                config.instance_name.clone(),

                Self::new_with_timeout(

                    config,

                    store_manager,

                    idle_stream_timeout,

                    remote_cache_compression_enabled,

                )?,

            );

        }

        Ok(Self {

            instance_infos,

            reported_digest_function_mismatches: Mutex::new(HashSet::new()),

        })

    }

    pub fn new_with_timeout(

        config: &WithInstanceName<ByteStreamConfig>,

        store_manager: &StoreManager,

        idle_stream_timeout: Duration,

        remote_cache_compression_enabled: bool,

    ) -> Result<InstanceInfo, Error> {

        let store = store_manager

            .get_store(&config.cas_store)

            .ok_or_else(|| 
make_input_err!0
("'cas_store': '{}' does not exist", 
config.cas_store0
))
?0
;

        let max_bytes_per_stream = if config.max_bytes_per_stream == 0 {

            DEFAULT_MAX_BYTES_PER_STREAM

        } else {

            config.max_bytes_per_stream

        };

        let active_uploads: Arc<Mutex<HashMap<UuidKey, BytesWrittenAndIdleStream>>> =

            Arc::new(Mutex::new(HashMap::new()));

        let metrics = Arc::new(ByteStreamMetrics::default());

        // Spawn a single global sweeper task that periodically cleans up expired idle streams.

        // This replaces per-stream timeout tasks, reducing task spawn overhead from O(n) to O(1).

        let sweeper_active_uploads = Arc::downgrade(&active_uploads);

        let sweeper_metrics = Arc::downgrade(&metrics);

        let sweep_interval = idle_stream_timeout / 2; // Check every half-timeout period

        let sweeper_handle = spawn!("bytestream_idle_stream_sweeper", async move 
{30

            loop {

                sleep(sweep_interval).await;

                let Some(active_uploads) = sweeper_active_uploads.upgrade() else {

                    // InstanceInfo has been dropped, exit the sweeper

                    break;

                };

                let metrics = sweeper_metrics.upgrade();

                let now = Instant::now();

                let mut expired_count = 0u64;

                // Lock and sweep expired entries

                {

                    let mut uploads = active_uploads.lock();

                    uploads.retain(|uuid, (_, maybe_idle)| {

                        if let Some(idle_stream) = maybe_idle

                            && now.duration_since(idle_stream.idle_since) >= idle_stream_timeout

                        {

                            info!(

                                msg = "Sweeping expired idle stream",

                                uuid = format!("{:032x}", uuid)

                            );

                            expired_count += 1;

                            return false; // Remove this entry

                        }

                        true // Keep this entry

                    });

                }

                // Update metrics outside the lock

                if expired_count > 0 {

                    if let Some(m) = &metrics {

                        m.idle_stream_timeouts

                            .fetch_add(expired_count, Ordering::Relaxed);

                        m.active_uploads.fetch_sub(expired_count, Ordering::Relaxed);

                    }

                    trace!(

                        msg = "Sweeper cleaned up expired streams",

                        count = expired_count

                    );

                }

            }

        });

        Ok(InstanceInfo {

            store,

            max_bytes_per_stream,

            active_uploads,

            idle_stream_timeout,

            metrics,

            _sweeper_handle: Arc::new(sweeper_handle),

            remote_cache_compression_enabled,

        })

    }

    pub fn into_service(self) -> Server<Self> {

        Server::new(self)

    }

    /// Creates or joins an upload stream for the given UUID.

    ///

    /// This function handles three scenarios:

    /// 1. UUID doesn't exist - creates a new upload stream

    /// 2. UUID exists but is idle - resumes the existing stream

    /// 3. UUID exists and is active - generates a unique UUID by appending a nanosecond

    ///    timestamp to avoid collision, then creates a new stream with that UUID

    ///

    /// The nanosecond timestamp ensures virtually zero probability of collision since

    /// two concurrent uploads would need to both collide on the original UUID AND

    /// generate the unique UUID in the exact same nanosecond.

    fn create_or_join_upload_stream(

        &self,

        uuid_str: &str,

        instance: &InstanceInfo,

        digest: DigestInfo,

        digest_function: DigestHasherFunc,

    ) -> Result<ActiveStreamGuard, Error> {

        // Bind the digest function to the retained store future itself. The

        // future can outlive this RPC while an upload is idle, so it must not

        // depend on the context of whichever request polls it next.

        let store_update_context = make_ctx_for_hash_func(digest_function)
?0
;

        // Parse UUID string to u128 key for efficient HashMap operations

        let uuid_key = parse_uuid_to_key(uuid_str);

        let (uuid, bytes_received, is_collision) = {

            let mut active_uploads = instance.active_uploads.lock();

            match active_uploads.entry(uuid_key) {

                Entry::Occupied(mut entry) => {

                    let maybe_idle_stream = entry.get_mut();

                    if let Some(
idle_stream4
) = maybe_idle_stream.1.as_ref()

                        && idle_stream.stream_state.digest_function != digest_function

                    {

                        return Err(make_input_err!(

                            "Cannot resume upload with digest function {} because it started with {}",

                            digest_function,

                            idle_stream.stream_state.digest_function,

                        ));

                    }

                    if let Some(
idle_stream4
) = maybe_idle_stream.1.take() {

                        // Case 2: Stream exists but is idle, we can resume it

                        let bytes_received = maybe_idle_stream.0.clone();

                        info!(

                            msg = "Joining existing stream",

                            uuid = format!("{:032x}", entry.key())

                        );

                        // Track resumed upload

                        instance

                            .metrics

                            .resumed_uploads

                            .fetch_add(1, Ordering::Relaxed);

                        return Ok(idle_stream.into_active_stream(bytes_received, instance));

                    }

                    // Case 3: Stream is active - generate a unique UUID to avoid collision.

                    // Using nanosecond timestamp makes collision probability essentially zero.

                    let original_key = *entry.key();

                    let unique_key = Self::generate_unique_uuid_key(original_key);

                    warn!(

                        msg = "UUID collision detected, generating unique UUID to prevent conflict",

                        original_uuid = format!("{:032x}", original_key),

                        unique_uuid = format!("{:032x}", unique_key)

                    );

                    // Release the Occupied entry's borrow so we can insert on the same guard.

                    let _ = entry;

                    let bytes_received = Arc::new(AtomicU64::new(0));

                    active_uploads.insert(unique_key, (bytes_received.clone(), None));

                    (unique_key, bytes_received, true)

                }

                Entry::Vacant(entry) => {

                    // Case 1: UUID doesn't exist, create new stream

                    let bytes_received = Arc::new(AtomicU64::new(0));

                    let uuid = *entry.key();

                    // Our stream is "in use" if the key is in the map, but the value is None.

                    entry.insert((bytes_received.clone(), None));

                    (uuid, bytes_received, false)

                }

            }

        };

        // Track metrics for new upload

        instance

            .metrics

            .active_uploads

            .fetch_add(1, Ordering::Relaxed);

        if is_collision {

            instance

                .metrics

                .uuid_collisions

                .fetch_add(1, Ordering::Relaxed);

        }

        // Important: Do not return an error from this point onwards without

        // removing the entry from the map, otherwise that UUID becomes

        // unusable.

        let (tx, rx) = make_buf_channel_pair();

        let store = instance.store.clone();

        let store_update_fut = Box::pin(

            async move {

                // We need to wrap `Store::update()` in a another future because we need to capture

                // `store` to ensure its lifetime follows the future and not the caller.

                store

                    // Bytestream always uses digest size as the actual byte size.

                    .update(digest, rx, UploadSizeInfo::ExactSize(digest.size_bytes()))

                    .await

                    .map(|_| ())

            }

            .with_context(store_update_context),

        );

        Ok(ActiveStreamGuard {

            stream_state: Some(StreamState {

                uuid,

                tx,

                store_update_fut,

                digest_function,

            }),

            bytes_received,

            active_uploads: instance.active_uploads.clone(),

            metrics: instance.metrics.clone(),

        })

    }

    async fn inner_read(

        &self,

        instance: &InstanceInfo,

        digest: DigestInfo,

        read_request: ReadRequest,

    ) -> Result<impl Stream<Item = Result<ReadResponse, Status>> + Send + use<>, Error> {

        struct ReaderState {

            max_bytes_per_stream: usize,

            rx: DropCloserReadHalf,

            maybe_get_part_result: Option<Result<(), Error>>,

            get_part_fut: Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>,

        }

        let read_limit = u64::try_from(read_request.read_limit)

            .err_tip(|| "Could not convert read_limit to u64")
?0
;

        let (tx, rx) = make_buf_channel_pair();

        let read_limit = if read_limit != 0 {

            Some(read_limit)

        } else {

            None

        };

        // This allows us to call a destructor when the the object is dropped.

        let store = instance.store.clone();

        let state = Some(ReaderState {

            rx,

            max_bytes_per_stream: instance.max_bytes_per_stream,

            maybe_get_part_result: None,

            get_part_fut: Box::pin(async move {

                store

                    .get_part(

                        digest,

                        tx,

                        u64::try_from(read_request.read_offset)

                            .err_tip(|| "Could not convert read_offset to u64")
?0
,

                        read_limit,

                    )

                    .await

            }),

        });

        let read_stream_span = error_span!("read_stream");

        Ok(
Box::pin8
(
unfold8
(
state8
, move |state| {

            async {

            let mut state = state
?0
; // If None our stream is done.

            let mut response = ReadResponse::default();

            {

                let consume_fut = state.rx.consume(Some(state.max_bytes_per_stream));

                tokio::pin!(consume_fut);

                loop {

                    tokio::select! {

                        
read_result9.79k
 = &mut consume_fut => {

                            match read_result {

                                Ok(bytes) => {

                                    if bytes.is_empty() {

                                        // EOF.

                                        return None;

                                    }

                                    if bytes.len() > state.max_bytes_per_stream {

                                        let err = make_err!(Code::Internal, "Returned store size was larger than read size");

                                        return Some((Err(err.into()), None));

                                    }

                                    response.data = bytes;

                                    trace!(response.data = format!("<redacted len({})>", response.data.len()));

                                    break;

                                }

                                Err(mut e) => {

                                    // We may need to propagate the error from reading the data through first.

                                    // For example, the NotFound error will come through `get_part_fut`, and

                                    // will not be present in `e`, but we need to ensure we pass NotFound error

                                    // code or the client won't know why it failed.

                                    let get_part_result = if let Some(result) = state.maybe_get_part_result {

                                        result

                                    } else {

                                        // This should never be `future::pending()` if maybe_get_part_result is

                                        // not set.

                                        state.get_part_fut.await

                                    };

                                    if let Err(err) = get_part_result {

                                        e = err.merge(e);

                                    
}0

                                    if e.code == Code::NotFound {

                                        // Trim the error code. Not Found is quite common and we don't want to send a large

                                        // error (debug) message for something that is common. We resize to just the last

                                        // message as it will be the most relevant.

                                        e.messages.truncate(1);

                                    
}0

                                    error!(response = ?e);

                                    return Some((Err(e.into()), None))

                                }

                            }

                        },

                        
result8
 = &mut state.get_part_fut => {

                            state.maybe_get_part_result = Some(result);

                            // It is non-deterministic on which future will finish in what order.

                            // It is also possible that the `state.rx.consume()` call above may not be able to

                            // respond even though the publishing future is done.

                            // Because of this we set the writing future to pending so it never finishes.

                            // The `state.rx.consume()` future will eventually finish and return either the

                            // data or an error.

                            // An EOF will terminate the `state.rx.consume()` future, but we are also protected

                            // because we are dropping the writing future, it will drop the `tx` channel

                            // which will eventually propagate an error to the `state.rx.consume()` future if

                            // the EOF was not sent due to some other error.

                            state.get_part_fut = Box::pin(pending());

                        },

                    }

                }

            }

            Some((Ok(response), Some(state)))

        }.instrument(read_stream_span.clone())

        })))

    }

    // We instrument tracing here as well as below because `stream` has a hash on it

    // that is extracted from the first stream message. If we only implemented it below

    // we would not have the hash available to us.

    #[instrument(

        ret(level = Level::DEBUG),

        level = Level::ERROR,

        skip(self, instance_info),

    )]

    async fn inner_write(

        &self,

        instance_info: &InstanceInfo,

        digest: DigestInfo,

        digest_function: DigestHasherFunc,

        stream: WriteRequestStreamWrapper<impl Stream<Item = Result<WriteRequest, Status>> + Unpin>,

    ) -> Result<Response<WriteResponse>, Error> {

        async fn process_client_stream(

            mut stream: WriteRequestStreamWrapper<

                impl Stream<Item = Result<WriteRequest, Status>> + Unpin,

            >,

            tx: &mut DropCloserWriteHalf,

            outer_bytes_received: &Arc<AtomicU64>,

            expected_size: u64,

        ) -> Result<(), Error> {

            loop {

                let 
write_request27
 = match stream.next().await {

                    // Code path for when client tries to gracefully close the stream.

                    // If this happens it means there's a problem with the data sent,

                    // because we always close the stream from our end before this point

                    // by counting the number of bytes sent from the client. If they send

                    // less than the amount they said they were going to send and then

                    // close the stream, we know there's a problem.

                    None => {

                        return Err(make_input_err!(

                            "Client closed stream before sending all data"

                        ));

                    }

                    // Code path for client stream error. Probably client disconnect.

                    Some(Err(err)) => return Err(err),

                    // Code path for received chunk of data.

                    Some(Ok(write_request)) => write_request,

                };

                if write_request.write_offset < 0 {

                    return Err(make_input_err!(

                        "Invalid negative write offset in write request: {}",

                        write_request.write_offset

                    ));

                }

                let write_offset = write_request.write_offset as u64;

                // If we get duplicate data because a client didn't know where

                // it left off from, then we can simply skip it.

                let 
data26
 = if write_offset < tx.get_bytes_written() {

                    if (write_offset + write_request.data.len() as u64) < tx.get_bytes_written() {

                        if write_request.finish_write {

                            return Err(make_input_err!(

                                "Resumed stream finished at {} bytes when we already received {} bytes.",

                                write_offset + write_request.data.len() as u64,

                                tx.get_bytes_written()

                            ));

                        }

                        continue;

                    }

                    write_request.data.slice(

                        usize::try_from(tx.get_bytes_written() - write_offset)

                            .unwrap_or(usize::MAX)..,

                    )

                } else {

                    if write_offset != tx.get_bytes_written() {

                        return Err(make_input_err!(

                            "Received out of order data. Got {}, expected {}",

                            write_offset,

                            tx.get_bytes_written()

                        ));

                    }

                    write_request.data

                };

                // Do not process EOF or weird stuff will happen.

                if !data.is_empty() {

                    // We also need to process the possible EOF branch, so we can't early return.

                    if let Err(
mut err0
) = tx.send(data).await {

                        err.code = Code::Internal;

                        return Err(err);

                    }

                    outer_bytes_received.store(tx.get_bytes_written(), Ordering::Release);

                }

                if expected_size < tx.get_bytes_written() {

                    return Err(make_input_err!("Received more bytes than expected"));

                }

                if write_request.finish_write {

                    // Gracefully close our stream.

                    tx.send_eof()

                        .err_tip(|| "Failed to send EOF in ByteStream::write")
?0
;

                    return Ok(());

                }

                // Continue.

            }

            // Unreachable.

        }

        let uuid = stream

            .resource_info

            .uuid

            .as_ref()

            .ok_or_else(|| make_input_err!("UUID must be set if writing data"))?;

        let mut active_stream_guard =

            self.create_or_join_upload_stream(uuid, instance_info, digest, digest_function)?;

        let expected_size = stream.resource_info.expected_size as u64;

        let active_stream = active_stream_guard.stream_state.as_mut().unwrap();

        try_join!(

            process_client_stream(

                stream,

                &mut active_stream.tx,

                &active_stream_guard.bytes_received,

                expected_size

            ),

            (&mut active_stream.store_update_fut)

                .map_err(|err| { err.append("Error updating inner store") })

        )?;

        // Close our guard and consider the stream no longer active.

        active_stream_guard.graceful_finish();

        Ok(Response::new(WriteResponse {

            committed_size: expected_size.try_into().unwrap_or(i64::MAX),

        }))

    }

    /// Fast-path write that bypasses channel overhead for stores that support direct Bytes updates.

    /// This buffers all data in memory and calls `update_oneshot` directly.

    async fn inner_write_oneshot(

        &self,

        instance_info: &InstanceInfo,

        digest: DigestInfo,

        digest_function: DigestHasherFunc,

        mut stream: WriteRequestStreamWrapper<

            impl Stream<Item = Result<WriteRequest, Status>> + Unpin,

        >,

    ) -> Result<Response<WriteResponse>, Error> {

        let expected_size = stream.resource_info.expected_size as u64;

        // Pre-allocate buffer for expected size (capped at reasonable limit to prevent DoS)

        let capacity =

            usize::try_from(expected_size.min(64 * 1024 * 1024)).unwrap_or(64 * 1024 * 1024);

        let mut buffer = BytesMut::with_capacity(capacity);

        let mut bytes_received: u64 = 0;

        // Collect all data from client stream

        loop {

            let 
write_request4
 = match stream.next().await {

                None => {

                    return Err(make_input_err!(

                        "Client closed stream before sending all data"

                    ));

                }