// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{convert::TryInto, fmt::Display, io, sync::Arc};

use async_trait::async_trait;
use cloud::{
    blob::{BlobConfig, BlobStorage, BucketConf, PutResource, StringNonEmpty, none_to_empty},
    metrics,
};
use futures_util::{
    future::TryFutureExt,
    io::{self as async_io, AsyncRead, Cursor},
    stream::{StreamExt, TryStreamExt},
};
use http::{HeaderValue, Version};
use hyper::{Body, Client, Request, Response, StatusCode, client::HttpConnector};
use hyper_tls::HttpsConnector;
pub use kvproto::brpb::{CloudDynamic, Gcs as InputConfig};
use tame_gcs::{
    common::{PredefinedAcl, StorageClass},
    objects::{InsertObjectOptional, Metadata, Object},
    types::{BucketName, ObjectId},
};
use tame_oauth::gcp::{ServiceAccountInfo, ServiceAccountProvider, TokenOrRequest, TokenProvider};
use tikv_util::{
    stream::{AsyncReadAsSyncStreamOfBytes, RetryError, error_stream},
    time::Instant,
};

use crate::utils::retry;

const GOOGLE_APIS: &str = "https://www.googleapis.com";
const HARDCODED_ENDPOINTS_SUFFIX: &[&str] = &["upload/storage/v1/", "storage/v1/"];

#[derive(Clone, Debug)]
pub struct Config {
    bucket: BucketConf,
    predefined_acl: Option<PredefinedAcl>,
    storage_class: Option<StorageClass>,
    svc_info: Option<ServiceAccountInfo>,
}

impl Config {
    #[cfg(test)]
    pub fn default(bucket: BucketConf) -> Self {
        Self {
            bucket,
            predefined_acl: None,
            storage_class: None,
            svc_info: None,
        }
    }

    pub fn missing_credentials() -> io::Error {
        io::Error::new(io::ErrorKind::InvalidInput, "missing credentials")
    }

    pub fn from_cloud_dynamic(cloud_dynamic: &CloudDynamic) -> io::Result<Config> {
        let bucket = BucketConf::from_cloud_dynamic(cloud_dynamic)?;
        let attrs = &cloud_dynamic.attrs;
        let def = &String::new();
        let predefined_acl = parse_predefined_acl(attrs.get("predefined_acl").unwrap_or(def))
            .or_invalid_input("invalid predefined_acl")?;
        let storage_class = parse_storage_class(&none_to_empty(bucket.storage_class.clone()))
            .or_invalid_input("invalid storage_class")?;

        let credentials_blob_opt = StringNonEmpty::opt(
            attrs
                .get("credentials_blob")
                .unwrap_or(&"".to_string())
                .to_string(),
        );
        let svc_info = if let Some(cred) = credentials_blob_opt {
            Some(deserialize_service_account_info(cred)?)
        } else {
            None
        };

        Ok(Config {
            bucket,
            predefined_acl,
            svc_info,
            storage_class,
        })
    }

    pub fn from_input(input: InputConfig) -> io::Result<Config> {
        let endpoint = StringNonEmpty::opt(input.endpoint);
        let bucket = BucketConf {
            endpoint,
            bucket: StringNonEmpty::required_field(input.bucket, "bucket")?,
            prefix: StringNonEmpty::opt(input.prefix),
            storage_class: StringNonEmpty::opt(input.storage_class),
            region: None,
        };
        let predefined_acl = parse_predefined_acl(&input.predefined_acl)
            .or_invalid_input("invalid predefined_acl")?;
        let storage_class = parse_storage_class(&none_to_empty(bucket.storage_class.clone()))
            .or_invalid_input("invalid storage_class")?;
        let svc_info = if let Some(cred) = StringNonEmpty::opt(input.credentials_blob) {
            Some(deserialize_service_account_info(cred)?)
        } else {
            None
        };
        Ok(Config {
            bucket,
            predefined_acl,
            svc_info,
            storage_class,
        })
    }
}

