// 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::pin::Pin;

use core::sync::atomic::{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,

};

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

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, 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::ClientHeaders;

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

use opentelemetry::context::Context;

use opentelemetry::global;

use opentelemetry::propagation::Injector;

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::{error, trace, warn};

use uuid::Uuid;

struct TonicMetadataInjector<'a>(&'a mut tonic::metadata::MetadataMap);

impl Injector for TonicMetadataInjector<'_> {

    fn set(&mut self, key: &str, value: String) {

        if let (Ok(k), Ok(v)) = (

            MetadataKey::from_bytes(key.as_bytes()),

            MetadataValue::try_from(&value),

        ) {

            self.0.insert(k, v);

        }

    }

}

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

    }

    if !forward_headers.is_empty()

        && let Some(client_headers) = Context::current().get::<ClientHeaders>()

    {

        for name in forward_headers {

            if let Some(value) = client_headers.0.get(&name.to_lowercase())

                && let (Ok(k), Ok(v)) = (

                    MetadataKey::from_bytes(name.as_bytes()),

                    MetadataValue::try_from(value.as_str()),

                )

            {

                request.metadata_mut().insert(k, v);

            
}0

        }

    }

    global::get_text_map_propagator(|propagator| {

        propagator.inject(&mut TonicMetadataInjector(request.metadata_mut()));

    });

    request

}

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

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

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

    /// to detached dispatcher tasks.

    weak_self: Weak<Self>,

}

impl GrpcStore {

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

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

    }

    pub async fn new_with_jitter(

        spec: &GrpcSpec,

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

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

        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_batcher15
 = 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::pin0
(
sleep0
(
duration0
))),

                jitter_fn.clone(),

                spec.retry.clone(),

            ),

            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,

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

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

            let channel = self

                .connection_manager

                .connection("batch_update_blobs".into())

                .await

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

            ContentAddressableStorageClient::new(channel)

                .batch_update_blobs(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,

                ))

                .await

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

        })

        .await

    }

    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")?;

            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_response = response

            .message()

            .await

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

        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,

    {

        error_if!(

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

            "CAS operation on AC store"

        );

        let local_state = Arc::new(Mutex::new(WriteState::new(

            self.instance_name.clone(),

            stream,

        )));

        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 result = 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

                            }

                        },

                    );

                    let result = if rpc_timeout > Duration::ZERO {

                        match tokio::time::timeout(rpc_timeout, rpc_fut).await {

                            Ok(res) => res,

                            Err(_elapsed) => {

                                warn!(

                                    instance_name = %instance_name,

                                    attempt,

                                    rpc_timeout_s = rpc_timeout.as_secs(),

                                    "GrpcStore::write: per-RPC timeout exceeded, cancelling",

                                );

                                #[allow(unused_qualifications)]

                                Err(nativelink_error::make_err!(

                                    nativelink_error::Code::DeadlineExceeded,

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

                                    rpc_timeout.as_secs()

                                ))

                            }

                        }

                    } else {

                        rpc_fut.await

                    };

                    // Get the state back from StateWrapper, this should be

                    // uncontended since write has returned.

                    let mut local_state_locked = local_state.lock();

                    let result = local_state_locked

                        .take_read_stream_error()

                        .map(|err| RetryResult::Err(
err0
.
append0
("Where read_stream_error was set")))

                        .unwrap_or_else(|| {

                            // No stream error, handle the original result

                            match result {

                                Ok(response) => RetryResult::Ok(response),

                                Err(ref err) => {

                                    warn!(

                                        instance_name = %instance_name,

                                        attempt,

                                        ?err,

                                        can_resume = local_state_locked.can_resume(),

                                        "GrpcStore::write: RPC failed",

                                    );

                                    if local_state_locked.can_resume() {

                                        local_state_locked.resume();

                                        RetryResult::Retry(err.clone())

                                    } else {

                                        RetryResult::Err(

                                            err.clone().append("Retry is not possible"),

                                        )

                                    }

                                }

                            }

                        });

                    drop(local_state_locked);

                    Some((result, local_state))

                }

            }))

            .await
?0
;

        let total_elapsed_ms = u64::try_from(write_start.elapsed().as_millis()).unwrap_or(u64::MAX);

        trace!(

            instance_name = %self.instance_name,

            total_elapsed_ms,

            "GrpcStore::write: completed successfully",

        );

        Ok(result)

    }

    pub async fn query_write_status(

        &self,

        grpc_request: Request<QueryWriteStatusRequest>,

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

        const IS_UPLOAD_TRUE: bool = true;

        error_if!(

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

            "CAS operation on AC store"

        );

        let mut request = grpc_request.into_inner();

        let mut request_info = ResourceInfo::new(&request.resource_name, IS_UPLOAD_TRUE)?;

        if request_info.instance_name != self.instance_name {

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

            request.resource_name = request_info.to_string(IS_UPLOAD_TRUE);

        }

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

            let channel = self

                .connection_manager

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

                .await

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

            ByteStreamClient::new(channel)

                .query_write_status(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,

                ))

                .await

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

        })

        .await

    }

    pub async fn get_action_result(

        &self,

        grpc_request: Request<GetActionResultRequest>,

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

        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_action_result: {:?}", request.action_digest))

                .await

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

            ActionCacheClient::new(channel)

                .get_action_result(enrich_request(

                    Request::new(request),

                    &self.headers,

                    &self.forward_headers,