use std::sync::Arc;
use std::time::Duration;
use parking_lot::Mutex;
use super::types::VectorizerValue;
pub const DEFAULT_RPC_PORT: u16 = 15503;
const MAX_FRAME_BYTES: usize = 512 * 1024 * 1024;
pub fn protocol_config() -> thunder::Config {
use thunder::wire::config::{ErrorConvention, Handshake, HelloStyle, PushPolicy};
thunder::Config::standard()
.scheme("vectorizer")
.port(DEFAULT_RPC_PORT)
.handshake(Handshake::AuthCommand)
.hello_style(HelloStyle::NotUsed)
.push(PushPolicy::Reserved)
.error_codes(ErrorConvention::Resp3Prefixes)
.max_frame_bytes(MAX_FRAME_BYTES)
}
#[derive(Debug, thiserror::Error)]
pub enum RpcClientError {
#[error("connection error: {0}")]
Connection(String),
#[error("server error: {0}")]
Server(String),
#[error("not authenticated: {0}")]
NotAuthenticated(String),
#[error("timed out")]
Timeout,
#[error("protocol error: {0}")]
Protocol(String),
}
impl From<thunder::ClientError> for RpcClientError {
fn from(err: thunder::ClientError) -> Self {
use thunder::ClientError;
match err {
ClientError::Auth { message } => Self::NotAuthenticated(message),
ClientError::Server { message, .. } => Self::Server(message),
ClientError::Connection { message } => Self::Connection(message),
ClientError::Timeout => Self::Timeout,
ClientError::FrameTooLarge { message } | ClientError::Decode { message } => {
Self::Protocol(message)
}
}
}
}
pub type Result<T> = std::result::Result<T, RpcClientError>;
#[derive(Debug, Clone, Default)]
pub struct HelloPayload {
pub token: Option<String>,
pub api_key: Option<String>,
pub client_name: Option<String>,
pub version: i64,
}
impl HelloPayload {
pub fn new(client_name: impl Into<String>) -> Self {
Self {
client_name: Some(client_name.into()),
version: 1,
..Default::default()
}
}
pub fn with_token(mut self, token: impl Into<String>) -> Self {
self.token = Some(token.into());
self.api_key = None;
self
}
pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
self.api_key = Some(api_key.into());
self.token = None;
self
}
fn credentials(&self) -> Option<thunder::client::Credentials> {
if let Some(token) = &self.token {
return Some(thunder::client::Credentials::Token(token.clone()));
}
self.api_key
.as_ref()
.map(|key| thunder::client::Credentials::ApiKey(key.clone()))
}
fn into_value(self) -> VectorizerValue {
let mut pairs = vec![(
VectorizerValue::Str("version".into()),
VectorizerValue::Int(self.version),
)];
if let Some(token) = self.token {
pairs.push((
VectorizerValue::Str("token".into()),
VectorizerValue::Str(token),
));
}
if let Some(api_key) = self.api_key {
pairs.push((
VectorizerValue::Str("api_key".into()),
VectorizerValue::Str(api_key),
));
}
if let Some(name) = self.client_name {
pairs.push((
VectorizerValue::Str("client_name".into()),
VectorizerValue::Str(name),
));
}
VectorizerValue::Map(pairs)
}
}
#[derive(Debug, Clone)]
pub struct HelloResponse {
pub server_version: String,
pub protocol_version: i64,
pub authenticated: bool,
pub admin: bool,
pub capabilities: Vec<String>,
}
impl HelloResponse {
fn parse(value: &VectorizerValue) -> Self {
let server_version = value
.map_get("server_version")
.and_then(|v| v.as_str())
.map(str::to_owned)
.unwrap_or_default();
let protocol_version = value
.map_get("protocol_version")
.and_then(|v| v.as_int())
.unwrap_or(0);
let authenticated = value
.map_get("authenticated")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let admin = value
.map_get("admin")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let capabilities = value
.map_get("capabilities")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(str::to_owned))
.collect()
})
.unwrap_or_default();
Self {
server_version,
protocol_version,
authenticated,
admin,
capabilities,
}
}
}
pub struct RpcClient {
endpoint: String,
client_config: Mutex<thunder::ClientConfig>,
client: Mutex<Arc<thunder::Client>>,
redial: tokio::sync::Mutex<()>,
}
impl RpcClient {
pub async fn connect_url(url: &str) -> Result<Self> {
use super::endpoint::{Endpoint, parse_endpoint};
match parse_endpoint(url).map_err(|e| RpcClientError::Connection(e.to_string()))? {
Endpoint::Rpc { host, port } => Self::connect(format!("{host}:{port}")).await,
Endpoint::Rest { url } => Err(RpcClientError::Connection(format!(
"RpcClient cannot dial REST URL '{url}'; \
use the HTTP client (`vectorizer_sdk::VectorizerClient`) instead, \
or pass a `vectorizer://` URL"
))),
}
}
pub async fn connect(addr: impl AsRef<str>) -> Result<Self> {
let endpoint = addr.as_ref().to_owned();
let client_config = thunder::ClientConfig::new()
.client_name(concat!("vectorizer-sdk-rust/", env!("CARGO_PKG_VERSION")));
let client = Self::dial(&endpoint, client_config.clone()).await?;
Ok(Self {
endpoint,
client_config: Mutex::new(client_config),
client: Mutex::new(client),
redial: tokio::sync::Mutex::new(()),
})
}
pub async fn with_timeout(&self, timeout: Duration) -> Result<()> {
let config = {
let current = self.client_config.lock().clone();
current.connect_timeout(timeout).call_timeout(timeout)
};
self.replace_connection(config).await
}
async fn dial(endpoint: &str, config: thunder::ClientConfig) -> Result<Arc<thunder::Client>> {
thunder::Client::connect_with(endpoint, protocol_config(), config)
.await
.map(Arc::new)
.map_err(RpcClientError::from)
}
async fn replace_connection(&self, config: thunder::ClientConfig) -> Result<()> {
let _guard = self.redial.lock().await;
let fresh = Self::dial(&self.endpoint, config.clone()).await?;
*self.client_config.lock() = config;
*self.client.lock() = fresh;
Ok(())
}
fn client(&self) -> Arc<thunder::Client> {
Arc::clone(&self.client.lock())
}
pub async fn hello(&self, payload: HelloPayload) -> Result<HelloResponse> {
if let Some(credentials) = payload.credentials() {
let mut config = self.client_config.lock().clone();
config.credentials = Some(credentials);
if let Some(name) = &payload.client_name {
config = config.client_name(name.clone());
}
self.replace_connection(config).await?;
}
let result = self.call("HELLO", vec![payload.into_value()]).await?;
Ok(HelloResponse::parse(&result))
}
pub async fn ping(&self) -> Result<String> {
let result = self.call("PING", vec![]).await?;
result
.as_str()
.map(str::to_owned)
.ok_or_else(|| RpcClientError::Server("PING returned non-string payload".into()))
}
pub async fn call(
&self,
command: impl Into<String>,
args: Vec<VectorizerValue>,
) -> Result<VectorizerValue> {
self.client()
.call(command.into(), args)
.await
.map_err(RpcClientError::from)
}
pub fn is_authenticated(&self) -> bool {
self.client().is_authenticated()
}
pub async fn close(self) {
self.client().close().await;
}
}