fn deserialize_service_account_info(
    cred: StringNonEmpty,
) -> std::result::Result<ServiceAccountInfo, RequestError> {
    ServiceAccountInfo::deserialize(cred.to_string())
        .map_err(|e| RequestError::OAuth(e, "deserialize ServiceAccountInfo".to_string()))
}

impl BlobConfig for Config {
    fn name(&self) -> &'static str {
        STORAGE_NAME
    }

    fn url(&self) -> io::Result<url::Url> {
        self.bucket
            .url("gcs")
            .map_err(|s| io::Error::other(format!("error creating bucket url: {}", s)))
    }
}

// GCS compatible storage
#[derive(Clone)]
pub struct GcsStorage {
    config: Config,
    svc_access: Option<Arc<ServiceAccountProvider>>,
    client: Client<HttpsConnector<HttpConnector>, Body>,
}

trait ResultExt {
    type Ok;

    // Maps the error of this result as an `std::io::Error` with `InvalidInput`
    // error kind.
    fn or_invalid_input<D: Display>(self, msg: D) -> io::Result<Self::Ok>;
}

impl<T, E: Display> ResultExt for Result<T, E> {
    type Ok = T;
    fn or_invalid_input<D: Display>(self, msg: D) -> io::Result<T> {
        self.map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("{}: {}", msg, e)))
    }
}

#[derive(Debug)]
enum RequestError {
    Hyper(hyper::Error, String),
    Http(http::Error, String),
    OAuth(tame_oauth::Error, String),
    Gcs(tame_gcs::Error),
    InvalidHeaderValue(http::header::InvalidHeaderValue, String),
    InvalidEndpoint(http::uri::InvalidUri),
}

impl From<http::uri::InvalidUri> for RequestError {
    fn from(err: http::uri::InvalidUri) -> Self {
        Self::InvalidEndpoint(err)
    }
}

fn status_code_error(code: StatusCode, msg: String) -> RequestError {
    RequestError::OAuth(tame_oauth::Error::HttpStatus(to_oauth_status(code)), msg)
}

fn to_oauth_status(code: StatusCode) -> http1::StatusCode {
    http1::StatusCode::from_u16(code.as_u16()).unwrap_or(http1::StatusCode::INTERNAL_SERVER_ERROR)
}

fn from_oauth_version(version: http1::Version) -> Version {
    match version {
        http1::Version::HTTP_09 => Version::HTTP_09,
        http1::Version::HTTP_10 => Version::HTTP_10,
        http1::Version::HTTP_11 => Version::HTTP_11,
        http1::Version::HTTP_2 => Version::HTTP_2,
        http1::Version::HTTP_3 => Version::HTTP_3,
        _ => Version::HTTP_11,
    }
}

fn to_oauth_version(version: Version) -> http1::Version {
    match version {
        Version::HTTP_09 => http1::Version::HTTP_09,
        Version::HTTP_10 => http1::Version::HTTP_10,
        Version::HTTP_11 => http1::Version::HTTP_11,
        Version::HTTP_2 => http1::Version::HTTP_2,
        Version::HTTP_3 => http1::Version::HTTP_3,
        _ => http1::Version::HTTP_11,
    }
}

fn to_hyper_request(request: http1::Request<Vec<u8>>) -> Result<Request<Body>, RequestError> {
    let (parts, body) = request.into_parts();
    let mut builder = Request::builder()
        .method(parts.method.as_str())
        .uri(parts.uri.to_string())
        .version(from_oauth_version(parts.version));

    for (name, value) in &parts.headers {
        builder = builder.header(name.as_str(), value.as_bytes());
    }

    builder
        .body(Body::from(body))
        .map_err(|e| RequestError::Http(e, "build OAuth token request".to_string()))
}

