// Copyright 2024 Trace Machina, Inc. 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::pin::Pin;

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

use core::time::Duration;

use std::borrow::Cow;

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

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

use async_trait::async_trait;

use bytes::{Bytes, BytesMut};

use futures::stream::{FuturesUnordered, unfold};

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

use nativelink_config::stores::{GrpcReadBatchingConfig, GrpcSpec};

use nativelink_error::{Error, ResultExt, error_if, make_err};

use nativelink_metric::MetricsComponent;

use nativelink_proto::build::bazel::remote::execution::v2::action_cache_client::ActionCacheClient;

use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_client::ContentAddressableStorageClient;

use nativelink_proto::build::bazel::remote::execution::v2::{

    ActionResult, BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest,

    BatchUpdateBlobsResponse, FindMissingBlobsRequest, FindMissingBlobsResponse,

    GetActionResultRequest, GetTreeRequest, GetTreeResponse, SpliceBlobRequest, SpliceBlobResponse,

    SplitBlobRequest, SplitBlobResponse, UpdateActionResultRequest, batch_update_blobs_request,

    compressor,

};

use nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient;

use nativelink_proto::google::bytestream::{

    QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, ReadResponse, WriteRequest,

    WriteResponse,

};

use nativelink_util::buf_channel::{

    DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair,

};

use nativelink_util::common::DigestInfo;

use nativelink_util::connection_manager::ConnectionManager;

use nativelink_util::digest_hasher::{DigestHasherFunc, default_digest_hasher_func};

use nativelink_util::health_utils::HealthStatusIndicator;

use nativelink_util::proto_stream_utils::{

    FirstStream, MAX_WRITE_REQUEST_DATA_BYTES, WriteRequestStreamWrapper, WriteState,

    WriteStateWrapper,

};

use nativelink_util::resource_info::ResourceInfo;

use nativelink_util::retry::{Retrier, RetryResult};

use nativelink_util::store_trait::{RemoveCallback, StoreDriver, StoreKey, UploadSizeInfo};

use nativelink_util::telemetry::{forward_client_headers, inject_context};

use nativelink_util::wire_compression::{

    stream_decode_compressed_upload, stream_encode_compressed_upload_from_reader,

};

use nativelink_util::{background_spawn, default_health_status_indicator, tls_utils};

use opentelemetry::context::Context;

use parking_lot::Mutex;

use prost::Message;

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

use tokio::time::sleep;

use tonic::metadata::{Ascii, MetadataKey, MetadataValue};

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

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

use uuid::Uuid;

/// Adds configured static headers, forwards nominated client request headers,

/// and injects the current OpenTelemetry trace context into an outgoing gRPC

/// request.

fn enrich_request<T>(

    mut request: Request<T>,

    headers: &[(MetadataKey<Ascii>, MetadataValue<Ascii>)],

    forward_headers: &[String],

) -> Request<T> {

    for (
key1
, 
value1
) in headers {

        request.metadata_mut().insert(key.clone(), value.clone());

    }

    forward_client_headers(&mut request, forward_headers);

    inject_context(&mut request);

    request

}

/// Minimum blob size for wire compression when

/// `experimental_remote_cache_compression` is enabled. Below this, zstd

/// framing overhead and per-blob CPU outweigh the wire savings (small blobs

/// are the batching paths' domain).

const WIRE_COMPRESSION_MIN_SIZE_BYTES: u64 = 64 * 1024;

/// Zstd level for `GrpcStore`'s own compressed transfers. Level 1: measured on

/// artifact-shaped corpora, higher levels bought no meaningful wire reduction

/// while encoding slower, and internal hops are throughput-sensitive.

const WIRE_COMPRESSION_ZSTD_LEVEL: i32 = 1;

/// Estimated per-entry protobuf and framing overhead charged against

/// `max_batch_bytes`, so that batches of many tiny blobs cannot push a

/// `BatchReadBlobs` response over the gRPC message size limit.

const BATCH_READ_PER_ENTRY_OVERHEAD_BYTES: u64 = 256;

/// A small-blob read waiting to be coalesced into a `BatchReadBlobs` RPC.

#[derive(Debug)]

struct PendingRead {

    digest: DigestInfo,

    digest_function: i32,

    tx: oneshot::Sender<Result<Bytes, Error>>,

}

/// The pending small-blob reads plus the total payload bytes they declare.

