use super::{BuildSnafu, Error, LOG_TARGET, RequestSnafu, SerdeSnafu, UriSnafu};
use axum::http::Method;
use axum::http::header::HeaderMap;
use axum::http::uri::Uri;
use bytes::Bytes;
use reqwest::Client as ReqwestClient;
use reqwest::RequestBuilder;
use scopeguard::defer;
use serde::Serialize;
use serde::de::DeserializeOwned;
use snafu::ResultExt;
use std::collections::HashMap;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Duration;
use tibba_util::{Stopwatch, json_get};
use tracing::info;
type Result<T> = std::result::Result<T, Error>;
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
const VERSION: &str = env!("CARGO_PKG_VERSION");
const EMPTY_QUERY: Option<&[(&str, &str)]> = None;
const EMPTY_BODY: Option<&[(&str, &str)]> = None;
#[derive(Clone, Debug, Default)]
pub struct Params<'a, Q, P>
where
Q: Serialize + ?Sized,
P: Serialize + ?Sized,
{
pub method: Method,
pub timeout: Option<Duration>,
pub query: Option<&'a Q>,
pub body: Option<&'a P>,
pub url: &'a str,
}
#[derive(Default, Clone, Debug)]
pub struct HttpStats {
pub method: String,
pub path: String,
pub remote_addr: String,
pub status: u16,
pub content_length: usize,
pub processing: u32,
pub transfer: u32,
pub serde: u32,
pub total: u32,
pub tls_version: String,
pub tls_not_before: String,
pub tls_not_after: String,
}
pub trait HttpInterceptor: Send + Sync {
fn fail(&self, _status: u16, _data: &Bytes) -> BoxFuture<'_, Result<()>> {
Box::pin(async { Ok(()) })
}
fn request(&self, req: RequestBuilder) -> BoxFuture<'_, Result<RequestBuilder>> {
Box::pin(async move { Ok(req) })
}
fn response(&self, data: Bytes) -> BoxFuture<'_, Result<Bytes>> {
Box::pin(async move { Ok(data) })
}
fn on_done(&self, _stats: &HttpStats, _err: Option<&Error>) -> BoxFuture<'_, Result<()>> {
Box::pin(async { Ok(()) })
}
}
pub fn handle_fail(service: &str, status: u16, data: &Bytes) -> Result<()> {
if status >= 400 {
let mut message = json_get(data, "message");
if message.is_empty() {
message = "unknown error".to_string();
}
return Err(Error::Common {
service: service.to_string(),
message,
});
}
Ok(())
}
pub struct CommonInterceptor {
service: String,
}
impl CommonInterceptor {
pub fn new(service: &str) -> Self {
Self {
service: service.to_string(),
}
}
}
impl HttpInterceptor for CommonInterceptor {
fn fail(&self, status: u16, data: &Bytes) -> BoxFuture<'_, Result<()>> {
let result = handle_fail(&self.service, status, data);
Box::pin(async move { result })
}
fn on_done(&self, stats: &HttpStats, err: Option<&Error>) -> BoxFuture<'_, Result<()>> {
let error = err.map(ToString::to_string);
let service = self.service.clone();
let method = stats.method.clone();
let path = stats.path.clone();
let status = stats.status;
let remote_addr = stats.remote_addr.clone();
let content_length = stats.content_length;
let processing = stats.processing;
let transfer = stats.transfer;
let serde = stats.serde;
let total = stats.total;
Box::pin(async move {
info!(
target: LOG_TARGET,
service,
method,
path,
status,
remote_addr,
content_length,
processing,
transfer,
serde,
total,
error,
);
Ok(())
})
}
}
struct ClientConfig {
service: String,
base_url: String,
read_timeout: Option<Duration>,
timeout: Option<Duration>,
connect_timeout: Option<Duration>,
pool_idle_timeout: Option<Duration>,
pool_max_idle_per_host: usize,
max_processing: Option<u32>,
headers: Option<HeaderMap>,
dns_overrides: Option<HashMap<String, Vec<SocketAddr>>>,
interceptors: Option<Vec<Box<dyn HttpInterceptor>>>,
}
pub struct ClientBuilder {
config: ClientConfig,
}
impl ClientBuilder {
pub fn new(service: &str) -> Self {
Self {
config: ClientConfig {
service: service.to_string(),
base_url: String::new(),
read_timeout: None,
timeout: None,
connect_timeout: None,
pool_idle_timeout: None,
pool_max_idle_per_host: 0,
headers: None,
interceptors: None,
max_processing: None,
dns_overrides: None,
},
}
}
#[must_use]
pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
self.config.base_url = base_url.into();
self
}
#[must_use]
pub fn with_interceptor(mut self, interceptor: Box<dyn HttpInterceptor>) -> Self {
self.config
.interceptors
.get_or_insert_with(Vec::new)
.push(interceptor);
self
}
#[must_use]
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.config.timeout = Some(timeout);
self
}
#[must_use]
pub fn with_read_timeout(mut self, read_timeout: Duration) -> Self {
self.config.read_timeout = Some(read_timeout);
self
}
#[must_use]
pub fn with_connect_timeout(mut self, connect_timeout: Duration) -> Self {
self.config.connect_timeout = Some(connect_timeout);
self
}
#[must_use]
pub fn with_pool_idle_timeout(mut self, pool_idle_timeout: Duration) -> Self {
self.config.pool_idle_timeout = Some(pool_idle_timeout);
self
}
#[must_use]
pub fn with_headers(mut self, headers: HeaderMap) -> Self {
self.config.headers = Some(headers);
self
}
#[must_use]
pub fn with_common_interceptor(self) -> Self {
let service = self.config.service.clone();
self.with_interceptor(Box::new(CommonInterceptor::new(&service)))
}
#[must_use]
pub fn with_pool_max_idle_per_host(mut self, pool_max_idle_per_host: usize) -> Self {
self.config.pool_max_idle_per_host = pool_max_idle_per_host;
self
}
#[must_use]
pub fn with_max_processing(mut self, max_processing: u32) -> Self {
self.config.max_processing = Some(max_processing);
self
}
#[must_use]
pub fn with_dns_overrides(mut self, dns_overrides: HashMap<String, Vec<SocketAddr>>) -> Self {
self.config.dns_overrides = Some(dns_overrides);
self
}
pub fn build(mut self) -> Result<Client> {
let mut builder = ReqwestClient::builder()
.user_agent(format!("tibba-request/{VERSION}"))
.referer(false);
if let Some(timeout) = self.config.timeout {
builder = builder.timeout(timeout);
}
if let Some(headers) = self.config.headers.take() {
builder = builder.default_headers(headers);
}
if let Some(read_timeout) = self.config.read_timeout {
builder = builder.read_timeout(read_timeout);
}
if let Some(connect_timeout) = self.config.connect_timeout {
builder = builder.connect_timeout(connect_timeout);
}
if let Some(pool_idle_timeout) = self.config.pool_idle_timeout {
builder = builder.pool_idle_timeout(pool_idle_timeout);
}
if self.config.pool_max_idle_per_host > 0 {
builder = builder.pool_max_idle_per_host(self.config.pool_max_idle_per_host);
}
if let Some(dns_overrides) = self.config.dns_overrides.take() {
for (host, addrs) in dns_overrides {
builder = builder.resolve_to_addrs(&host, &addrs);
}
}
builder = builder.tls_info(true);
let client = builder.build().context(BuildSnafu {
service: self.config.service.clone(),
})?;
Ok(Client {
client,
config: self.config,
processing: AtomicU32::new(0),
})
}
}
pub struct Client {
client: ReqwestClient,
config: ClientConfig,
processing: AtomicU32,
}
impl Client {
fn get_url(&self, url: &str) -> String {
if url.starts_with("http") {
url.to_string()
} else {
self.config.base_url.to_string() + url
}
}
async fn raw<Q, P>(&self, stats: &mut HttpStats, params: Params<'_, Q, P>) -> Result<Bytes>
where
Q: Serialize + ?Sized,
P: Serialize + ?Sized,
{
let processing = self.processing.fetch_add(1, Ordering::Relaxed) + 1;
defer! {
self.processing.fetch_sub(1, Ordering::Relaxed);
};
if let Some(max_processing) = self.config.max_processing
&& processing > max_processing
{
return Err(Error::Common {
service: self.config.service.clone(),
message: "too many requests".to_string(),
});
}
let url = self.get_url(params.url);
let uri = url.parse::<Uri>().context(UriSnafu {
service: self.config.service.clone(),
})?;
stats.path = uri.path().to_string();
stats.method = params.method.to_string();
let mut req = match params.method {
Method::POST => self.client.post(url),
Method::PUT => self.client.put(url),
Method::PATCH => self.client.patch(url),
Method::DELETE => self.client.delete(url),
_ => self.client.get(url),
};
if let Some(value) = params.timeout {
req = req.timeout(value);
}
if let Some(value) = params.query {
req = req.query(value);
}
if let Some(value) = params.body {
req = req.json(value);
}
if let Some(interceptors) = &self.config.interceptors {
for interceptor in interceptors {
req = interceptor.request(req).await?;
}
}
let process_done = Stopwatch::new();
let res = req.send().await.context(RequestSnafu {
service: self.config.service.clone(),
path: stats.path.clone(),
})?;
stats.processing = process_done.elapsed_ms();
if let Some(remote_addr) = res.remote_addr() {
stats.remote_addr = remote_addr.to_string();
}
let status = res.status().as_u16();
let transfer_done = Stopwatch::new();
let mut full = res.bytes().await.context(RequestSnafu {
service: self.config.service.clone(),
path: stats.path.clone(),
})?;
stats.transfer = transfer_done.elapsed_ms();
stats.content_length = full.len();
stats.status = status;
if let Some(interceptors) = &self.config.interceptors {
if status >= 400 {
for interceptor in interceptors {
interceptor.fail(status, &full).await?;
}
}
for interceptor in interceptors {
full = interceptor.response(full).await?;
}
}
Ok(full)
}
async fn do_request<Q, P, T>(
&self,
stats: &mut HttpStats,
params: Params<'_, Q, P>,
) -> Result<T>
where
Q: Serialize + ?Sized,
P: Serialize + ?Sized,
T: DeserializeOwned,
{
let full = self.raw(stats, params).await?;
let serde_done = Stopwatch::new();
let data = serde_json::from_slice(&full).context(SerdeSnafu {
service: self.config.service.clone(),
})?;
stats.serde = serde_done.elapsed_ms();
Ok(data)
}
async fn request<Q, P, T>(&self, params: Params<'_, Q, P>) -> Result<T>
where
Q: Serialize + ?Sized,
P: Serialize + ?Sized,
T: DeserializeOwned,
{
let mut stats = HttpStats::default();
let done = Stopwatch::new();
let result = self.do_request(&mut stats, params).await;
stats.total = done.elapsed_ms();
let err = result.as_ref().err();
if let Some(interceptors) = &self.config.interceptors {
for interceptor in interceptors {
interceptor.on_done(&stats, err).await?;
}
}
result
}
pub async fn request_raw<Q, P>(&self, params: Params<'_, Q, P>) -> Result<Bytes>
where
Q: Serialize + ?Sized,
P: Serialize + ?Sized,
{
let mut stats = HttpStats::default();
let done = Stopwatch::new();
let result = self.raw(&mut stats, params).await;
stats.total = done.elapsed_ms();
let err = result.as_ref().err();
if let Some(interceptors) = &self.config.interceptors {
for interceptor in interceptors {
interceptor.on_done(&stats, err).await?;
}
}
result
}
pub async fn get<T>(&self, url: &str) -> Result<T>
where
T: DeserializeOwned,
{
self.request(Params {
timeout: None,
method: Method::GET,
url,
query: EMPTY_QUERY,
body: EMPTY_BODY,
})
.await
}
pub async fn get_with_query<P, T>(&self, url: &str, query: &P) -> Result<T>
where
P: Serialize + ?Sized,
T: DeserializeOwned,
{
self.request(Params {
timeout: None,
method: Method::GET,
url,
query: Some(query),
body: EMPTY_BODY,
})
.await
}
pub async fn post<P, T>(&self, url: &str, json: &P) -> Result<T>
where
P: Serialize + ?Sized,
T: DeserializeOwned,
{
self.request(Params {
timeout: None,
method: Method::POST,
url,
query: EMPTY_QUERY,
body: Some(json),
})
.await
}
pub async fn post_with_query<P, Q, T>(&self, url: &str, json: &P, query: &Q) -> Result<T>
where
P: Serialize + ?Sized,
Q: Serialize + ?Sized,
T: DeserializeOwned,
{
self.request(Params {
timeout: None,
method: Method::POST,
url,
query: Some(query),
body: Some(json),
})
.await
}
pub async fn put<P, T>(&self, url: &str, json: &P) -> Result<T>
where
P: Serialize + ?Sized,
T: DeserializeOwned,
{
self.request(Params {
timeout: None,
method: Method::PUT,
url,
query: EMPTY_QUERY,
body: Some(json),
})
.await
}
pub async fn patch<P, T>(&self, url: &str, json: &P) -> Result<T>
where
P: Serialize + ?Sized,
T: DeserializeOwned,
{
self.request(Params {
timeout: None,
method: Method::PATCH,
url,
query: EMPTY_QUERY,
body: Some(json),
})
.await
}
pub async fn delete<T>(&self, url: &str) -> Result<T>
where
T: DeserializeOwned,
{
self.request(Params {
timeout: None,
method: Method::DELETE,
url,
query: EMPTY_QUERY,
body: EMPTY_BODY,
})
.await
}
pub fn get_processing(&self) -> u32 {
self.processing.load(Ordering::Relaxed)
}
}