fn to_oauth_response(
    parts: http::response::Parts,
    body: hyper::body::Bytes,
) -> Result<http1::Response<hyper::body::Bytes>, RequestError> {
    let mut builder = http1::Response::builder()
        .status(parts.status.as_u16())
        .version(to_oauth_version(parts.version));

    for (name, value) in &parts.headers {
        builder = builder.header(name.as_str(), value.as_bytes());
    }

    builder.body(body).map_err(|e| {
        RequestError::OAuth(
            tame_oauth::Error::Http(e),
            "build OAuth token response".to_string(),
        )
    })
}

fn to_hyper_header_value(
    value: http1::HeaderValue,
    msg: &'static str,
) -> Result<HeaderValue, RequestError> {
    HeaderValue::from_bytes(value.as_bytes())
        .map_err(|e| RequestError::InvalidHeaderValue(e, msg.to_string()))
}

impl From<RequestError> for io::Error {
    fn from(err: RequestError) -> Self {
        match err {
            RequestError::Hyper(e, msg) => {
                Self::new(io::ErrorKind::InvalidInput, format!("HTTP {}: {}", msg, e))
            }
            RequestError::Http(e, msg) => {
                Self::new(io::ErrorKind::InvalidInput, format!("HTTP {}: {}", msg, e))
            }
            RequestError::OAuth(tame_oauth::Error::Io(e), _) => e,
            RequestError::OAuth(tame_oauth::Error::HttpStatus(sc), msg) => {
                let fmt = format!("GCS OAuth: {}: {}", msg, sc);
                match sc.as_u16() {
                    401 | 403 => Self::new(io::ErrorKind::PermissionDenied, fmt),
                    404 => Self::new(io::ErrorKind::NotFound, fmt),
                    _ if sc.is_server_error() => Self::new(io::ErrorKind::Interrupted, fmt),
                    _ => Self::new(io::ErrorKind::InvalidInput, fmt),
                }
            }
            RequestError::OAuth(tame_oauth::Error::Auth(e), msg) => Self::new(
                io::ErrorKind::PermissionDenied,
                format!("authorization failed: {}: {}", msg, e),
            ),
            RequestError::OAuth(e, msg) => Self::new(
                io::ErrorKind::InvalidInput,
                format!("oauth failed: {}: {}", msg, e),
            ),
            RequestError::Gcs(e) => Self::new(
                io::ErrorKind::InvalidInput,
                format!("invalid GCS request: {}", e),
            ),
            RequestError::InvalidHeaderValue(e, msg) => Self::new(
                io::ErrorKind::InvalidInput,
                format!("invalid HTTP header value: {}: {}", msg, e),
            ),
            RequestError::InvalidEndpoint(e) => Self::new(
                io::ErrorKind::InvalidInput,
                format!("invalid GCS endpoint URI: {}", e),
            ),
        }
    }
}

impl RetryError for RequestError {
    fn is_retryable(&self) -> bool {
        match self {
            // FIXME: Inspect the error source?
            Self::Hyper(e, _) => {
                e.is_closed()
                    || e.is_connect()
                    || e.is_incomplete_message()
                    || e.is_body_write_aborted()
            }
            // See https://cloud.google.com/storage/docs/exponential-backoff.
            Self::OAuth(tame_oauth::Error::HttpStatus(status), _) => {
                status.as_u16() == StatusCode::TOO_MANY_REQUESTS.as_u16()
                    || status.as_u16() == StatusCode::REQUEST_TIMEOUT.as_u16()
                    || status.is_server_error()
            }
            // Consider everything else not retryable.
            _ => false,
        }
    }
}

impl GcsStorage {
    pub fn from_input(input: InputConfig) -> io::Result<Self> {
        Self::new(Config::from_input(input)?)
    }

    pub fn from_cloud_dynamic(cloud_dynamic: &CloudDynamic) -> io::Result<Self> {
        Self::new(Config::from_cloud_dynamic(cloud_dynamic)?)
    }

    /// Create a new GCS storage for the given config.
    pub fn new(config: Config) -> io::Result<GcsStorage> {
        let svc_access = if let Some(si) = &config.svc_info {
            Some(
                ServiceAccountProvider::new(si.clone())
                    .or_invalid_input("invalid credentials_blob")?,
            )
        } else {
            None
        };

        let client = Client::builder().build(HttpsConnector::new());
        Ok(GcsStorage {
            config,
            svc_access: svc_access.map(Arc::new),
            client,
        })
    }