#[derive(Debug, Default)]

struct ReadQueue {

    items: VecDeque<PendingRead>,

    bytes: u64,

}

/// State for coalescing small-blob reads into `BatchReadBlobs` RPCs.

///

/// This uses a slot-based group commit scheme: callers enqueue their read and

/// then try to start a detached dispatcher task by acquiring one of

/// `dispatch_slots` semaphore permits. A dispatcher repeatedly drains up to

/// `max_batch_bytes` worth of pending reads into a single `BatchReadBlobs`

/// request until the queue is empty. This is work-conserving (a read never

/// waits while a dispatch slot is free) and uses no timers.

#[derive(Debug, MetricsComponent)]

struct ReadBatcher {

    max_blob_size_bytes: u64,

    max_batch_bytes: u64,

    max_queued_bytes: u64,

    queue: Mutex<ReadQueue>,

    dispatch_slots: Arc<Semaphore>,

    #[metric(help = "Number of BatchReadBlobs RPCs sent by the read coalescer")]

    batches_sent: AtomicU64,

    #[metric(help = "Number of blob reads coalesced into BatchReadBlobs RPCs")]

    blobs_batched: AtomicU64,

    #[metric(

        help = "Number of reads that bypassed batching because the queue byte budget was full"

    )]

    queue_bypasses: AtomicU64,

    #[metric(help = "Number of batched reads that resolved to a per-entry error")]

    batched_read_errors: AtomicU64,

    #[metric(help = "Payload bytes currently waiting in the read coalescer queue")]

    queued_bytes: AtomicU64,

}

impl ReadBatcher {

    fn new(config: &GrpcReadBatchingConfig) -> Self {

        Self {

            max_blob_size_bytes: config.max_blob_size_bytes,

            max_batch_bytes: config.max_batch_bytes,

            max_queued_bytes: config.max_queued_bytes,

            queue: Mutex::new(ReadQueue::default()),

            dispatch_slots: Arc::new(Semaphore::new(config.dispatch_slots)),

            batches_sent: AtomicU64::new(0),

            blobs_batched: AtomicU64::new(0),

            queue_bypasses: AtomicU64::new(0),

            batched_read_errors: AtomicU64::new(0),

            queued_bytes: AtomicU64::new(0),

        }

    }

}

/// Mirrors the default retryable-code classification used by

/// `nativelink_util::retry::Retrier::should_retry`: only codes that are

/// always terminal are considered non-retryable.

const fn is_retryable_code(code: Code) -> bool {

    !matches!(

        code,

        Code::Ok

            | Code::InvalidArgument

            | Code::FailedPrecondition

            | Code::OutOfRange

            | Code::Unimplemented

            | Code::NotFound

            | Code::AlreadyExists

            | Code::PermissionDenied

            | Code::Unauthenticated

    )

}

// This store is usually a pass-through store, but can also be used as a CAS store. Using it as an

// AC store has one major side-effect... The has() function may not give the proper size of the

// underlying data. This might cause issues if embedded in certain stores.

/// Starting capacity for the stitched-together batch response. Only a hint:

/// most batches are one chunk, and the vec grows if not.

const MAX_BATCH_UPDATE_ENTRIES_HINT: usize = 256;

/// Groups blobs so no single `BatchUpdateBlobs` RPC carries more than

/// `MAX_WRITE_REQUEST_DATA_BYTES` of payload.

///

/// An entry larger than the cap on its own still goes out alone. Nothing can

/// be done for it here, since a batch entry cannot be split across messages,

/// and REAPI already tells clients to use `ByteStream` for blobs that size.

fn split_batch_update_requests(

    requests: Vec<batch_update_blobs_request::Request>,

) -> Vec<Vec<batch_update_blobs_request::Request>> {

    let mut chunks = Vec::new();

    let mut current: Vec<batch_update_blobs_request::Request> = Vec::new();

    let mut current_bytes = 0usize;

    for request in 
requests1
 {

        let len = request.data.len();

        if !current.is_empty() && 
current_bytes11
.saturating_add(len) > MAX_WRITE_REQUEST_DATA_BYTES {

            chunks.push(core::mem::take(&mut current));

            current_bytes = 0;

        }

        current_bytes = current_bytes.saturating_add(len);

        current.push(request);

    }

    // An empty batch still needs one RPC: the caller expects the upstream's

    // answer to an empty request, not a synthesised one.

    if !current.is_empty() || 
chunks0
.
is_empty0
() {

        chunks.push(current);

    
}0

    chunks

}

