feat: add websocket backend for TransactorClient (#3)

Author: Serial-ATA

Cargo.toml

@@ -10,10 +10,8 @@ reqwest = { version = "0.12.15", default-features = false, features = [
     "json",
     "rustls-tls",
 ] }
-reqwest-middleware = { version = "0.4.2", features = ["json", "rustls-tls"] }
-reqwest-retry = { version = "0.7.0" }
-reqwest-ratelimit = "0.4.1"
 governor = { version = "0.10.0", features = ["std"] }
+reqwest-websocket = { version = "0.5.0", features = ["json"] }
 serde = "1.0.219"
 serde_json = "1.0.140"
 thiserror = "2.0.12"
@@ -27,6 +25,9 @@ config = "0.15.11"
 secrecy = { version = "0.10.3", features = ["serde"] }
 serde_with = "3.12.0"
 rand = "0.9.1"
+futures = "0.3.31"
+tokio_with_wasm = { version = "0.8.6", features = ["rt", "sync", "macros"] }
+tokio-stream = { version = "0.1.17", features = ["sync"] }
 
 actix-web = { version = "4.10.2", optional = true, features = ["rustls"] }
 rdkafka = { version = "0.38.0", optional = true, features = [
@@ -38,6 +39,14 @@ num-traits = "0.2.19"
 itoa = "1.0.15"
 ryu = "1.0.20"
 
+# Middleware
+reqwest-middleware = { version = "0.4.2", features = ["json", "rustls-tls"] }
+reqwest-retry = { version = "0.7.0", optional = true }
+reqwest-ratelimit = { version = "0.4.1", optional = true }
+
+[target.'cfg(target_family = "wasm")'.dependencies]
+wasmtimer = { version = "0.4.1" }
+
 [dev-dependencies]
 anyhow = "1.0.98"
 tokio = { version = "1", features = ["full"] }
@@ -46,4 +55,7 @@ tokio = { version = "1", features = ["full"] }
 default = ["reqwest_middleware"]
 actix = ["dep:actix-web"]
 kafka = ["dep:rdkafka"]
-reqwest_middleware = []
+reqwest_middleware = ["dep:reqwest-retry", "dep:reqwest-ratelimit"]
+
+[lints.clippy]
+result_large_err = "allow"

src/lib.rs

@@ -32,13 +32,21 @@ pub enum Error {
     #[error(transparent)]
     Reqwest(#[from] reqwest::Error),
 
+    #[error(transparent)]
+    Ws(#[from] reqwest_websocket::Error),
+
     #[error(transparent)]
     ReqwestMiddleware(#[from] reqwest_middleware::Error),
 
     #[cfg(feature = "kafka")]
     #[error(transparent)]
     Kafka(#[from] rdkafka::error::KafkaError),
 
+    #[error("Subscription task panicked")]
+    SubscriptionFailed,
+    #[error("Subscription task lagged and was forcibly disconnected")]
+    SubscriptionLagged,
+
     #[error(transparent)]
     Url(#[from] url::ParseError),
 

src/services/event.rs

@@ -0,0 +1,11 @@
+use serde_json::Value;
+
+pub trait Class {
+    const CLASS: &'static str;
+}
+
+pub trait Event: Class {
+    fn matches(value: &Value) -> bool {
+        value.get("_class").and_then(|v| v.as_str()) == Some(Self::CLASS)
+    }
+}

src/services/mod.rs

@@ -15,8 +15,10 @@
 
 pub mod account;
 pub mod core;
+pub mod event;
 pub mod jwt;
 pub mod kvs;
+mod rpc;
 pub mod transactor;
 
 pub use reqwest_middleware::{ClientWithMiddleware as HttpClient, RequestBuilder};
@@ -27,20 +29,28 @@ use std::time::Duration;
 use reqwest::{self, Response, Url};
 use reqwest::{StatusCode, header::HeaderValue};
 use reqwest_middleware::ClientBuilder;
-use reqwest_retry::{
-    RetryTransientMiddleware, Retryable, RetryableStrategy, default_on_request_failure,
-    policies::ExponentialBackoff,
-};
 use secrecy::SecretString;
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use serde_json::{self as json, Value};
 use tracing::*;
 
 use crate::services::core::{AccountUuid, WorkspaceUuid};
+use crate::services::transactor::backend::http::HttpBackend;
+use crate::services::transactor::backend::ws::{WsBackend, WsBackendOpts};
+use crate::{Error, Result, config::Config};
+use account::AccountClient;
+use jwt::Claims;
+use kvs::KvsClient;
+use transactor::TransactorClient;
+
 #[cfg(feature = "kafka")]
 use crate::services::transactor::kafka;
-use crate::{Error, Result, config::Config};
-use {account::AccountClient, jwt::Claims, kvs::KvsClient, transactor::TransactorClient};
+
+#[cfg(feature = "reqwest_middleware")]
+use reqwest_retry::{
+    RetryTransientMiddleware, Retryable, RetryableStrategy, default_on_request_failure,
+    policies::ExponentialBackoff,
+};
 
 pub trait RequestBuilderExt {
     fn send_ext(self) -> impl Future<Output = Result<Response>>;
@@ -54,11 +64,12 @@ pub trait BasePathProvider {
     fn provide_base_path(&self) -> &Url;
 }
 
-pub trait ForceHttpScheme {
+pub trait ForceScheme {
     fn force_http_scheme(self) -> Url;
+    fn force_ws_scheme(self) -> Url;
 }
 
-impl ForceHttpScheme for Url {
+impl ForceScheme for Url {
     fn force_http_scheme(mut self) -> Url {
         match self.scheme() {
             "ws" => {
@@ -74,6 +85,24 @@ impl ForceHttpScheme for Url {
 
         self
     }
+
+    fn force_ws_scheme(mut self) -> Url {
+        match self.scheme() {
+            "http" => {
+                self.set_scheme("ws").unwrap();
+            }
+
+            "https" => {
+                self.set_scheme("wss").unwrap();
+            }
+
+            "ws" | "wss" => {}
+
+            _ => panic!(),
+        };
+
+        self
+    }
 }
 
 impl RequestBuilderExt for RequestBuilder {
@@ -116,13 +145,13 @@ fn from_value<T: DeserializeOwned>(value: Value) -> Result<T> {
 pub trait JsonClient {
     fn get<U: TokenProvider, R: DeserializeOwned>(
         &self,
-        user: U,
+        user: &U,
         url: Url,
     ) -> impl Future<Output = Result<R>>;
 
     fn post<U: TokenProvider, Q: Serialize, R: DeserializeOwned>(
         &self,
-        user: U,
+        user: &U,
         url: Url,
         body: &Q,
     ) -> impl Future<Output = Result<R>>;
@@ -134,7 +163,7 @@ impl JsonClient for HttpClient {
         skip(self, user, url),
         fields(%url, method = "get", type = "json")
     )]
-    async fn get<U: TokenProvider, R: DeserializeOwned>(&self, user: U, url: Url) -> Result<R> {
+    async fn get<U: TokenProvider, R: DeserializeOwned>(&self, user: &U, url: Url) -> Result<R> {
         trace!("request");
 
         let mut request = self.get(url.clone());
@@ -148,7 +177,7 @@ impl JsonClient for HttpClient {
 
     async fn post<U: TokenProvider, Q: Serialize, R: DeserializeOwned>(
         &self,
-        user: U,
+        user: &U,
         url: Url,
         body: &Q,
     ) -> Result<R> {
@@ -170,7 +199,7 @@ impl JsonClient for HttpClient {
     }
 }
 
-#[derive(Deserialize, Debug, Clone, strum::Display)]
+#[derive(Serialize, Deserialize, Debug, Clone, strum::Display)]
 #[serde(rename_all = "UPPERCASE")]
 pub enum Severity {
     Ok,
@@ -179,11 +208,11 @@ pub enum Severity {
     Error,
 }
 
-#[derive(Deserialize, Debug, Clone, thiserror::Error)]
+#[derive(Serialize, Deserialize, Debug, Clone, thiserror::Error)]
 pub struct Status {
     pub severity: Severity,
     pub code: String,
-    pub params: HashMap<String, String>,
+    pub params: HashMap<String, Value>,
 }
 
 impl std::fmt::Display for Status {
@@ -431,7 +460,11 @@ impl ServiceFactory {
         )
     }
 
-    pub fn new_transactor_client(&self, base: Url, claims: &Claims) -> Result<TransactorClient> {
+    pub fn new_transactor_client(
+        &self,
+        base: Url,
+        claims: &Claims,
+    ) -> Result<TransactorClient<HttpBackend>> {
         TransactorClient::new(
             self.transactor_http.clone(),
             base,
@@ -445,15 +478,45 @@ impl ServiceFactory {
         )
     }
 
+    pub async fn new_transactor_client_ws(
+        &self,
+        base: Url,
+        claims: &Claims,
+        opts: WsBackendOpts,
+    ) -> Result<TransactorClient<WsBackend>> {
+        TransactorClient::new_ws(
+            base,
+            claims.workspace()?,
+            claims.encode(
+                self.config
+                    .token_secret
+                    .as_ref()
+                    .ok_or(Error::Other("NoSecret"))?,
+            )?,
+            opts,
+        )
+        .await
+    }
+
     pub fn new_transactor_client_from_token(
         &self,
         base: Url,
         workspace: WorkspaceUuid,
         token: impl Into<SecretString>,
-    ) -> Result<TransactorClient> {
+    ) -> Result<TransactorClient<HttpBackend>> {
         TransactorClient::new(self.transactor_http.clone(), base, workspace, token)
     }
 
+    pub async fn new_transactor_client_ws_from_token(
+        &self,
+        base: Url,
+        workspace: WorkspaceUuid,
+        token: impl Into<SecretString>,
+        opts: WsBackendOpts,
+    ) -> Result<TransactorClient<WsBackend>> {
+        TransactorClient::new_ws(base, workspace, token, opts).await
+    }
+
     #[cfg(feature = "kafka")]
     pub fn new_kafka_publisher(&self, topic: &str) -> Result<kafka::KafkaProducer> {
         kafka::KafkaProducer::new(&self.config, topic)

src/services/rpc/mod.rs

@@ -0,0 +1,105 @@
+pub mod util;
+
+use crate::services::Status;
+use crate::services::core::Account;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)]
+#[serde(untagged, rename_all = "camelCase")]
+pub enum ReqId {
+    Str(String),
+    Num(i32),
+}
+
+impl From<String> for ReqId {
+    fn from(s: String) -> Self {
+        ReqId::Str(s)
+    }
+}
+
+impl From<i32> for ReqId {
+    fn from(i: i32) -> Self {
+        ReqId::Num(i)
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct RateLimitInfo {
+    pub remaining: u32,
+    pub limit: u32,
+    pub current: u32,
+    pub reset: f64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub retry_after: Option<u32>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct Chunk {
+    pub index: u32,
+    pub r#final: bool,
+}
+
+#[derive(Serialize, Deserialize, Default, Debug, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct Response<R> {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub result: Option<R>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub id: Option<ReqId>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub error: Option<Status>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub terminate: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub rate_limit: Option<RateLimitInfo>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub chunk: Option<Chunk>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub time: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bfst: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub queue: Option<u32>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct Request<P> {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub id: Option<ReqId>,
+    pub method: String,
+    pub params: Vec<P>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub time: Option<f64>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct HelloRequest {
+    #[serde(flatten)]
+    pub request: Request<()>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub binary: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub compression: Option<bool>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct HelloResponse {
+    #[serde(flatten)]
+    pub response: Response<String>,
+    pub binary: bool,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub reconnect: Option<bool>,
+    pub server_version: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub last_tx: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub last_hash: Option<String>,
+    pub account: Account,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub use_compression: Option<bool>,
+}