// 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::time::Duration;

use core::{cmp, mem};

use std::borrow::Cow;

use std::sync::Arc;

use std::time::Instant;

use async_trait::async_trait;

use azure_core::credentials::TokenCredential;

use azure_core::error::ErrorKind;

use azure_core::http::headers::{RETRY_AFTER, RETRY_AFTER_MS, X_MS_RETRY_AFTER_MS};

use azure_core::http::{RequestContent, RetryOptions, StatusCode, Transport, Url};

use azure_identity::WorkloadIdentityCredential;

use azure_storage_blob::clients::{BlobContainerClient, BlobContainerClientOptions};

use azure_storage_blob::models::{

    BlobClientDownloadOptions, BlobClientGetPropertiesResultHeaders, BlockListType,

    BlockLookupList, HttpRange, StorageErrorCode,

};

use futures::future::FusedFuture;

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

use futures::{FutureExt, StreamExt, TryStreamExt};

use nativelink_config::stores::ExperimentalAzureSpec;

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

use nativelink_metric::MetricsComponent;

use nativelink_util::adaptive_permits::{AdaptivePermit, AdaptivePermits};

use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};

use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};

use nativelink_util::instant_wrapper::InstantWrapper;

use nativelink_util::metrics::{record_object_store_request_wait, record_object_store_throttle};

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

use nativelink_util::store_trait::{

    RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo,

};

use nativelink_util::{background_spawn, fs};

use parking_lot::Mutex;

use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};

use tokio::time::sleep;

use tracing::{Level, event};

use crate::cas_utils::is_zero_digest;

use crate::common_s3_utils::install_default_rustls_crypto_provider;

// Check the below doc for the limits specific to Azure.

// https://learn.microsoft.com/en-us/azure/storage/blobs/scalability-targets#scale-targets-for-blob-storage

// Maximum number of blocks in a block blob or append blob

const MAX_BLOCKS: usize = 50_000;

// Maximum size of a block in a block blob (4,000 MiB)

const MAX_BLOCK_SIZE: u64 = 4_000 * 1024 * 1024; // 4,000 MiB = 4 GiB

// Default block size for uploads (5 MiB)

const DEFAULT_BLOCK_SIZE: u64 = 5 * 1024 * 1024; // 5 MiB

// Default maximum retry buffer per request

const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 5 * 1024 * 1024; // 5 MiB

// Default maximum number of concurrent uploads

const DEFAULT_MAX_CONCURRENT_UPLOADS: usize = 10;

// Default public Azure Blob Storage endpoint suffix.

const DEFAULT_BLOB_ENDPOINT_SUFFIX: &str = "blob.core.windows.net";

/// Successful reads in a row before the adaptive read limit grows by one.

const READ_LIMIT_GROW_AFTER: usize = 16;

/// Shortest time between two warnings from one store that the host ran out

/// of local ports.

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

#[derive(MetricsComponent)]

pub struct AzureBlobStore<NowFn> {

    client: Arc<BlobContainerClient>,

    now_fn: NowFn,

    #[metric(help = "The container name for the Azure store")]

    container: String,

    #[metric(help = "The blob prefix for the Azure store")]

    blob_prefix: String,

    retrier: Retrier,

    #[metric(help = "The number of seconds to consider an object expired")]

    consider_expired_after_s: i64,

    #[metric(help = "The number of bytes to buffer for retrying requests")]

    max_retry_buffer_per_request: usize,

    #[metric(help = "The number of concurrent uploads allowed")]

    max_concurrent_uploads: usize,

    throttle_retry_delay: Duration,

    throttle_retry_max_delay: Duration,

    /// Bounds `GET` and `HEAD` requests in flight; `None` leaves them

    /// unbounded. See `CommonObjectSpec::max_concurrent_reads`.

    read_permits: Option<Arc<AdaptivePermits>>,

    /// Bounds every request in flight, of any kind; `None` leaves them

    /// unbounded. See `CommonObjectSpec::max_concurrent_requests`.

    request_permits: Option<Arc<Semaphore>>,

    #[metric(help = "The most requests this store keeps in flight at once, 0 for unbounded")]

    max_concurrent_requests: usize,

    /// When this store last warned that the host ran out of local ports.

    last_port_hint: Mutex<Option<Instant>>,

}

impl<NowFn> core::fmt::Debug for AzureBlobStore<NowFn> {

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

        f.debug_struct("AzureBlobStore")

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

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

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