#[derive(Debug, MetricsComponent)]

pub struct GrpcStore {

    #[metric(help = "Instance name for the store")]

    instance_name: String,

    store_type: nativelink_config::stores::StoreType,

    retrier: Retrier,

    connection_manager: ConnectionManager,

    /// Per-RPC timeout. `Duration::ZERO` means disabled.

    rpc_timeout: Duration,

    use_legacy_resource_names: bool,

    headers: Vec<(MetadataKey<Ascii>, MetadataValue<Ascii>)>,

    forward_headers: Vec<String>,

    /// When configured, coalesces small-blob reads into `BatchReadBlobs`

    /// RPCs. `None` means reads always use the `ByteStream` `Read` path.

    #[metric(group = "read_batcher")]

    read_batcher: Option<ReadBatcher>,

    remote_cache_compression_enabled: bool,

    /// Used by the read coalescer to hand a strong reference of this store

    /// to detached dispatcher tasks.

    weak_self: Weak<Self>,

}

impl GrpcStore {

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

        Self::new_with_jitter(spec, spec.retry.make_jitter_fn())

    }

    pub fn new_with_jitter(

        spec: &GrpcSpec,

        jitter_fn: Arc<dyn Fn(Duration) -> Duration + Send + Sync>,

    ) -> Result<Arc<Self>, Error> {

        error_if!(

            matches!(

                spec.store_type,

                nativelink_config::stores::StoreType::EventSink

            ) && spec.experimental_read_batching.is_some(),

            "An event sink cannot batch reads"

        );

        error_if!(

            spec.endpoints.is_empty(),

            "Expected at least 1 endpoint in GrpcStore"

        );

        let mut endpoints = Vec::with_capacity(spec.endpoints.len());

        for endpoint_config in &spec.endpoints {

            let endpoint = tls_utils::endpoint(endpoint_config).map_err(|e| 
{0

                Error::from_std_err(Code::InvalidArgument, &e)

                    .append("Invalid URI for GrpcStore endpoint")

            })?;

            endpoints.push(endpoint);

        }

        let rpc_timeout = Duration::from_secs(spec.rpc_timeout_s);

        let 
read_batcher47
 = match &spec.experimental_read_batching {

            Some(config) => {

                error_if!(

                    config.dispatch_slots == 0,

                    "experimental_read_batching.dispatch_slots must be greater than zero"

                );

                // Batched reads share one upstream RPC across many client

                // requests, so per-client forwarded headers (e.g. credentials)

                // cannot be attached correctly.

                error_if!(

                    !spec.forward_headers.is_empty(),

                    "experimental_read_batching is incompatible with forward_headers"

                );

                Some(ReadBatcher::new(config))

            }

            None => None,

        };

        let mut headers = Vec::with_capacity(spec.headers.len());

        for (
name1
, 
value1
) in &spec.headers {

            // We lowercase keys as HTTP headers are case-insensitive so we should match all cases

            let key = MetadataKey::from_bytes(name.to_lowercase().as_bytes()).map_err(|_| 
{0

                make_err!(Code::InvalidArgument, "Invalid gRPC metadata key: {name}")

            })?;

            let val = MetadataValue::try_from(value.as_str()).map_err(|_| 
{0

                make_err!(

                    Code::InvalidArgument,

                    "Invalid gRPC metadata value for key: {name}"

                )

            })?;

            headers.push((key, val));

        }

        Ok(Arc::new_cyclic(|weak_self| Self {

            weak_self: weak_self.clone(),

            instance_name: spec.instance_name.clone(),

            store_type: spec.store_type,

            retrier: Retrier::new(

                Arc::new(|duration| 
Box::pin2
(
sleep2
(
duration2
))),

                jitter_fn.clone(),

                spec.retry.clone(),

            )

            .with_client("grpc_store"),

            connection_manager: ConnectionManager::new(

                endpoints,

                spec.connections_per_endpoint,

                spec.max_concurrent_requests,

                spec.retry.clone(),

                jitter_fn,

            ),

            rpc_timeout,

            use_legacy_resource_names: spec.use_legacy_resource_names,

            read_batcher,

            remote_cache_compression_enabled: spec.remote_cache_compression_enabled(),

            headers,

            // We lowercase keys as HTTP headers are case-insensitive so we should match all cases

            forward_headers: spec

                .forward_headers

                .iter()

                .map(|s| 
s1
.
to_lowercase1
())

                .collect(),

        }))

    }

    async fn perform_request<F, Fut, R, I>(&self, input: I, mut request: F) -> Result<R, Error>

    where

        F: FnMut(I) -> Fut + Send + Copy,

        Fut: Future<Output = Result<R, Error>> + Send,

        R: Send,

        I: Send + Clone,

    {

        self.retrier

            .retry(unfold(input, move |input| async move {

                let input_clone = input.clone();

                Some((

                    request(input_clone)

                        .await

                        .map_or_else(RetryResult::Retry, RetryResult::Ok),

                    input,

                ))

            }))

            .await

    }

    pub async fn find_missing_blobs(

        &self,

        grpc_request: Request<FindMissingBlobsRequest>,

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

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let mut request = grpc_request.into_inner();

        // Some builds (Chromium for example) do lots of empty requests for some reason, so shortcut them

        if request.blob_digests.is_empty() {

            return Ok(Response::new(FindMissingBlobsResponse {

                missing_blob_digests: vec![],

            }));

        }

        request.instance_name.clone_from(&self.instance_name);

        self.perform_request(request, |request| async move {

            let channel = self

                .connection_manager

                .connection(format!(

                    "find_missing_blobs: ({}) {:?}",

                    request.blob_digests.len(),

                    request.blob_digests

                ))

                .await

                .err_tip(|| "in find_missing_blobs")?;

            ContentAddressableStorageClient::new(channel)

                .find_missing_blobs(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,

                ))

                .await

                .err_tip(|| "in GrpcStore::find_missing_blobs")

        })

        .await

    }

    pub async fn batch_update_blobs(

        &self,

        grpc_request: Request<BatchUpdateBlobsRequest>,

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

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let mut request = grpc_request.into_inner();

        request.instance_name.clone_from(&self.instance_name);

        // A batch is one message, so an oversized one cannot be chunked the

        // way a ByteStream write can; it has to become several RPCs. Split on

        // the accumulated payload and stitch the responses back together in

        // request order, which is what the caller matches results by.

        let mut responses =

            Vec::with_capacity(request.requests.len().min(MAX_BATCH_UPDATE_ENTRIES_HINT));

        for chunk in 
split_batch_update_requests1
(
core::mem::take1
(
&mut request.requests1
)) {

            let chunk_request = BatchUpdateBlobsRequest {

                instance_name: request.instance_name.clone(),

                requests: chunk,

                digest_function: request.digest_function,

            };

            let response = self

                .perform_request(chunk_request, |chunk_request| async move {

                    let channel = self

                        .connection_manager

                        .connection("batch_update_blobs".into())

                        .await

                        .err_tip(|| "in batch_update_blobs")
?0
;

                    ContentAddressableStorageClient::new(channel)

                        .batch_update_blobs(enrich_request(

                            Request::new(chunk_request),

                            &self.headers,

                            &self.forward_headers,

                        ))

                        .await

                        .err_tip(|| "in GrpcStore::batch_update_blobs")

                })

                .await
?0
;

            responses.extend(response.into_inner().responses);

        }

        Ok(Response::new(BatchUpdateBlobsResponse { responses }))

    }

    pub async fn batch_read_blobs(

        &self,

        grpc_request: Request<BatchReadBlobsRequest>,

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

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let mut request = grpc_request.into_inner();

        request.instance_name.clone_from(&self.instance_name);

        self.perform_request(request, |request| async move {

            let channel = self

                .connection_manager

                .connection("batch_read_blobs".into())

                .await

                .err_tip(|| "in batch_read_blobs")
?0
;

            ContentAddressableStorageClient::new(channel)

                .batch_read_blobs(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,

                ))

                .await

                .err_tip(|| "in GrpcStore::batch_read_blobs")

        })

        .await

    }

    /// Enqueues a small-blob read for coalescing into a `BatchReadBlobs` RPC

    /// and waits for its result. Returns `None` when the queue is over its

    /// byte budget, in which case the caller must fall back to the

    /// `ByteStream` `Read` path.

    async fn batched_read(

        &self,

        batcher: &ReadBatcher,

        digest: DigestInfo,

    ) -> Option<Result<Bytes, Error>> {

        // Capture the digest function from the caller's ambient context now;

        // the dispatcher runs on a detached task with no such context.

        let digest_function: i32 = Context::current()

            .get::<DigestHasherFunc>()

            .map_or_else(default_digest_hasher_func, |v| *v)

            .proto_digest_func()

            .into();

        let (tx, rx) = oneshot::channel();

        {

            let mut queue = batcher.queue.lock();

            // Admission control. On overflow gracefully degrade to the

            // stream path instead of blocking.

            let new_bytes = queue.bytes.saturating_add(digest.size_bytes());

            if new_bytes > batcher.max_queued_bytes {

                batcher.queue_bypasses.fetch_add(1, Ordering::Relaxed);

                return None;

            }

            queue.bytes = new_bytes;

            batcher.queued_bytes.store(new_bytes, Ordering::Relaxed);

            queue.items.push_back(PendingRead {

                digest,

                digest_function,

                tx,

            });

        }

        self.maybe_dispatch_read_batches(batcher);

        Some(rx.await.
unwrap_or_else207
(|_| 
{0

            Err(make_err!(

                Code::Internal,

                "Read batch dispatcher dropped result in GrpcStore::batched_read"

            ))

        }))

    }

    /// Tries to start a read batch dispatcher. This is work-conserving: if a

    /// dispatch slot is free a dispatcher is started immediately, otherwise

    /// one of the active dispatchers is responsible for every currently

    /// queued item. The dispatcher runs as a detached task so that

    /// cancellation of any individual reader can neither abort an in-flight

    /// `BatchReadBlobs` RPC nor strand still-queued waiters.

    fn maybe_dispatch_read_batches(&self, batcher: &ReadBatcher) {

        let Ok(
mut permit4
) = batcher.dispatch_slots.clone().try_acquire_owned() else {

            return;

        };

        let Some(store) = self.weak_self.upgrade() else {

            return;

        };

        background_spawn!("grpc_store_read_batch_dispatch", async move {

            let Some(batcher) = &store.read_batcher else {

                return;

            };

            loop {

                store.dispatch_read_batches(batcher).await;

                drop(permit);

                // Items may have been enqueued between the last drain and

                // the permit release. Re-check so they are not stranded

                // with no active dispatcher.

                if batcher.queue.lock().items.is_empty() {

                    return;

                }

                match batcher.dispatch_slots.clone().try_acquire_owned() {

                    Ok(new_permit) => permit = new_permit,

                    // Another dispatcher is active and will observe these

                    // items (or re-check after releasing its own permit).

                    Err(_) => return,

                }

            }

        });

    }

    /// Drains the pending read queue, sending one `BatchReadBlobs` RPC per

    /// drained batch, until the queue is empty.

    async fn dispatch_read_batches(&self, batcher: &ReadBatcher) {

        loop {

            let batch = {

                let mut queue = batcher.queue.lock();

                let Some(
head4
) = queue.items.front() else {

                    return;

                };

                // All digests in one BatchReadBlobsRequest must use the same

                // digest function. Partition-drain: take items matching the

                // head's digest function from anywhere in the queue (up to

                // the batch budget) and keep the rest in relative order.

                let digest_function = head.digest_function;

                let mut batch = Vec::new();

                let mut batch_bytes = 0u64;

                let mut rest = VecDeque::with_capacity(queue.items.len());

                while let Some(
item208
) = queue.items.pop_front() {

                    let item_cost = item

                        .digest

                        .size_bytes()

                        .saturating_add(BATCH_READ_PER_ENTRY_OVERHEAD_BYTES);

                    if item.digest_function == digest_function

                        && (batch.is_empty()

                            || batch_bytes.saturating_add(item_cost) <= batcher.max_batch_bytes)

                    {

                        batch_bytes = batch_bytes.saturating_add(item_cost);

                        queue.bytes = queue.bytes.saturating_sub(item.digest.size_bytes());

                        batch.push(item);

                    } else {

                        rest.push_back(item);

                    }

                }

                queue.items = rest;

                batcher.queued_bytes.store(queue.bytes, Ordering::Relaxed);

                batch

            };

            self.send_read_batch(batcher, batch).await;

        }

    }

    /// Sends one `BatchReadBlobs` RPC for `batch` and demultiplexes the

    /// per-blob responses back to the waiting readers. One failed item does

    /// not affect its batch-mates; failure of the whole RPC is broadcast to

    /// every item in the batch.

    async fn send_read_batch(&self, batcher: &ReadBatcher, batch: Vec<PendingRead>) {

        let Some(digest_function) = batch.first().map(|item| item.digest_function) else {

            return;

        };

        // Servers may dedupe duplicate digests within one request, so group

        // the waiters per digest and request each digest exactly once,

        // fanning the (refcounted) data out to every waiter.

        let batch_len = u64::try_from(batch.len()).unwrap_or(u64::MAX);

        let mut waiters: HashMap<DigestInfo, Vec<PendingRead>> = HashMap::new();

        for item in 
batch4
 {

            waiters.entry(item.digest).or_default().push(item);

        }

        let request = BatchReadBlobsRequest {

            // batch_read_blobs() overwrites the instance name, so there is

            // no need to set it here.

            instance_name: String::new(),

            digests: 
waiters4
.
keys4
().
map4
(|digest| (*digest).into()).
collect4
(),

            acceptable_compressors: vec![],

            digest_function,

        };

        batcher.batches_sent.fetch_add(1, Ordering::Relaxed);

        batcher

            .blobs_batched

            .fetch_add(batch_len, Ordering::Relaxed);

        let response = match self.batch_read_blobs(Request::new(request)).await {

            Ok(response) => response.into_inner(),

            Err(err) => {

                // The whole RPC failed, so every waiter in this batch gets

                // the error. Waiters may have gone away, ignore send errors.

                for item in waiters.into_values().flatten() {

                    drop(item.tx.send(Err(err.clone())));

                }

                return;

            }

        };

        for entry in 
response.responses4
 {

            let Some(Ok(entry_digest)) = entry.digest.map(DigestInfo::try_from) else {

                continue;

            };

            let Some(items) = waiters.remove(&entry_digest) else {

                continue;

            };

            let entry_len = u64::try_from(entry.data.len()).unwrap_or(u64::MAX);

            let result = if let Some(
status1
) = entry.status.filter(|status| status.code != 0) {

                Err(Error::from(status)

                    .append("Batch read entry failed in GrpcStore::send_read_batch"))

            } else if entry.compressor != 0 {

                // We requested no acceptable compressors, so data must be

                // returned with the identity compressor.

                Err(make_err!(

                    Code::Internal,

                    "BatchReadBlobs entry for {entry_digest} used unsupported compressor {}",

                    entry.compressor

                ))

            } else if entry_len != entry_digest.size_bytes() {

                Err(make_err!(

                    Code::Internal,

                    "BatchReadBlobs entry for {entry_digest} returned {entry_len} bytes, expected {}",

                    entry_digest.size_bytes()

                ))

            } else {

                Ok(entry.data)

            };

            if result.is_err() {

                batcher.batched_read_errors.fetch_add(

                    u64::try_from(items.len()).unwrap_or(u64::MAX),

                    Ordering::Relaxed,

                );

            }

            for item in 
items207
 {

                drop(item.tx.send(result.clone()));

            }

        }

        // Any waiter with no matching response entry is missing upstream.

        for 
item0
 in waiters.into_values().flatten() {

            batcher.batched_read_errors.fetch_add(1, Ordering::Relaxed);

            let err = make_err!(

                Code::NotFound,

                "Blob {} not found in BatchReadBlobs response",

                item.digest

            );

            drop(item.tx.send(Err(err)));

        }

    }

    pub async fn get_tree(

        &self,

        grpc_request: Request<GetTreeRequest>,

    ) -> Result<Response<Streaming<GetTreeResponse>>, Error> {

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let mut request = grpc_request.into_inner();

        request.instance_name.clone_from(&self.instance_name);

        self.perform_request(request, |request| async move {

            let channel = self

                .connection_manager

                .connection(format!("get_tree: {:?}", request.root_digest))

                .await

                .err_tip(|| "in get_tree")
?0
;

            ContentAddressableStorageClient::new(channel)

                .get_tree(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,

                ))

                .await

                .err_tip(|| "in GrpcStore::get_tree")

        })

        .await

    }

    pub async fn split_blob(

        &self,

        grpc_request: Request<SplitBlobRequest>,

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

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let mut request = grpc_request.into_inner();

        request.instance_name.clone_from(&self.instance_name);

        self.perform_request(request, |request| async move {

            let channel = self

                .connection_manager

                .connection(format!("split_blob: {:?}", request.blob_digest))

                .await

                .err_tip(|| "in split_blob")
?0
;

            ContentAddressableStorageClient::new(channel)

                .split_blob(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,

                ))

                .await

                .err_tip(|| "in GrpcStore::split_blob")

        })

        .await

    }

    pub async fn splice_blob(

        &self,

        grpc_request: Request<SpliceBlobRequest>,

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

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let mut request = grpc_request.into_inner();

        request.instance_name.clone_from(&self.instance_name);

        self.perform_request(request, |request| async move {

            let channel = self

                .connection_manager

                .connection(format!("splice_blob: {:?}", request.blob_digest))

                .await

                .err_tip(|| "in splice_blob")
?0
;

            ContentAddressableStorageClient::new(channel)

                .splice_blob(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,

                ))

                .await

                .err_tip(|| "in GrpcStore::splice_blob")

        })

        .await

    }

    fn get_read_request(&self, mut request: ReadRequest) -> Result<ReadRequest, Error> {

        const IS_UPLOAD_FALSE: bool = false;

        let mut resource_info = ResourceInfo::new(&request.resource_name, IS_UPLOAD_FALSE)?;

        if resource_info.instance_name != self.instance_name {

            resource_info.instance_name = Cow::Borrowed(&self.instance_name);

            request.resource_name = resource_info.to_string(IS_UPLOAD_FALSE);

        }

        Ok(request)

    }

    async fn read_internal(

        &self,

        request: ReadRequest,

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

        let channel = self

            .connection_manager

            .connection(format!("read_internal: {}", request.resource_name))

            .await

            .err_tip(|| "in read_internal")
?0
;

        let mut response = ByteStreamClient::new(channel)

            .read(enrich_request(

                Request::new(request),

                &self.headers,

                &self.forward_headers,

            ))

            .await

            .err_tip(|| "in GrpcStore::read")
?0

            .into_inner();

        let 
first_response14
 = response

            .message()

            .await

            .err_tip(|| "Fetching first chunk in GrpcStore::read()")
?1
;

        Ok(FirstStream::new(first_response, response))

    }

    pub async fn read<R>(

        &self,

        grpc_request: R,

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

    where

        R: IntoRequest<ReadRequest>,

    {

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let request = self.get_read_request(grpc_request.into_request().into_inner())?;

        self.perform_request(request, |request| async move {

            self.read_internal(request).await

        })

        .await

    }

    pub async fn write<T, E>(

        &self,

        stream: WriteRequestStreamWrapper<T>,

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

    where

        T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,

        E: Into<Error> + 'static,

    {

        const RESUMABLE: bool = true;

        self.write_internal(stream, RESUMABLE).await

    }

    /// Writes one event to an event sink. A sink takes an upload only from

    /// offset zero, so where `write` resumes a retry from its last cached

    /// messages (mid-event, for an event over one `WriteRequest`), this

    /// resends the whole event on every attempt. Events are bounded, so the

    /// event is held in memory for its retries.

    async fn update_event(

        &self,

        resource_name: String,

        mut reader: DropCloserReadHalf,

        digest: DigestInfo,

    ) -> Result<u64, Error> {

        let size = digest.size_bytes();

        let data = reader

            .consume(Some(

                usize::try_from(size).err_tip(|| "Event size in GrpcStore::update_event")
?0
,

            ))

            .await

            .err_tip(|| "In GrpcStore::update_event")
?0
;

        let response = self

            .retrier

            .
retry4
(
unfold4
(
()4
, move |()| {

                let mut requests: Vec<WriteRequest> = data

                    .chunks(MAX_WRITE_REQUEST_DATA_BYTES)

                    .
scan6
(0usize, |offset, chunk| {

                        let start = *offset;

                        *offset += chunk.len();

                        Some(WriteRequest {

                            resource_name: resource_name.clone(),

                            write_offset: i64::try_from(start).unwrap_or(i64::MAX),

                            finish_write: false,

                            data: data.slice(start..*offset),

                        })

                    })

                    .collect();

                requests.push(WriteRequest {

                    resource_name: resource_name.clone(),

                    write_offset: i64::try_from(data.len()).unwrap_or(i64::MAX),

                    finish_write: true,

                    data: Bytes::new(),

                });

                async move {

                    let rpc = self

                        .connection_manager

                        .connection("update_event".into())

                        .and_then(|channel| async move {

                            ByteStreamClient::new(channel)

                                .write(enrich_request(

                                    Request::new(futures::stream::iter(requests)),

                                    &self.headers,

                                    &self.forward_headers,

                                ))

                                .await

                                .err_tip(|| "in GrpcStore::update_event")

                        });

                    let result = if self.rpc_timeout > Duration::ZERO {

                        tokio::time::timeout(self.rpc_timeout, rpc)

                            .await

                            .unwrap_or_else(|_elapsed| 
{0

                                Err(make_err!(

                                    Code::DeadlineExceeded,

                                    "GrpcStore::update_event RPC timed out after {}s",

                                    self.rpc_timeout.as_secs()

                                ))

                            })

                    } else {

                        rpc.await

                    };

                    if let Err(
err2
) = &result {

                        warn!(instance_name = %self.instance_name, ?err, "GrpcStore::update_event: RPC failed; a retry resends the whole event");

                    }

                    Some((result.map_or_else(RetryResult::Retry, RetryResult::Ok), ()))

                }

            }))

            .await
?0
;

        if u64::try_from(response.get_ref().committed_size).ok() != Some(size) {

            return Err(make_err!(

                Code::DataLoss,

                "Event sink acknowledged {} bytes; expected {}",

                response.get_ref().committed_size,

                size

            ));

        }

        Ok(size)

    }

    async fn write_internal<T, E>(

        &self,

        stream: WriteRequestStreamWrapper<T>,

        resumable: bool,

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

    where

        T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,

        E: Into<Error> + 'static,

    {

        error_if!(

            matches!(self.store_type, nativelink_config::stores::StoreType::Ac),

            "CAS operation on AC store"

        );

        let mut write_state = WriteState::new(self.instance_name.clone(), stream);

        if !resumable {

            write_state.set_non_resumable();

        
}1

        let local_state = Arc::new(Mutex::new(write_state));

        let write_start = std::time::Instant::now();

        let instance_name = self.instance_name.clone();

        let rpc_timeout = self.rpc_timeout;

        trace!(

            instance_name = %instance_name,

            rpc_timeout_s = rpc_timeout.as_secs(),

            "GrpcStore::write: starting ByteStream write",

        );

        let mut attempt: u32 = 0;

        let 
result9
 = self

            .retrier

            .retry(unfold(local_state, move |local_state| {

                attempt += 1;

                let instance_name = instance_name.clone();

                async move {

                    // The client write may occur on a separate thread and

                    // therefore in order to share the state with it we have to

                    // wrap it in a Mutex and retrieve it after the write

                    // has completed.  There is no way to get the value back

                    // from the client.

                    trace!(

                        instance_name = %instance_name,

                        attempt,

                        "GrpcStore::write: requesting connection from pool",

                    );

                    let conn_start = std::time::Instant::now();

                    let rpc_fut = self.connection_manager.connection("write".into()).and_then(

                        |channel| {

                            let conn_elapsed = conn_start.elapsed();

                            let instance_for_rpc = instance_name.clone();

                            let conn_elapsed_ms =

                                u64::try_from(conn_elapsed.as_millis()).unwrap_or(u64::MAX);

                            trace!(

                                instance_name = %instance_for_rpc,

                                conn_elapsed_ms,

                                "GrpcStore::write: got connection, starting ByteStream.Write RPC",

                            );

                            let rpc_start = std::time::Instant::now();

                            let local_state_for_rpc = local_state.clone();

                            async move {

                                let res = ByteStreamClient::new(channel)

                                    .write(enrich_request(

                                        Request::new(WriteStateWrapper::new(local_state_for_rpc)),

                                        &self.headers,

                                        &self.forward_headers,

                                    ))

                                    .await

                                    .err_tip(|| "in GrpcStore::write");

                                let rpc_elapsed_ms = u64::try_from(rpc_start.elapsed().as_millis())

                                    .unwrap_or(u64::MAX);

                                trace!(

                                    instance_name = %instance_for_rpc,

                                    rpc_elapsed_ms,

                                    success = res.is_ok(),

                                    "GrpcStore::write: ByteStream.Write RPC returned",

                                );

                                res

                            }