use super::config::Config;
use crate::{
entities::ws::{
DocumentSubscribeResponse, Empty, EventSubscription, PhxReply, ReceivePhoenixMessage,
SendPhoenixMessage,
},
Auth, Subscription, WebsocketConnectionError,
};
use async_tungstenite::{tokio::connect_async, tungstenite::Message};
use backoff::{backoff::Backoff, ExponentialBackoff};
use futures::{future::BoxFuture, FutureExt, SinkExt, Stream};
use serde::{de::DeserializeOwned, Serialize};
use serde_json::{json, Value};
use std::{collections::HashMap, fmt::Debug, time::Duration};
use tokio::{
select,
sync::{broadcast, mpsc},
task,
time::{sleep, timeout},
};
use tokio_stream::{
wrappers::{BroadcastStream, ReceiverStream},
StreamExt,
};
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use uuid::Uuid;
#[derive(Debug)]
enum SubOp {
AddSubscription(String, SubscriptionRef),
RemoveSubscription(String),
}
type CloseState = (
mpsc::Receiver<Message>,
mpsc::UnboundedReceiver<SubOp>,
HashMap<String, SubscriptionRef>,
);
pub(super) struct Socket {
auth: Auth,
config: Config,
join_ref: Uuid,
outgoing_messages: (mpsc::Sender<Message>, Option<mpsc::Receiver<Message>>),
incoming_messages: (
broadcast::Sender<ReceivePhoenixMessage<Value>>,
broadcast::Receiver<ReceivePhoenixMessage<Value>>,
),
subscriptions: Option<HashMap<String, SubscriptionRef>>,
sub_ops: (
mpsc::UnboundedSender<SubOp>,
Option<mpsc::UnboundedReceiver<SubOp>>,
),
cancellation_token: CancellationToken,
handle: Option<BoxFuture<'static, Result<CloseState, WebsocketConnectionError>>>,
}
impl Debug for Socket {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Socket")
.field("auth", &self.auth)
.field("config", &self.config)
.field("join_ref", &self.join_ref)
.finish()
}
}
impl Socket {
pub fn new(auth: Auth, config: Config) -> Self {
let (outgoing_messages_sender, outgoing_messages_receiver) =
mpsc::channel(config.outgoing_capacity);
let incoming_messages = broadcast::channel(config.incoming_capacity);
let (sub_ops_sender, sub_ops_receiver) = mpsc::unbounded_channel();
Self {
auth,
config,
join_ref: Uuid::new_v4(),
outgoing_messages: (outgoing_messages_sender, Some(outgoing_messages_receiver)),
incoming_messages,
subscriptions: Some(Default::default()),
sub_ops: (sub_ops_sender, Some(sub_ops_receiver)),
cancellation_token: CancellationToken::new(),
handle: None,
}
}
pub fn client(&self) -> SocketClient {
SocketClient {
join_ref: self.join_ref,
outgoing_messages: self.outgoing_messages.0.clone(),
incoming_messages: self.incoming_messages.0.clone(),
sub_ops: self.sub_ops.0.clone(),
request_timeout: self.config.request_timeout,
cancellation_token: self.cancellation_token.clone(),
}
}
pub async fn connect(&mut self) -> Result<(), WebsocketConnectionError> {
let mut query = vec![("vsn", self.config.version.clone())];
match &self.auth {
Auth::ClientId(client_id) => query.push(("client_id", client_id.clone())),
Auth::AccessToken(token) => query.push(("token", token.clone())),
Auth::RefreshableAccessToken(token) => {
let access_token = token.access_token().await?;
query.push(("token", access_token.access_token));
}
Auth::ClientCredentials(client_credentials) => {
let access_token = client_credentials.access_token().await?;
query.push(("token", access_token.access_token));
}
}
let query_str = serde_urlencoded::to_string(query.as_slice())?;
let connection_url = format!("{}?{}", self.config.api_url, query_str);
let (ws_stream, _) = connect_async(&connection_url).await?;
let (mut ws_tx, mut ws_rx) = futures::StreamExt::split(ws_stream);
let cancellation_token = self.cancellation_token.child_token();
let outgoing_messages_handle = {
let mut outgoing_messages_receiver = self
.outgoing_messages
.1
.take()
.ok_or(WebsocketConnectionError::AlreadyConnected)?;
let cancellation_token = cancellation_token.clone();
task::spawn(async move {
loop {
select! {
_ = cancellation_token.cancelled() => {
tracing::trace!("received cancellation signal");
break;
}
msg = outgoing_messages_receiver.recv() => {
match msg {
Some(msg) => {
tracing::trace!(?msg, "sending message");
if let Err(err) = ws_tx.send(msg).await {
tracing::error!(?err, "failed to send message on the socket");
cancellation_token.cancel();
break;
}
}
None => {
tracing::trace!("all senders were dropped");
cancellation_token.cancel();
break;
}
}
}
}
}
outgoing_messages_receiver
})
.instrument(tracing::trace_span!("outgoing_messages"))
};
let incoming_messages_handle = {
let cancellation_token = cancellation_token.clone();
let incoming_messages_sender = self.incoming_messages.0.clone();
task::spawn(async move {
loop {
select! {
_ = cancellation_token.cancelled() => {
tracing::trace!("received cancellation signal");
break;
}
msg = ws_rx.next() => {
match msg {
Some(Ok(Message::Text(text))) => {
match serde_json::from_str::<ReceivePhoenixMessage<Value>>(&text) {
Ok(msg) => {
if msg.event == "phx_error" {
tracing::error!(?msg.payload, "error on socket");
cancellation_token.cancel();
break;
}
tracing::trace!(?msg, "incoming message");
if let Err(err) = incoming_messages_sender.send(msg) {
tracing::error!(?text, ?err, "failed to broadcast incoming message");
}
}
Err(err) => {
tracing::error!(?text, ?err, "failed to deserialize glimesh message");
}
}
}
Some(Ok(Message::Close(reason))) => {
tracing::error!(?reason, "socket closed");
cancellation_token.cancel();
break;
}
Some(Ok(frame)) => {
tracing::error!(?frame, "unexpected frame type");
cancellation_token.cancel();
break;
}
Some(Err(err)) => {
tracing::error!(?err, "socket error");
cancellation_token.cancel();
break;
}
None => {
tracing::error!("no more socket messages");
cancellation_token.cancel();
break;
}
}
}
}
}
})
.instrument(tracing::trace_span!("incoming_messages"))
};
let socket_client = self.client();
if let Err(err) = socket_client
.request::<_, Empty>("__absinthe__:control".into(), "phx_join".into(), Empty {})
.await
{
tracing::error!(?err, "join request failed");
cancellation_token.cancel();
return Err(err);
}
let pinger_handle = {
let ping_interval = Duration::from_secs(30);
let cancellation_token = cancellation_token.clone();
task::spawn(async move {
loop {
select! {
_ = cancellation_token.cancelled() => {
tracing::trace!("received cancellation signal");
break;
}
_ = sleep(ping_interval) => {
if let Err(err) = socket_client.request::<_, Empty>(
"phoenix".into(),
"heartbeat".into(),
Empty {},
)
.await {
tracing::error!(?err, "failed to send ping");
cancellation_token.cancel();
break;
}
}
};
}
})
.instrument(tracing::trace_span!("pinger"))
};
let subscriptions_handle = {
let socket_client = self.client();
let mut sub_ops_receiver = self
.sub_ops
.1
.take()
.ok_or(WebsocketConnectionError::AlreadyConnected)?;
let mut subscriptions = self
.subscriptions
.take()
.ok_or(WebsocketConnectionError::AlreadyConnected)?;
task::spawn(async move {
let sub_ids = subscriptions.keys().cloned().collect::<Vec<_>>();
for old_sub_id in sub_ids {
if cancellation_token.is_cancelled() {
break;
}
let sub = subscriptions.remove(&old_sub_id).unwrap();
let op = || async {
let res = socket_client
.request::<_, DocumentSubscribeResponse>(
"__absinthe__:control".into(),
"doc".into(),
&sub.payload,
)
.await;
match res {
Ok(subscription) => {
Ok(subscription.response.subscription_id)
}
Err(err) => {
tracing::debug!(?err, ?sub, "failed to resubscribe");
if cancellation_token.is_cancelled() {
Err(backoff::Error::permanent(err))
} else {
Err(backoff::Error::transient(err))
}
}
}
};
match backoff::future::retry(ExponentialBackoff::default(), op).await {
Ok(sub_id) => {
tracing::debug!(?sub, "resubscribed");
subscriptions.insert(sub_id, sub);
}
Err(err) => {
tracing::error!(?err, "fatal error trying to resubscribe to subscriptions (did the socket die?)");
subscriptions.insert(old_sub_id, sub);
break;
}
}
}
if !cancellation_token.is_cancelled() {
let mut messages = socket_client.filter_messages::<EventSubscription, _>(|msg| {
msg.event == "subscription:data" && msg.topic.starts_with("__absinthe__:doc")
});
loop {
select! {
_ = cancellation_token.cancelled() => {
tracing::trace!("received cancellation signal");
break;
}
sub = sub_ops_receiver.recv() => {
match sub {
Some(SubOp::AddSubscription(sub_id, sub)) => {
subscriptions.insert(sub_id, sub);
}
Some(SubOp::RemoveSubscription(sub_id)) => {
subscriptions.remove(&sub_id);
let socket_client = socket_client.clone();
task::spawn(async move {
let payload = json!({ "subscriptionId": sub_id });
if let Err(err) = socket_client.send_message(
"__absinthe__:control".into(),
"unsubscribe".into(),
payload
).await {
tracing::error!(?err, "failed to send unsubscribe request");
}
});
}
None => {
tracing::trace!("all senders were dropped");
cancellation_token.cancel();
break;
}
}
}
msg = messages.next() => {
match msg {
Some(EventSubscription{ result, subscription_id }) => {
if let Some(subscriber) = subscriptions.get(&subscription_id) {
match serde_json::from_value::<graphql_client::Response<Value>>(result) {
Ok(msg) => {
if let Err(err) = subscriber.sender.send(msg).await {
tracing::error!(?err, "failed to notify subscriber of event");
}
}
Err(err) => {
tracing::error!(?err, "invalid subscription message received");
}
}
}
}
None => {
tracing::trace!("all senders were dropped");
cancellation_token.cancel();
break;
}
}
}
}
}
}
(sub_ops_receiver, subscriptions)
})
.instrument(tracing::trace_span!("subscriptions"))
};
self.handle.replace(
async move {
incoming_messages_handle
.await
.map_err(anyhow::Error::from)?;
pinger_handle.await.map_err(anyhow::Error::from)?;
let outgoing_messages_receiver = outgoing_messages_handle
.await
.map_err(anyhow::Error::from)?;
let (sub_ops_receiver, subscriptions) =
subscriptions_handle.await.map_err(anyhow::Error::from)?;
Ok::<_, WebsocketConnectionError>((
outgoing_messages_receiver,
sub_ops_receiver,
subscriptions,
))
}
.boxed(),
);
tracing::debug!("connected to socket");
Ok(())
}
pub fn stay_conected(mut self) {
task::spawn(async move {
loop {
if let Err(err) = self.wait().await {
tracing::error!(?err, "irrecoverable connecton error");
break;
}
if self.cancellation_token.is_cancelled() {
break;
}
let mut backoff = ExponentialBackoff::default();
while let Err(err) = self.connect().await {
match backoff.next_backoff() {
Some(backoff_time) => {
tracing::error!(
?err,
"failed to reconnect, retrying in {:?}",
backoff_time
);
sleep(backoff_time).await;
}
None => {
tracing::error!(?err, "failed to reconnect, after many attempts");
return;
}
}
}
tracing::info!("successfully reconnected")
}
});
}
async fn wait(&mut self) -> Result<(), WebsocketConnectionError> {
let handle = self
.handle
.take()
.ok_or(WebsocketConnectionError::SocketClosed)?;
let (outgoing_messages_receiver, sub_ops_receiver, subscriptions) = handle.await?;
self.outgoing_messages.1.replace(outgoing_messages_receiver);
self.sub_ops.1.replace(sub_ops_receiver);
self.subscriptions.replace(subscriptions);
Ok(())
}
}
#[derive(Debug)]
struct SubscriptionRef {
payload: Value,
sender: mpsc::Sender<graphql_client::Response<Value>>,
}
#[derive(Debug, Clone)]
pub(super) struct SocketClient {
join_ref: Uuid,
outgoing_messages: mpsc::Sender<Message>,
incoming_messages: broadcast::Sender<ReceivePhoenixMessage<Value>>,
sub_ops: mpsc::UnboundedSender<SubOp>,
request_timeout: Duration,
cancellation_token: CancellationToken,
}
impl SocketClient {
pub async fn send_message<T>(
&self,
topic: String,
event: String,
payload: T,
) -> Result<Uuid, WebsocketConnectionError>
where
T: Serialize,
{
let msg_ref = Uuid::new_v4();
let msg = serde_json::to_string(&SendPhoenixMessage {
join_ref: self.join_ref,
msg_ref,
topic,
event,
payload,
})?;
self.outgoing_messages.send(msg.into()).await?;
Ok(msg_ref)
}
pub async fn request<T, U>(
&self,
topic: String,
event: String,
payload: T,
) -> Result<PhxReply<U>, WebsocketConnectionError>
where
T: Serialize,
U: DeserializeOwned,
{
let msg_ref = self.send_message(topic, event, payload).await?;
timeout(
self.request_timeout,
self.filter_messages::<PhxReply<U>, _>(move |msg| msg.msg_ref == Some(msg_ref))
.take(1)
.next(),
)
.await?
.ok_or(WebsocketConnectionError::SocketClosed)
}
pub async fn subscribe<T, U>(
&self,
payload: T,
) -> Result<Subscription<U>, WebsocketConnectionError>
where
T: Serialize,
U: DeserializeOwned,
{
let subscription: PhxReply<DocumentSubscribeResponse> = self
.request("__absinthe__:control".into(), "doc".into(), &payload)
.await?;
let payload = serde_json::to_value(&payload)?;
let (sender, receiver) = mpsc::channel(10);
let sub_id = subscription.response.subscription_id;
self.sub_ops
.send(SubOp::AddSubscription(
sub_id.clone(),
SubscriptionRef { payload, sender },
))
.map_err(anyhow::Error::from)?;
let this = self.clone();
Ok(Subscription::wrap(
ReceiverStream::new(receiver).filter_map(|res| serde_json::from_value(res.data?).ok()),
Some(move || {
if let Err(err) = this.sub_ops.send(SubOp::RemoveSubscription(sub_id)) {
tracing::error!(?err, "failed to notify unsubscribe");
}
}),
))
}
pub fn filter_messages<T, F>(&self, mut predicate: F) -> impl Stream<Item = T>
where
T: DeserializeOwned,
F: FnMut(&ReceivePhoenixMessage<Value>) -> bool,
{
BroadcastStream::new(self.incoming_messages.subscribe()).filter_map(move |msg| match msg {
Ok(msg) => {
if predicate(&msg) {
serde_json::from_value::<T>(msg.payload).ok()
} else {
None
}
}
Err(_) => None,
})
}
pub fn close(self) {
self.cancellation_token.cancel();
}
}