    fn maybe_prefix_key(&self, key: &str) -> String {
        if let Some(prefix) = &self.config.bucket.prefix {
            return format!("{}/{}", prefix, key);
        }
        key.to_owned()
    }

    async fn set_auth(
        &self,
        req: &mut Request<Body>,
        scope: tame_gcs::Scopes,
        svc_access: Arc<ServiceAccountProvider>,
    ) -> Result<(), RequestError> {
        let token_or_request = svc_access
            .get_token(&[scope])
            .map_err(|e| RequestError::OAuth(e, "get_token".to_string()))?;
        let token = match token_or_request {
            TokenOrRequest::Token(token) => token,
            TokenOrRequest::Request {
                request,
                scope_hash,
                ..
            } => {
                let res = self
                    .client
                    .request(to_hyper_request(request)?)
                    .await
                    .map_err(|e| RequestError::Hyper(e, "set auth request".to_owned()))?;
                if !res.status().is_success() {
                    return Err(status_code_error(
                        res.status(),
                        "set auth request".to_string(),
                    ));
                }
                let (parts, body) = res.into_parts();
                let body = hyper::body::to_bytes(body)
                    .await
                    .map_err(|e| RequestError::Hyper(e, "set auth body".to_owned()))?;
                svc_access
                    .parse_token_response(scope_hash, to_oauth_response(parts, body)?)
                    .map_err(|e| RequestError::OAuth(e, "set auth parse token".to_string()))?
            }
        };
        let auth_header = token
            .try_into()
            .map_err(|e| RequestError::OAuth(e, "set auth add auth header".to_string()))
            .and_then(|value| to_hyper_header_value(value, "set auth add auth header"))?;
        req.headers_mut()
            .insert(http::header::AUTHORIZATION, auth_header);

        Ok(())
    }

    async fn make_request(
        &self,
        mut req: Request<Body>,
        scope: tame_gcs::Scopes,
    ) -> Result<Response<Body>, RequestError> {
        // replace the hard-coded GCS endpoint by the custom one.

        if let Some(endpoint) = &self.config.bucket.endpoint {
            let uri = req.uri().to_string();
            let new_url_opt = change_host(endpoint, &uri);
            if let Some(new_url) = new_url_opt {
                *req.uri_mut() = new_url.parse()?;
            }
        }

        if let Some(svc_access) = &self.svc_access {
            self.set_auth(&mut req, scope, svc_access.clone()).await?;
        }
        let uri = req.uri().to_string();
        let res = self
            .client
            .request(req)
            .await
            .map_err(|e| RequestError::Hyper(e, uri.clone()))?;
        if !res.status().is_success() {
            return Err(status_code_error(res.status(), uri));
        }
        Ok(res)
    }

    fn error_to_async_read<E>(kind: io::ErrorKind, e: E) -> cloud::blob::BlobStream<'static>
    where
        E: Into<Box<dyn std::error::Error + Send + Sync>>,
    {
        Box::new(error_stream(io::Error::new(kind, e)).into_async_read())
    }