            .finish_non_exhaustive()

    }

}

/// A request that cannot connect, or reads nothing for this long (response

/// headers included), fails and is retried. Without a bound, a request that

/// hangs keeps its permit forever, and `max_concurrent_requests` of them

/// wedge the store. The read timer runs only while the body is polled, so a

/// download that a slow reader holds back is not cut off. The values are the

/// Azure SDK's own defaults, which this store's client replaces.

const AZURE_CONNECT_TIMEOUT: Duration = Duration::from_secs(20);

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

/// How long a cancelled upload's cleanup waits before it looks at the block

/// list, so a block whose request was already in flight has landed and is

/// seen rather than raced.

const STAGED_BLOCK_DISCARD_DELAY: Duration = Duration::from_secs(5);

/// Discards the blocks an upload staged if the upload never commits, from a

/// background task in `Drop`: after an error, so the caller gets it without

/// waiting for request permits, and when the upload future is dropped (the

/// `ByteStream` it fed died and the server swept it), since a dropped future

/// runs no code of its own. Only a committed upload disarms it.

struct StagedBlocksGuard {

    client: Arc<BlobContainerClient>,

    requests: Option<Arc<Semaphore>>,

    blob_path: String,

    armed: bool,

}

impl StagedBlocksGuard {

    const fn disarm(&mut self) {

        self.armed = false;

    }

}

impl Drop for StagedBlocksGuard {

    fn drop(&mut self) {

        if !self.armed {

            return;

        }

        let client = Arc::clone(&self.client);

        let requests = self.requests.take();

        let blob_path = mem::take(&mut self.blob_path);

        background_spawn!("azure_discard_staged_blocks", async move 
{8

            sleep(STAGED_BLOCK_DISCARD_DELAY).await;

            discard_staged_blocks(&client, requests.as_ref(), &blob_path).await;

        });

    }

}

/// Best effort: drop the blocks an upload left staged under this name.

/// Azure keeps uncommitted blocks, and bills them, for a week before

/// collecting them, which adds up under an upload failure storm. A

/// committed blob under the same name (a re-upload of a key that already

/// exists) must survive: re-committing its own block list is what discards

/// the staged blocks next to it. With nothing committed the name holds only

/// staged blocks and deleting it removes them. Failure here is logged, never

/// returned: the upload has already failed, and the blocks are collected

/// eventually either way. Each request takes a permit from `requests`

/// like any other the store sends.

async fn discard_staged_blocks(

    client: &BlobContainerClient,

    requests: Option<&Arc<Semaphore>>,

    blob_path: &str,

) {

    let blob_client = client.blob_client(blob_path);

    let block_list_result = {

        let _request_permit = request_permit(requests).await;

        blob_client

            .block_blob_client()

            .get_block_list(BlockListType::All, None)

            .await

    };

    let block_list = match block_list_result {

        Ok(response) => match response.into_model() {

            Ok(block_list) => block_list,

            Err(err) => {

                event!(

                    Level::WARN,

                    ?err,

                    blob_path,

                    "Could not read the block list after a failed Azure upload"

                );

                return;

            }

        },

        Err(err) => {

            if err.http_status() != Some(StatusCode::NotFound) {

                event!(

                    Level::WARN,

                    ?err,

                    blob_path,

                    "Could not read the block list after a failed Azure upload"

                );

            }

            return;

        }

    };

    let uncommitted = block_list.uncommitted_blocks.unwrap_or_default();

    if uncommitted.is_empty() {

        return;

    }

    let committed: Vec<Vec<u8>> = block_list

        .committed_blocks

        .unwrap_or_default()

        .into_iter()

        .filter_map(|block| block.name)

        .collect();

    let _request_permit = request_permit(requests).await;

    let result = if committed.is_empty() {

        blob_client.delete(None).await.map(|_| ())

    } else {

        let keep = BlockLookupList {

            committed: Some(committed),

            ..Default::default()

        };

        match RequestContent::try_from(keep) {

            Ok(blocks) => blob_client

                .block_blob_client()

                .commit_block_list(blocks, None)

                .await

                .map(|_| ()),

            Err(err) => Err(err),

        }

    };

    match result {

        Ok(()) => event!(

            Level::INFO,

            blob_path,

            staged_blocks = uncommitted.len(),

            "Discarded the blocks a failed Azure upload left staged"

        ),

        Err(err) => event!(

            Level::WARN,

            ?err,

            blob_path,

            staged_blocks = uncommitted.len(),

            "Could not discard the blocks a failed Azure upload left staged"

        ),

    }

}

