2021-01-16 04:47:32 +00:00
|
|
|
use std::future::Future;
|
2020-12-04 04:13:52 +00:00
|
|
|
use std::str::FromStr;
|
|
|
|
use std::time::{Duration, Instant};
|
|
|
|
|
2020-03-17 09:26:59 +00:00
|
|
|
use actix::{
|
2020-09-17 18:22:54 +00:00
|
|
|
Actor, ActorContext, ActorFuture, ActorStream, AsyncContext, ContextFutureSpawner,
|
|
|
|
StreamHandler, WrapFuture, WrapStream,
|
2020-03-17 09:26:59 +00:00
|
|
|
};
|
2020-12-04 04:13:52 +00:00
|
|
|
use actix_http::error::PayloadError;
|
|
|
|
use actix_http::{ws, Error};
|
|
|
|
use actix_web::web::Bytes;
|
|
|
|
use actix_web::{HttpRequest, HttpResponse};
|
2021-02-06 03:02:41 +00:00
|
|
|
use actix_web_actors::ws::{CloseReason, Message, ProtocolError, WebsocketContext};
|
|
|
|
use async_graphql::http::{WebSocket, WebSocketProtocols, WsMessage};
|
2020-10-01 01:07:19 +00:00
|
|
|
use async_graphql::{Data, ObjectType, Result, Schema, SubscriptionType};
|
2021-01-16 04:47:32 +00:00
|
|
|
use futures_util::future::Ready;
|
2020-12-04 04:13:52 +00:00
|
|
|
use futures_util::stream::Stream;
|
2020-03-17 09:26:59 +00:00
|
|
|
|
2020-04-07 06:30:46 +00:00
|
|
|
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
|
|
|
|
const CLIENT_TIMEOUT: Duration = Duration::from_secs(10);
|
|
|
|
|
2020-04-14 01:53:17 +00:00
|
|
|
/// Actor for subscription via websocket
|
2021-01-16 04:47:32 +00:00
|
|
|
pub struct WSSubscription<Query, Mutation, Subscription, F> {
|
2020-12-04 04:13:52 +00:00
|
|
|
schema: Schema<Query, Mutation, Subscription>,
|
|
|
|
protocol: WebSocketProtocols,
|
2020-09-17 18:22:54 +00:00
|
|
|
last_heartbeat: Instant,
|
2020-10-15 09:33:38 +00:00
|
|
|
messages: Option<async_channel::Sender<Vec<u8>>>,
|
2021-01-16 04:47:32 +00:00
|
|
|
initializer: Option<F>,
|
2020-09-17 18:22:54 +00:00
|
|
|
continuation: Vec<u8>,
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
|
|
|
|
2021-01-16 04:47:32 +00:00
|
|
|
impl<Query, Mutation, Subscription>
|
|
|
|
WSSubscription<Query, Mutation, Subscription, fn(serde_json::Value) -> Ready<Result<Data>>>
|
2020-03-17 09:26:59 +00:00
|
|
|
where
|
2021-01-14 04:41:59 +00:00
|
|
|
Query: ObjectType + 'static,
|
|
|
|
Mutation: ObjectType + 'static,
|
|
|
|
Subscription: SubscriptionType + 'static,
|
2020-03-17 09:26:59 +00:00
|
|
|
{
|
2020-12-04 04:13:52 +00:00
|
|
|
/// Start an actor for subscription connection via websocket.
|
|
|
|
pub fn start<T>(
|
|
|
|
schema: Schema<Query, Mutation, Subscription>,
|
|
|
|
request: &HttpRequest,
|
|
|
|
stream: T,
|
|
|
|
) -> Result<HttpResponse, Error>
|
|
|
|
where
|
|
|
|
T: Stream<Item = Result<Bytes, PayloadError>> + 'static,
|
2020-12-04 04:35:35 +00:00
|
|
|
{
|
2021-01-16 04:47:32 +00:00
|
|
|
Self::start_with_initializer(schema, request, stream, |_| {
|
|
|
|
futures_util::future::ready(Ok(Default::default()))
|
|
|
|
})
|
2020-12-04 04:35:35 +00:00
|
|
|
}
|
2021-01-16 04:47:32 +00:00
|
|
|
}
|
2020-12-04 04:35:35 +00:00
|
|
|
|
2021-01-16 04:47:32 +00:00
|
|
|
impl<Query, Mutation, Subscription, F, R> WSSubscription<Query, Mutation, Subscription, F>
|
|
|
|
where
|
|
|
|
Query: ObjectType + 'static,
|
|
|
|
Mutation: ObjectType + 'static,
|
|
|
|
Subscription: SubscriptionType + 'static,
|
|
|
|
F: FnOnce(serde_json::Value) -> R + Unpin + Send + 'static,
|
|
|
|
R: Future<Output = Result<Data>> + Send + 'static,
|
|
|
|
{
|
2020-12-04 04:35:35 +00:00
|
|
|
/// Start an actor for subscription connection via websocket with an initialization function.
|
2021-01-16 04:47:32 +00:00
|
|
|
pub fn start_with_initializer<T>(
|
2020-12-04 04:35:35 +00:00
|
|
|
schema: Schema<Query, Mutation, Subscription>,
|
|
|
|
request: &HttpRequest,
|
|
|
|
stream: T,
|
|
|
|
initializer: F,
|
|
|
|
) -> Result<HttpResponse, Error>
|
|
|
|
where
|
|
|
|
T: Stream<Item = Result<Bytes, PayloadError>> + 'static,
|
2021-01-16 04:47:32 +00:00
|
|
|
F: FnOnce(serde_json::Value) -> R + Unpin + Send + 'static,
|
|
|
|
R: Future<Output = Result<Data>> + Send + 'static,
|
2020-12-04 04:13:52 +00:00
|
|
|
{
|
|
|
|
let protocol = match request
|
|
|
|
.headers()
|
|
|
|
.get("sec-websocket-protocol")
|
|
|
|
.and_then(|value| value.to_str().ok())
|
2020-12-04 17:16:14 +00:00
|
|
|
.and_then(|protocols| {
|
|
|
|
protocols
|
|
|
|
.split(',')
|
|
|
|
.find_map(|p| WebSocketProtocols::from_str(p.trim()).ok())
|
|
|
|
}) {
|
2020-12-04 04:13:52 +00:00
|
|
|
Some(protocol) => protocol,
|
|
|
|
None => {
|
|
|
|
// default to the prior standard
|
|
|
|
WebSocketProtocols::SubscriptionsTransportWS
|
|
|
|
}
|
|
|
|
};
|
|
|
|
|
|
|
|
actix_web_actors::ws::start_with_protocols(
|
|
|
|
Self {
|
|
|
|
schema,
|
|
|
|
protocol,
|
|
|
|
last_heartbeat: Instant::now(),
|
|
|
|
messages: None,
|
2021-01-16 04:47:32 +00:00
|
|
|
initializer: Some(initializer),
|
2020-12-04 04:13:52 +00:00
|
|
|
continuation: Vec::new(),
|
|
|
|
},
|
|
|
|
&["graphql-transport-ws", "graphql-ws"],
|
|
|
|
request,
|
|
|
|
stream,
|
|
|
|
)
|
2020-04-23 07:30:12 +00:00
|
|
|
}
|
|
|
|
|
2020-09-17 18:22:54 +00:00
|
|
|
fn send_heartbeats(&self, ctx: &mut WebsocketContext<Self>) {
|
2020-04-07 06:30:46 +00:00
|
|
|
ctx.run_interval(HEARTBEAT_INTERVAL, |act, ctx| {
|
2020-09-17 18:22:54 +00:00
|
|
|
if Instant::now().duration_since(act.last_heartbeat) > CLIENT_TIMEOUT {
|
2020-03-19 09:20:12 +00:00
|
|
|
ctx.stop();
|
|
|
|
}
|
2020-04-07 06:30:46 +00:00
|
|
|
ctx.ping(b"");
|
2020-03-19 09:20:12 +00:00
|
|
|
});
|
|
|
|
}
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
|
|
|
|
2021-01-16 04:47:32 +00:00
|
|
|
impl<Query, Mutation, Subscription, F, R> Actor for WSSubscription<Query, Mutation, Subscription, F>
|
2020-03-17 09:26:59 +00:00
|
|
|
where
|
2021-01-14 04:41:59 +00:00
|
|
|
Query: ObjectType + 'static,
|
|
|
|
Mutation: ObjectType + 'static,
|
|
|
|
Subscription: SubscriptionType + 'static,
|
2021-01-16 04:47:32 +00:00
|
|
|
F: FnOnce(serde_json::Value) -> R + Unpin + Send + 'static,
|
|
|
|
R: Future<Output = Result<Data>> + Send + 'static,
|
2020-03-17 09:26:59 +00:00
|
|
|
{
|
|
|
|
type Context = WebsocketContext<Self>;
|
|
|
|
|
|
|
|
fn started(&mut self, ctx: &mut Self::Context) {
|
2020-09-17 18:22:54 +00:00
|
|
|
self.send_heartbeats(ctx);
|
|
|
|
|
2020-10-15 09:33:38 +00:00
|
|
|
let (tx, rx) = async_channel::unbounded();
|
2020-09-17 18:22:54 +00:00
|
|
|
|
2020-12-04 04:13:52 +00:00
|
|
|
WebSocket::with_data(
|
|
|
|
self.schema.clone(),
|
|
|
|
rx,
|
2021-01-16 04:47:32 +00:00
|
|
|
self.initializer.take().unwrap(),
|
2020-12-04 04:13:52 +00:00
|
|
|
self.protocol,
|
|
|
|
)
|
|
|
|
.into_actor(self)
|
2021-02-06 03:02:41 +00:00
|
|
|
.map(|response, _act, ctx| match response {
|
|
|
|
WsMessage::Text(text) => ctx.text(text),
|
|
|
|
WsMessage::Close(code, msg) => ctx.close(Some(CloseReason {
|
|
|
|
code: code.into(),
|
|
|
|
description: Some(msg),
|
|
|
|
})),
|
2020-12-04 04:13:52 +00:00
|
|
|
})
|
|
|
|
.finish()
|
|
|
|
.spawn(ctx);
|
2020-09-17 18:22:54 +00:00
|
|
|
|
|
|
|
self.messages = Some(tx);
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-01-16 04:47:32 +00:00
|
|
|
impl<Query, Mutation, Subscription, F, R> StreamHandler<Result<Message, ProtocolError>>
|
|
|
|
for WSSubscription<Query, Mutation, Subscription, F>
|
2020-03-17 09:26:59 +00:00
|
|
|
where
|
2021-01-14 04:41:59 +00:00
|
|
|
Query: ObjectType + 'static,
|
|
|
|
Mutation: ObjectType + 'static,
|
|
|
|
Subscription: SubscriptionType + 'static,
|
2021-01-16 04:47:32 +00:00
|
|
|
F: FnOnce(serde_json::Value) -> R + Unpin + Send + 'static,
|
|
|
|
R: Future<Output = Result<Data>> + Send + 'static,
|
2020-03-17 09:26:59 +00:00
|
|
|
{
|
|
|
|
fn handle(&mut self, msg: Result<Message, ProtocolError>, ctx: &mut Self::Context) {
|
|
|
|
let msg = match msg {
|
|
|
|
Err(_) => {
|
|
|
|
ctx.stop();
|
|
|
|
return;
|
|
|
|
}
|
|
|
|
Ok(msg) => msg,
|
|
|
|
};
|
|
|
|
|
2020-09-17 18:22:54 +00:00
|
|
|
let message = match msg {
|
2020-03-17 09:26:59 +00:00
|
|
|
Message::Ping(msg) => {
|
2020-09-17 18:22:54 +00:00
|
|
|
self.last_heartbeat = Instant::now();
|
2020-03-17 09:26:59 +00:00
|
|
|
ctx.pong(&msg);
|
2020-09-17 18:22:54 +00:00
|
|
|
None
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
|
|
|
Message::Pong(_) => {
|
2020-09-17 18:22:54 +00:00
|
|
|
self.last_heartbeat = Instant::now();
|
|
|
|
None
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
2020-09-17 18:22:54 +00:00
|
|
|
Message::Continuation(item) => match item {
|
|
|
|
ws::Item::FirstText(bytes) | ws::Item::FirstBinary(bytes) => {
|
|
|
|
self.continuation = bytes.to_vec();
|
|
|
|
None
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
2020-09-17 18:22:54 +00:00
|
|
|
ws::Item::Continue(bytes) => {
|
|
|
|
self.continuation.extend_from_slice(&bytes);
|
|
|
|
None
|
|
|
|
}
|
|
|
|
ws::Item::Last(bytes) => {
|
|
|
|
self.continuation.extend_from_slice(&bytes);
|
|
|
|
Some(std::mem::take(&mut self.continuation))
|
|
|
|
}
|
|
|
|
},
|
|
|
|
Message::Text(s) => Some(s.into_bytes()),
|
|
|
|
Message::Binary(bytes) => Some(bytes.to_vec()),
|
|
|
|
Message::Close(_) => {
|
2020-03-17 09:26:59 +00:00
|
|
|
ctx.stop();
|
2020-09-17 18:22:54 +00:00
|
|
|
None
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
2020-09-17 18:22:54 +00:00
|
|
|
Message::Nop => None,
|
|
|
|
};
|
2020-03-17 09:26:59 +00:00
|
|
|
|
2020-09-17 18:22:54 +00:00
|
|
|
if let Some(message) = message {
|
2020-10-15 09:33:38 +00:00
|
|
|
let sender = self.messages.as_ref().unwrap().clone();
|
2020-09-17 18:22:54 +00:00
|
|
|
|
|
|
|
async move { sender.send(message).await }
|
|
|
|
.into_actor(self)
|
|
|
|
.map(|res, _actor, ctx| match res {
|
|
|
|
Ok(()) => {}
|
|
|
|
Err(_) => ctx.stop(),
|
|
|
|
})
|
|
|
|
.spawn(ctx)
|
|
|
|
}
|
2020-03-17 09:26:59 +00:00
|
|
|
}
|
|
|
|
}
|