    fn get_range(&self, name: &str, range: Option<String>) -> cloud::blob::BlobStream<'_> {
        let bucket = self.config.bucket.bucket.to_string();
        let name = self.maybe_prefix_key(name);
        debug!("read file from GCS storage"; "key" => %name);
        let oid = match ObjectId::new(bucket, name) {
            Ok(oid) => oid,
            Err(e) => return GcsStorage::error_to_async_read(io::ErrorKind::InvalidInput, e),
        };
        let mut request = match Object::download(&oid, None /* optional */) {
            Ok(request) => request.map(|_: io::Empty| Body::empty()),
            Err(e) => return GcsStorage::error_to_async_read(io::ErrorKind::Other, e),
        };
        if let Some(r) = range {
            let header_value = match HeaderValue::from_str(&r) {
                Ok(v) => v,
                Err(e) => return GcsStorage::error_to_async_read(io::ErrorKind::Other, e),
            };
            request.headers_mut().insert("Range", header_value);
        }
        Box::new(
            self.make_request(request, tame_gcs::Scopes::ReadOnly)
                .and_then(|response| async {
                    if response.status().is_success() {
                        Ok(response.into_body().map_err(|e| {
                            io::Error::other(format!("download from GCS error: {}", e))
                        }))
                    } else {
                        Err(status_code_error(
                            response.status(),
                            "bucket read".to_string(),
                        ))
                    }
                })
                .err_into::<io::Error>()
                .try_flatten_stream()
                .boxed() // this `.boxed()` pin the stream.
                .into_async_read(),
        )
    }
}

fn change_host(host: &StringNonEmpty, url: &str) -> Option<String> {
    let new_host = (|| {
        for hardcoded in HARDCODED_ENDPOINTS_SUFFIX {
            if let Some(res) = host.strip_suffix(hardcoded) {
                return StringNonEmpty::opt(res.to_owned()).unwrap();
            }
        }
        host.to_owned()
    })();
    if let Some(res) = url.strip_prefix(GOOGLE_APIS) {
        return Some([new_host.trim_end_matches('/'), res].concat());
    }
    None
}

// Convert manually since they don't implement FromStr.
fn parse_storage_class(sc: &str) -> Result<Option<StorageClass>, &str> {
    Ok(Some(match sc {
        "" => return Ok(None),
        "STANDARD" => StorageClass::Standard,
        "NEARLINE" => StorageClass::Nearline,
        "COLDLINE" => StorageClass::Coldline,
        "DURABLE_REDUCED_AVAILABILITY" => StorageClass::DurableReducedAvailability,
        "REGIONAL" => StorageClass::Regional,
        "MULTI_REGIONAL" => StorageClass::MultiRegional,
        _ => return Err(sc),
    }))
}

fn parse_predefined_acl(acl: &str) -> Result<Option<PredefinedAcl>, &str> {
    Ok(Some(match acl {
        "" => return Ok(None),
        "authenticatedRead" => PredefinedAcl::AuthenticatedRead,
        "bucketOwnerFullControl" => PredefinedAcl::BucketOwnerFullControl,
        "bucketOwnerRead" => PredefinedAcl::BucketOwnerRead,
        "private" => PredefinedAcl::Private,
        "projectPrivate" => PredefinedAcl::ProjectPrivate,
        "publicRead" => PredefinedAcl::PublicRead,
        _ => return Err(acl),
    }))
}

/// Like AsyncReadExt::read_to_end, but only try to initialize the buffer once.
/// Check https://github.com/rust-lang/futures-rs/issues/2658 for the reason we cannot
/// directly use it.
async fn read_to_end<R: AsyncRead>(r: R, v: &mut Vec<u8>) -> std::io::Result<u64> {
    let mut c = Cursor::new(v);
    async_io::copy(r, &mut c).await
}

const STORAGE_NAME: &str = "gcs";

#[async_trait]
impl BlobStorage for GcsStorage {
    fn config(&self) -> Box<dyn BlobConfig> {
        Box::new(self.config.clone()) as Box<dyn BlobConfig>
    }