/// A permit from the store's request pool when requests are bounded,

/// counted on `store.object.request_waits` when none was free. Every path

/// takes it after the read permit and before the process-wide open-file

/// permit, so no two requests wait on each other's pools in opposite order.

async fn request_permit(requests: Option<&Arc<Semaphore>>) -> Option<OwnedSemaphorePermit> {

    let 
requests97
 = requests
?194
;

    if requests.available_permits() == 0 {

        record_object_store_request_wait("azure");

    
}20

    Some(

        Arc::clone(requests)

            .acquire_owned()

            .await

            .expect("The request permit pool is never closed"),

    )

}

/// Whether `err` failed because this host had no local port left to

/// connect from: `connect()` answers `EADDRNOTAVAIL` when every source port

/// toward the destination is held, by a request in flight, an idle

/// connection kept for reuse, or a closed one still in `TIME_WAIT`.

pub fn is_out_of_local_ports(err: &(dyn core::error::Error + 'static)) -> bool {

    let mut next = Some(err);

    while let Some(
err43
) = next {

        if err

            .downcast_ref::<std::io::Error>()

            .is_some_and(|err| 
err.kind()8
 == 
std::io::ErrorKind::AddrNotAvailable8
)

        {

            return true;

        }

        next = err.source();

    }

    false

}

impl<I, NowFn> AzureBlobStore<NowFn>

where

    I: InstantWrapper,

    NowFn: Fn() -> I + Send + Sync + Unpin + 'static,

{

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

        let jitter_fn = spec.common.retry.make_jitter_fn();

        let client = Self::build_container_client(spec)?;

        Self::new_with_client_and_jitter(spec, client, jitter_fn, now_fn)

    }

    /// Builds the container URL and selects the auth strategy:

    ///   * `sas_url` set    -> use it verbatim as the container URL with no credential.

    ///   * otherwise        -> `https://{account}.{endpoint}/{container}` authenticated with

    ///     Entra ID via Workload Identity (keyless).

    fn build_container_client(spec: &ExperimentalAzureSpec) -> Result<BlobContainerClient, Error> {

        let mut options = BlobContainerClientOptions::default();

        options.client_options.retry = RetryOptions::none();

        // Hand the SDK an HTTP client with an explicit rustls (ring) config.

        options.client_options.transport = Some(Self::build_http_transport()?);

        let (container_url, credential): (Url, Option<Arc<dyn TokenCredential>>) =

            if let Some(sas_url) = spec.sas_url.as_ref() {

                let url = Url::parse(sas_url)

                    .map_err(|e| make_err!(Code::InvalidArgument, "Invalid Azure sas_url: {e}"))?;

                (url, None)

            } else {

                let endpoint = spec.endpoint.clone().unwrap_or_else(|| {

                    format!(

                        "https://{}.{DEFAULT_BLOB_ENDPOINT_SUFFIX}",

                        spec.account_name

                    )

                });

                let mut url = Url::parse(&endpoint)

                    .map_err(|e| make_err!(Code::InvalidArgument, "Invalid Azure endpoint: {e}"))?;

                url.path_segments_mut()

                    .map_err(|()| {

                        make_err!(

                            Code::InvalidArgument,

                            "Azure endpoint is not a valid base URL: {endpoint}"

                        )

                    })?

                    .pop_if_empty()

                    .push(&spec.container);

                let credential: Arc<dyn TokenCredential> = WorkloadIdentityCredential::new(None)

                    .map_err(|e| {

                        make_err!(

                            Code::FailedPrecondition,

                            "Failed to create Azure Workload Identity credential: {e}"

                        )

                    })?;

                (url, Some(credential))

            };

        BlobContainerClient::new(container_url, credential, Some(options))

            .map_err(|e| make_err!(Code::Unavailable, "Failed to create Azure client: {e}"))

    }

    /// Builds an HTTP transport for the Azure SDK backed by a reqwest client with

    /// an explicit rustls config using `NativeLink`'s ring crypto provider, so the

    /// SDK never falls back to guessing a provider (which breaks HTTPS here).

    fn build_http_transport() -> Result<Transport, Error> {

        install_default_rustls_crypto_provider();

        let mut roots = rustls::RootCertStore::empty();

        roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());

        let tls_config = rustls::ClientConfig::builder()

            .with_root_certificates(roots)

            .with_no_client_auth();

        let client = reqwest::Client::builder()

            .use_preconfigured_tls(tls_config)

            .connect_timeout(AZURE_CONNECT_TIMEOUT)

            .read_timeout(AZURE_READ_TIMEOUT)

            .build()

            .map_err(|e| make_err!(Code::Unavailable, "Failed to build Azure HTTP client: {e}"))?;

        Ok(Transport::new(Arc::new(client)))

    }

    pub fn new_with_client_and_jitter(

        spec: &ExperimentalAzureSpec,

        client: BlobContainerClient,

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

        now_fn: NowFn,

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

        let throttle_retry_delay = Duration::try_from_secs_f32(spec.common.throttle_retry_delay_s)

            .map_err(|e| 
make_err!0
(
Code::InvalidArgument0
, "Invalid throttle_retry_delay_s: {e}"))
?0
;

        let throttle_retry_max_delay =

            Duration::try_from_secs_f32(spec.common.throttle_retry_max_delay_s).map_err(|e| 
{0

                make_err!(

                    Code::InvalidArgument,

                    "Invalid throttle_retry_max_delay_s: {e}"

                )

            })?;

        if throttle_retry_max_delay < throttle_retry_delay {

            return Err(make_err!(

                Code::InvalidArgument,

                "throttle_retry_max_delay_s ({}) is less than throttle_retry_delay_s ({})",

                spec.common.throttle_retry_max_delay_s,

                spec.common.throttle_retry_delay_s

            ));

        }

        let read_permits = (spec.common.max_concurrent_reads > 0).then(|| 
{11

            let max = spec.common.max_concurrent_reads;

            let (min, grow_after) = if spec.common.adaptive_read_concurrency {

                (

                    spec.common.min_concurrent_reads.min(max),

                    READ_LIMIT_GROW_AFTER,

                )

            } else {

                (max, 0)

            };

            AdaptivePermits::new("azure", max, min, max, grow_after)

        });

        let max_concurrent_requests = spec.common.max_concurrent_requests;

        let max_concurrent_reads = spec.common.max_concurrent_reads;

        if max_concurrent_requests > 0

            && (max_concurrent_reads == 0 || 
max_concurrent_reads >= max_concurrent_requests7
)

        {

            return Err(make_err!(

                Code::InvalidArgument,

                "max_concurrent_requests ({max_concurrent_requests}) needs max_concurrent_reads between 1 and {}: a download holds its request permit while its body streams, so without that gap downloads can take every permit and stall existence checks and uploads (max_concurrent_reads is {max_concurrent_reads})",

                max_concurrent_requests - 1

            ));

        }

        let request_permits = (max_concurrent_requests > 0)

            .then(|| 
Arc::new5
(
Semaphore::new5
(
max_concurrent_requests5
)));

        let throttle_retry_budget =

            Duration::try_from_secs_f32(spec.common.throttle_retry_budget_s).map_err(|e| 
{0

                make_err!(

                    Code::InvalidArgument,

                    "Invalid throttle_retry_budget_s: {e}"

                )

            })?;

        Ok(Arc::new(Self {

            client: Arc::new(client),

            now_fn,

            container: spec.container.clone(),

            blob_prefix: spec

                .common

                .key_prefix

                .as_ref()

                .unwrap_or(&String::new())

                .clone(),

            retrier: Retrier::new(

                Arc::new(|duration| 
Box::pin21
(
sleep21
(
duration21
))),

                jitter_fn,

                spec.common.retry.clone(),

            )

            .with_client("azure")

            .with_throttle_budget(throttle_retry_budget)

            .with_throttle_max_wait(throttle_retry_max_delay),

            consider_expired_after_s: i64::from(spec.common.consider_expired_after_s),

            max_retry_buffer_per_request: spec

                .common

                .max_retry_buffer_per_request

                .unwrap_or(DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST),

            max_concurrent_uploads: spec

                .common

                .multipart_max_concurrent_uploads

                .unwrap_or(DEFAULT_MAX_CONCURRENT_UPLOADS),

            throttle_retry_delay,

            throttle_retry_max_delay,

            read_permits,

            request_permits,

            max_concurrent_requests,

            last_port_hint: Mutex::new(None),

        }))

    }

    /// The read limit the permit pool stands at, or `None` with reads

    /// unbounded.

    pub fn read_limit(&self) -> Option<usize> {

        self.read_permits.as_ref().map(|pool| pool.limit())

    }

    /// A download permit when reads are bounded. Held for the request,

    /// body included, so what is in flight is what the limit says. Only

    /// `get_part` takes one: a `HEAD` moves no bytes and must not wait

    /// behind a body still being sent to a slow client, and a throttle

    /// on either side says nothing about the other.

    async fn read_permit(&self) -> Option<AdaptivePermit> {

        match &self.read_permits {

            Some(pool) => Some(pool.acquire().await),

            None => None,

        }

    }

    /// A request permit when requests are bounded; see [`request_permit`].

    async fn request_permit(&self) -> Option<OwnedSemaphorePermit> {

        request_permit(self.request_permits.as_ref()).await

    }

    /// Names the likely cause, and the setting that bounds it, when a

    /// request failed because the host ran out of local ports. At most once

    /// per `PORT_HINT_INTERVAL`: under exhaustion most requests fail alike.

    fn hint_if_out_of_local_ports(&self, err: &azure_core::Error) {

        if !is_out_of_local_ports(err) {

            return;

        }

        {

            let now = Instant::now();

            let mut last = self.last_port_hint.lock();

            if last.is_some_and(|last| 
now4
.duration_since(last) < 
PORT_HINT_INTERVAL4
) {

                return;

            }

            *last = Some(now);

        }

        if self.max_concurrent_requests == 0 {

            event!(

                Level::WARN,

                container = %self.container,

                "Out of local ports toward the Azure endpoint (AddrNotAvailable), and this store's requests are unbounded. Set max_concurrent_requests on it to cap its connections; idle connections and ones in TIME_WAIT hold ports too"

            );

        } else {

            event!(

                Level::WARN,

                container = %self.container,

                max_concurrent_requests = self.max_concurrent_requests,

                "Out of local ports toward the Azure endpoint (AddrNotAvailable). Lower max_concurrent_requests, and count every store in a shard, each of which holds up to that many; idle connections and ones in TIME_WAIT hold ports too"

            );

        }

    }

    /// A read was answered (found or not): one toward growing the limit.

    fn note_read_ok(&self) {

        if let Some(
pool25
) = &self.read_permits {

            pool.record_success();

        
}14

    }

    /// A read was throttled: the limit is cut.

    fn note_read_throttled(&self) {

        if let Some(
pool3
) = &self.read_permits {

            pool.record_throttle();

        }

    }

    fn make_blob_path(&self, key: &StoreKey<'_>) -> String {

        format!("{}{}", self.blob_prefix, key.as_str())

    }

    /// How long a throttled response (503 `ServerBusy` or 429) asked us to

    /// wait, or `None` when the error is anything else. The wait is the

    /// `Retry-After` header when present, never less than the configured

    /// floor and never more than the configured cap. Every failed request

    /// passes through here, so this is also where running out of local

    /// ports is named.

    fn throttle_wait(&self, err: &azure_core::Error, op: &'static str) -> Option<Duration> {

        self.hint_if_out_of_local_ports(err);

        let ErrorKind::HttpResponse {

            status,

            error_code,

            raw_response,

        } = err.kind()

        else {

            return None;

        };

        let throttled = 
matches!6
(

            status,

            StatusCode::ServiceUnavailable | StatusCode::TooManyRequests

        ) || error_code.as_deref() == Some(StorageErrorCode::ServerBusy.as_ref());

        if !throttled {

            return None;

        }

        record_object_store_throttle("azure", op);

        let header_wait = raw_response.as_ref().and_then(|response| {

            let headers = response.headers();

            let millis = headers

                .get_optional_str(&RETRY_AFTER_MS)

                .or_else(|| headers.get_optional_str(&X_MS_RETRY_AFTER_MS))

                .and_then(|value| 
value.trim()0
.
parse0
::<u64>().
ok0
())

                .map(Duration::from_millis);

            millis.or_else(|| {

                headers

                    .get_optional_str(&RETRY_AFTER)

                    .and_then(|value| 
value.trim()2
.
parse2
::<u64>().
ok2
())

                    .map(Duration::from_secs)

            })

        });

        let asked = header_wait

            .unwrap_or(Duration::ZERO)

            .max(self.throttle_retry_delay);

        if asked > self.throttle_retry_max_delay {

            event!(

                Level::WARN,

                ?op,

                asked_s = asked.as_secs_f32(),

                cap_s = self.throttle_retry_max_delay.as_secs_f32(),

                "Throttled response asked for a longer wait than throttle_retry_max_delay_s; waiting the cap"

            );

            return Some(self.throttle_retry_max_delay);

        }

        Some(asked)

    }

    fn retry_later<T>(

        &self,

        sdk_err: &azure_core::Error,

        op: &'static str,

        err: Error,

    ) -> RetryResult<T> {

        match self.throttle_wait(sdk_err, op) {

            Some(wait) => RetryResult::RetryAfter(err, wait),

            None => RetryResult::Retry(err),

        }

    }

    async fn has(self: Pin<&Self>, digest: &StoreKey<'_>) -> Result<Option<u64>, Error> {

        let blob_path = self.make_blob_path(digest);

        let declared_size = match digest.borrow() {

            StoreKey::Digest(digest) => Some(digest.size_bytes()),

            StoreKey::Str(_) => None,

        };

        self.retrier

            .
retry159
(
unfold159
(
()159
, move |state| {

                let blob_path = blob_path.clone();

                let client = Arc::clone(&self.client);

                async move {

                    let _request_permit = self.request_permit().await;

                    let _permit = match fs::get_permit().await {

                        Ok(permit) => permit,

                        Err(e) => {

                            return Some((

                                RetryResult::Retry(make_err!(

                                    Code::Unavailable,

                                    "Failed to acquire permit: {e}"

                                )),

                                state,

                            ));

                        }

                    };

                    let result = client.blob_client(&blob_path).get_properties(None).await;

                    match result {

                        Ok(props) => {

                            if self.consider_expired_after_s > 0

                                && let Some(last_modified) = props.last_modified().ok().flatten()

                            {

                                let now = (self.now_fn)()

                                    .unix_timestamp()

                                    .try_into()

                                    .unwrap_or(i64::MAX);

                                if last_modified.unix_timestamp() + self.consider_expired_after_s

                                    <= now

                                {

                                    return Some((RetryResult::Ok(None), state));

                                }

                            }

                            let blob_size = props.content_length().ok().flatten().unwrap_or(0);

                            // An empty object cannot be a blob whose digest

                            // names bytes, compressed or not: a failed upload

                            // published it. Reported present it is read as

                            // DataLoss forever and never uploaded again, so

                            // say it is missing and let the client replace it.

                            if blob_size == 0 && 
declared_size1
.
is_some_and1
(|size| 
size1
 > 0) {

                                event!(

                                    Level::WARN,

                                    blob_path,

                                    "Azure holds an empty object for a digest that names bytes; reporting it missing"

                                );

                                return Some((RetryResult::Ok(None), state));

                            }

                            Some((RetryResult::Ok(Some(blob_size)), state))

                        }

                        Err(err) => {

                            if err.http_status() == Some(StatusCode::NotFound) {

                                // Distinguish a missing container (a config error) from a

                                // missing blob (a normal cache miss).

                                if let ErrorKind::HttpResponse {

                                    error_code: Some(error_code),

                                    ..

                                } = err.kind()

                                    && error_code == StorageErrorCode::ContainerNotFound.as_ref()

                                {

                                    return Some((

                                        RetryResult::Err(make_err!(

                                            Code::InvalidArgument,

                                            "Container not found: {err}"

                                        )),

                                        state,

                                    ));

                                }

                                Some((RetryResult::Ok(None), state))

                            } else {

                                Some((

                                    self.retry_later(

                                        &err,

                                        "has",

                                        make_err!(

                                            Code::Unavailable,

                                            "Failed to get blob properties: {err:?}"

                                        ),

                                    ),

                                    state,

                                ))

                            }

                        }

                    }

                }

            }))

            .await

    }

    async fn upload_staged_blocks(

        self: Pin<&Self>,

        blob_path: &str,

        mut reader: DropCloserReadHalf,

        upload_size: UploadSizeInfo,

    ) -> Result<u64, Error> {

        let blob_path = blob_path.to_owned();

        let max_size = match upload_size {

            UploadSizeInfo::ExactSize(
size10
) | UploadSizeInfo::MaxSize(
size2
) => size,

        };

        let block_size =

            cmp::min(max_size / (MAX_BLOCKS as u64 - 1), MAX_BLOCK_SIZE).max(DEFAULT_BLOCK_SIZE);

        let (tx, mut rx) = mpsc::channel(self.max_concurrent_uploads);

        let mut block_ids: Vec<Vec<u8>> = Vec::with_capacity(MAX_BLOCKS);

        let retrier = self.retrier.clone();

        let read_stream_fut = {

            let tx = tx.clone();

            let blob_path = blob_path.clone();

            async move {

                let mut total_uploaded = 0;

                for block_id in 
0..MAX_BLOCKS12
 {

                    let 
write_buf32
 = reader

                        .consume(Some(

                            usize::try_from(block_size)

                                .err_tip(|| "Could not convert block_size to usize")
?0
,

                        ))

                        .await

                        .err_tip(|| "Failed to read chunk in azure_store")
?0
;

                    if write_buf.is_empty() {

                        break;

                    }

                    total_uploaded += write_buf.len() as u64;

                    // Fixed-width, zero-padded ids keep the committed block list ordered

                    // after a lexicographic sort.

                    let block_id = format!("{block_id:032}").into_bytes();

                    let blob_path = blob_path.clone();

                    tx.send(async move {

                        self.retrier

                            .retry(unfold(

                                (write_buf, block_id),

                                move |(write_buf, block_id)| {

                                    let client = Arc::clone(&self.client);

                                    let blob_path = blob_path.clone();

                                    async move {

                                        let _request_permit = self.request_permit().await;

                                        let _permit = match fs::get_permit().await {

                                            Ok(permit) => permit,

                                            Err(e) => {

                                                return Some((

                                                    RetryResult::Retry(make_err!(

                                                        Code::Unavailable,

                                                        "Failed to acquire permit: {e}"

                                                    )),

                                                    (write_buf, block_id),

                                                ));

                                            }

                                        };

                                        let content_length = write_buf.len() as u64;

                                        let retry_result = client

                                            .blob_client(&blob_path)

                                            .block_blob_client()

                                            .stage_block(

                                                &block_id,

                                                content_length,

                                                RequestContent::from(write_buf.to_vec()),

                                                None,

                                            )

                                            .await

                                            .map_or_else(

                                                |e| {

                                                    self.retry_later(

                                                        &e,

                                                        "put",

                                                        make_err!(

                                                            Code::Aborted,

                                                            "Failed to upload block in Azure store: {e:?}"

                                                        ),

                                                    )

                                                },

                                                |_| RetryResult::Ok(block_id.clone()),

                                            );

                                        Some((retry_result, (write_buf, block_id)))

                                    }

                                },

                            ))

                            .await

                    })

                    .await

                    .map_err(|err| 
{0

                        Error::from_std_err(Code::Internal, &err)

                            .append("Failed to send block to channel")

                    })?;

                }

                Ok::<_, Error>(total_uploaded)

            }

            .fuse()

        };

        let mut upload_futures = FuturesUnordered::new();

        let mut total_uploaded = 0;

        tokio::pin!(read_stream_fut);

        let mut failure: Option<Error> = None;

        loop {

            if read_stream_fut.is_terminated() && 
rx49
.
is_empty49
() && 
upload_futures26
.
is_empty26
() {

                break;

            }

            tokio::select! {

                
result11
 = &mut read_stream_fut => match 
result11
 {

                    Ok(uploaded) => total_uploaded = uploaded,

                    Err(err) => {

                        failure = Some(err);

                        break;

                    }

                },

                Some(
block_id21
) = upload_futures.next() => match 
block_id21
 {

                    Ok(block_id) => block_ids.push(block_id),

                    Err(err) => {

                        failure = Some(err);

                        break;

                    }

                },

                Some(
fut21
) = rx.recv() => 
upload_futures21
.
push21
(
fut21
),

            }

        }

        if let Some(
err2
) = failure {

            // The read future is never polled again, so nothing new is

            // staged; blocks already on their way to Azure are allowed to

            // land so the discard that follows sees every one of them

            // instead of racing a request still in flight.

            drop(rx);

            while upload_futures.next().await.is_some() 
{}0

            return Err(err);

        }

        // Reject an incomplete transfer before publishing its staged blocks.

        // The caller discards them using the same cleanup as other failures.

        upload_size.validate_size(total_uploaded)
?6
;

        // Sorting block IDs to ensure consistent ordering of the committed blob.

        block_ids.sort_unstable();

        let block_list = BlockLookupList {

            latest: Some(block_ids),

            ..Default::default()

        };

        retrier

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

                let client = Arc::clone(&self.client);

                let blob_path = blob_path.clone();

                async move {

                    let _request_permit = self.request_permit().await;

                    let _permit = match fs::get_permit().await {

                        Ok(permit) => permit,

                        Err(e) => {

                            return Some((

                                RetryResult::Retry(make_err!(

                                    Code::Unavailable,

                                    "Failed to acquire permit: {e}"

                                )),

                                block_list,

                            ));

                        }

                    };

                    let blocks = match RequestContent::try_from(block_list.clone()) {

                        Ok(blocks) => blocks,

                        Err(e) => {

                            return Some((

                                RetryResult::Err(make_err!(

                                    Code::Internal,

                                    "Failed to serialize block list in Azure store: {e:?}"

                                )),

                                block_list,

                            ));

                        }

                    };

                    let retry_result = client

                        .blob_client(&blob_path)

                        .block_blob_client()

                        .commit_block_list(blocks, None)

                        .await

                        .map_or_else(

                            |e| {

                                self.retry_later(

                                    &e,

                                    "put",

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

                                        .append("Failed to commit block list in Azure store:"),

                                )

                            },

                            |_| RetryResult::Ok(total_uploaded),

                        );

                    Some((retry_result, block_list))

                }

            }))

            .await

    }

}

