hyperfile 0.11.2

use std::io::Result;
use std::sync::Arc;
use log::debug;
use aws_sdk_s3::Client;
use btree_ondisk::BlockLoader;
use crate::{BlockPtr, BMapUserData};
use crate::meta_format::BlockPtrFormat;
use crate::segment::Segment;
use crate::s3uri::S3Uri;
use crate::file::ReadTiming;
use crate::s3commons::S3Ops;

pub struct S3BlockLoader {
    pub client: Client,
    pub bucket: String,
    pub root_path: String,
    /// Shared with the file's staging handle so meta-block requests
    /// land in the same counters as data requests.
    read_timing: Arc<ReadTiming>,
    /// The most node requests of one `read_batch` in flight at once.
    concurrency: usize,
}

impl Clone for S3BlockLoader {
    fn clone(&self) -> Self {
        Self {
            client: self.client.to_owned(),
            bucket: self.bucket.clone(),
            root_path: self.root_path.clone(),
            read_timing: self.read_timing.clone(),
            concurrency: self.concurrency,
        }
    }
}

impl S3BlockLoader {
    pub fn new(client: &Client, bucket: &str, root_path: &str, read_timing: Arc<ReadTiming>) -> Self {
        Self {
            client: client.to_owned(),
            bucket: bucket.to_string(),
            root_path: root_path.to_string(),
            read_timing,
            concurrency: crate::DEFAULT_META_READ_MAX_CONCURRENCY,
        }
    }

    /// At most `n` node requests of one batch in flight at once; zero is one.
    pub fn with_concurrency(mut self, n: usize) -> Self {
        self.concurrency = n.max(1);
        self
    }
}

impl BlockLoader<BlockPtr> for S3BlockLoader {
    async fn read(&self, v: BlockPtr, buf: &mut [u8], user_data: u32) -> Result<Vec<(BlockPtr, Vec<u8>)>> {
        let meta_block_size = buf.len();
        let ud = BMapUserData::from_u32(user_data);
        // Metadata blocks live in the summary part, so a pointer to one carries
        // that part and the key follows it.
        let (segid, offset) = BlockPtrFormat::decode(&v, &ud.blk_ptr_format);
        let key = format!("{}/{}", self.root_path, Segment::segid_to_staging_file_id(segid));
        let end = offset + meta_block_size - 1;
        let range = format!("bytes={}-{}", offset, end);
        debug!("read s3://{}/{} from s3 at offset {} for range {}", self.bucket, &key, offset, range);
        let start = std::time::Instant::now();
        let res = S3Ops::do_get_object(&self.client, &self.bucket, &key, buf, Some(&range), false).await;
        self.read_timing.add_meta_get(meta_block_size, start.elapsed().as_nanos() as u64);
        let _ = res?;
        // TODO return more
        Ok(Vec::new())
    }

    /// The nodes of one tree level, at most `concurrency` requests at a time.
    async fn read_batch(&self, ids: &[BlockPtr], bufs: &mut [Vec<u8>], user_data: u32) -> Result<Vec<(BlockPtr, Vec<u8>)>> {
        use futures::stream::{self, StreamExt, TryStreamExt};
        assert_eq!(ids.len(), bufs.len(), "read_batch: ids and bufs length must match");
        // Made first and then run: a stream mapped by a closure here trips the
        // compiler's check that the whole future is `Send`.
        let reads: Vec<_> = ids.iter().copied().zip(bufs.iter_mut())
            .map(|(id, buf)| self.read(id, buf.as_mut_slice(), user_data)).collect();
        stream::iter(reads)
            .buffer_unordered(self.concurrency)
            .try_concat()
            .await
    }

    fn dup_from_new_path(self, new_path: &str) -> Self {
        let s3uri = S3Uri::parse(new_path).expect("input new path is not a valid s3 uri");
        let mut clone = self.clone();
        clone.bucket = s3uri.bucket.to_string();
        clone.root_path = s3uri.key.to_string();
        clone
    }
}