    async fn put(&self, name: &str, reader: PutResource, content_length: u64) -> io::Result<()> {
        if content_length == 0 {
            // It is probably better to just write the empty file
            // However, currently going forward results in a body write aborted error
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "no content to write",
            ));
        }
        use std::convert::TryFrom;

        let key = self.maybe_prefix_key(name);
        debug!("save file to GCS storage"; "key" => %key);
        let bucket = BucketName::try_from(self.config.bucket.bucket.to_string())
            .or_invalid_input(format_args!("invalid bucket {}", self.config.bucket.bucket))?;

        let metadata = Metadata {
            name: Some(key),
            storage_class: self.config.storage_class,
            ..Default::default()
        };

        // FIXME: Switch to upload() API so we don't need to read the entire data into
        // memory in order to retry.
        let begin = Instant::now_coarse();
        let mut data = Vec::with_capacity(content_length as usize);
        read_to_end(reader, &mut data).await?;
        metrics::CLOUD_REQUEST_HISTOGRAM_VEC
            .with_label_values(&["gcp", "read_local"])
            .observe(begin.saturating_elapsed_secs());
        let begin = Instant::now_coarse();
        retry(
            || async {
                let data = Cursor::new(data.clone());
                let req = Object::insert_multipart(
                    &bucket,
                    data,
                    content_length,
                    &metadata,
                    Some(InsertObjectOptional {
                        predefined_acl: self.config.predefined_acl,
                        ..Default::default()
                    }),
                )
                .map_err(RequestError::Gcs)?
                .map(|reader| Body::wrap_stream(AsyncReadAsSyncStreamOfBytes::new(reader)));
                self.make_request(req, tame_gcs::Scopes::ReadWrite).await
            },
            "insert_multipart",
        )
        .await?;
        metrics::CLOUD_REQUEST_HISTOGRAM_VEC
            .with_label_values(&["gcp", "insert_multipart"])
            .observe(begin.saturating_elapsed_secs());
        Ok::<_, io::Error>(())
    }

    fn get(&self, name: &str) -> cloud::blob::BlobStream<'_> {
        self.get_range(name, None)
    }

    fn get_part(&self, name: &str, off: u64, len: u64) -> cloud::blob::BlobStream<'_> {
        // inclusive, bytes=0-499 -> [0, 499]
        self.get_range(name, Some(format!("bytes={}-{}", off, off + len - 1)))
    }
}

#[cfg(test)]
mod tests {
    extern crate test;
    use std::task::Poll;

    use futures_util::AsyncReadExt;
    use matches::assert_matches;

    use super::*;
    const HARDCODED_ENDPOINTS: &[&str] = &[
        "https://www.googleapis.com/upload/storage/v1",
        "https://www.googleapis.com/storage/v1",
    ];

    #[test]
    fn test_change_host() {
        let host = StringNonEmpty::static_str("http://localhost:4443");
        assert_eq!(
            &change_host(&host, &format!("{}/storage/v1/foo", GOOGLE_APIS)).unwrap(),
            "http://localhost:4443/storage/v1/foo"
        );

        let h1 = url::Url::parse(HARDCODED_ENDPOINTS[0]).unwrap();
        let h2 = url::Url::parse(HARDCODED_ENDPOINTS[1]).unwrap();

        let endpoint = StringNonEmpty::static_str("http://example.com");
        assert_eq!(
            &change_host(&endpoint, h1.as_str()).unwrap(),
            "http://example.com/upload/storage/v1"
        );
        assert_eq!(
            &change_host(&endpoint, h2.as_str()).unwrap(),
            "http://example.com/storage/v1"
        );
        assert_eq!(
            &change_host(&endpoint, &format!("{}/foo", h2)).unwrap(),
            "http://example.com/storage/v1/foo"
        );
        assert_matches!(&change_host(&endpoint, "foo"), None);

        // if we get the endpoint with suffix "/storage/v1/"
        let endpoint = StringNonEmpty::static_str("http://example.com/storage/v1/");
        assert_eq!(
            &change_host(&endpoint, &format!("{}/foo", h2)).unwrap(),
            "http://example.com/storage/v1/foo"
        );

        let endpoint = StringNonEmpty::static_str("http://example.com/upload/storage/v1/");
        assert_eq!(
            &change_host(&endpoint, &format!("{}/foo", h2)).unwrap(),
            "http://example.com/storage/v1/foo"
        );
    }

    #[test]
    fn test_parse_storage_class() {
        assert_matches!(
            parse_storage_class("STANDARD"),
            Ok(Some(StorageClass::Standard))
        );
        assert_matches!(parse_storage_class(""), Ok(None));
        assert_matches!(
            parse_storage_class("NOT_A_STORAGE_CLASS"),
            Err("NOT_A_STORAGE_CLASS")
        );
    }