#[async_trait]

impl<I, NowFn> StoreDriver for AzureBlobStore<NowFn>

where

    I: InstantWrapper,

    NowFn: Fn() -> I + Send + Sync + Unpin + 'static,

{

    async fn post_init(self: Arc<Self>) -> Result<(), Error> {

        Ok(())

    }

    async fn has_with_results(

        self: Pin<&Self>,

        keys: &[StoreKey<'_>],

        results: &mut [Option<u64>],

    ) -> Result<(), Error> {

        keys.iter()

            .zip(results.iter_mut())

            .map(|(key, result)| async move {

                if is_zero_digest(key.borrow()) {

                    *result = Some(0);

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

                }

                *result = self.has(key).await
?6
;

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

            })

            .collect::<FuturesUnordered<_>>()

            .try_collect()

            .await

    }

    fn optimized_for(&self, optimization: StoreOptimizations) -> bool {

        matches!(optimization, StoreOptimizations::LazyExistenceOnSync)

    }

    fn max_upload_size(&self) -> u64 {

        MAX_BLOCKS as u64 * MAX_BLOCK_SIZE

    }

    async fn update(

        self: Pin<&Self>,

        digest: StoreKey<'_>,

        mut reader: DropCloserReadHalf,

        upload_size: UploadSizeInfo,

    ) -> Result<u64, Error> {

        let blob_path = self.make_blob_path(&digest);

        // Handling zero-sized content check

        if is_zero_digest(digest.borrow()) {

            let size = reader.drain().await?;

            // The empty digest constrains the content; the caller's transfer

            // declaration must independently agree with the received size.

            UploadSizeInfo::ExactSize(0).validate_size(size)?;

            upload_size.validate_size(size)?;

            return Ok(0);

        }

        let max_size = match upload_size {

            UploadSizeInfo::ExactSize(sz) | UploadSizeInfo::MaxSize(sz) => sz,

        };

        // For small files of a known size we buffer to `Bytes` and upload in a single request.

        if max_size < DEFAULT_BLOCK_SIZE && matches!(upload_size, UploadSizeInfo::ExactSize(_)) {

            let UploadSizeInfo::ExactSize(sz) = upload_size else {

                unreachable!("upload_size must be UploadSizeInfo::ExactSize here");

            };

            reader.set_max_recent_data_size(

                u64::try_from(self.max_retry_buffer_per_request)

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

            );

            return self

                .retrier

                .retry(unfold(reader, move |mut reader| {

                    let client = Arc::clone(&self.client);

                    let blob_path = blob_path.clone();

                    async move {

                        // Collect through EOF before issuing a PUT. A short or failed

                        // producer must never replace a valid object with partial data.

                        // No permit is held meanwhile: the bytes arrive at the

                        // producer's pace, and no connection is open yet.

                        let 
content37
 = match reader

                            .consume(Some(usize::try_from(sz).unwrap_or(usize::MAX).saturating_add(1)))

                            .await

                        {

                            Ok(content) => content,

                            Err(err) => return Some((RetryResult::Err(err), reader)),

                        };

                        if let Err(
err2
) = upload_size.validate_size(content.len() as u64) {

                            return Some((RetryResult::Err(err), reader));