    #[test]
    fn test_parse_acl() {
        // can't use assert_matches!(), PredefinedAcl doesn't even implement Debug.
        assert!(matches!(
            parse_predefined_acl("private"),
            Ok(Some(PredefinedAcl::Private))
        ));
        assert!(matches!(parse_predefined_acl(""), Ok(None)));
        assert!(matches!(parse_predefined_acl("notAnACL"), Err("notAnACL")));
    }

    #[test]
    fn test_url_of_backend() {
        let bucket_name = StringNonEmpty::static_str("bucket");
        let mut bucket = BucketConf::default(bucket_name);
        bucket.prefix = Some(StringNonEmpty::static_str("/backup 02/prefix/"));
        let gcs = Config::default(bucket.clone());
        // only 'bucket' and 'prefix' should be visible in url_of_backend()
        assert_eq!(
            gcs.url().unwrap().to_string(),
            "gcs://bucket/backup%2002/prefix/"
        );
        bucket.endpoint = Some(StringNonEmpty::static_str("http://endpoint.com"));
        assert_eq!(
            &Config::default(bucket).url().unwrap().to_string(),
            "http://endpoint.com/bucket/backup%2002/prefix/"
        );
    }

    enum ThrottleReadState {
        Spawning,
        Emitting,
    }
    /// ThrottleRead throttles a `Read` -- make it emits 2 chars for each
    /// `read` call. This is copy & paste from the implmentation from s3.rs.
    #[pin_project::pin_project]
    struct ThrottleRead<R> {
        #[pin]
        inner: R,
        state: ThrottleReadState,
    }
    impl<R: AsyncRead> AsyncRead for ThrottleRead<R> {
        fn poll_read(
            self: std::pin::Pin<&mut Self>,
            cx: &mut std::task::Context<'_>,
            buf: &mut [u8],
        ) -> Poll<io::Result<usize>> {
            let this = self.project();
            match this.state {
                ThrottleReadState::Spawning => {
                    *this.state = ThrottleReadState::Emitting;
                    cx.waker().wake_by_ref();
                    Poll::Pending
                }
                ThrottleReadState::Emitting => {
                    *this.state = ThrottleReadState::Spawning;
                    this.inner.poll_read(cx, &mut buf[..2])
                }
            }
        }
    }
    impl<R> ThrottleRead<R> {
        fn new(r: R) -> Self {
            Self {
                inner: r,
                state: ThrottleReadState::Spawning,
            }
        }
    }

    const BENCH_READ_SIZE: usize = 128 * 1024;

    // 255,120,895 ns/iter (+/- 73,332,249) (futures-util 0.3.15)
    #[bench]
    fn bench_read_to_end(b: &mut test::Bencher) {
        let mut v = [0; BENCH_READ_SIZE];
        let mut dst = Vec::with_capacity(BENCH_READ_SIZE);
        let rt = tokio::runtime::Builder::new_current_thread()
            .build()
            .unwrap();

        b.iter(|| {
            let mut r = ThrottleRead::new(Cursor::new(&mut v));
            dst.clear();

            rt.block_on(r.read_to_end(&mut dst)).unwrap();
            assert_eq!(dst.len(), BENCH_READ_SIZE)
        })
    }

    // 5,850,042 ns/iter (+/- 3,787,438)
    #[bench]
    fn bench_manual_read_to_end(b: &mut test::Bencher) {
        let mut v = [0; BENCH_READ_SIZE];
        let mut dst = Vec::with_capacity(BENCH_READ_SIZE);
        let rt = tokio::runtime::Builder::new_current_thread()
            .build()
            .unwrap();
        b.iter(|| {
            let r = ThrottleRead::new(Cursor::new(&mut v));
            dst.clear();

            rt.block_on(read_to_end(r, &mut dst)).unwrap();
            assert_eq!(dst.len(), BENCH_READ_SIZE)
        })
    }
}