Add websocket subprotocol for actix-web.
This commit is contained in:
parent
e25bcf1c98
commit
4b00f9393a
|
@ -1,19 +1,26 @@
|
||||||
|
use std::str::FromStr;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use actix::{
|
use actix::{
|
||||||
Actor, ActorContext, ActorFuture, ActorStream, AsyncContext, ContextFutureSpawner,
|
Actor, ActorContext, ActorFuture, ActorStream, AsyncContext, ContextFutureSpawner,
|
||||||
StreamHandler, WrapFuture, WrapStream,
|
StreamHandler, WrapFuture, WrapStream,
|
||||||
};
|
};
|
||||||
use actix_http::ws;
|
use actix_http::error::PayloadError;
|
||||||
|
use actix_http::{ws, Error};
|
||||||
|
use actix_web::web::Bytes;
|
||||||
|
use actix_web::{HttpRequest, HttpResponse};
|
||||||
use actix_web_actors::ws::{Message, ProtocolError, WebsocketContext};
|
use actix_web_actors::ws::{Message, ProtocolError, WebsocketContext};
|
||||||
use async_graphql::http::WebSocket;
|
use async_graphql::http::{WebSocket, WebSocketProtocols};
|
||||||
use async_graphql::{Data, ObjectType, Result, Schema, SubscriptionType};
|
use async_graphql::{Data, ObjectType, Result, Schema, SubscriptionType};
|
||||||
use std::time::{Duration, Instant};
|
use futures_util::stream::Stream;
|
||||||
|
|
||||||
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
|
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
|
||||||
const CLIENT_TIMEOUT: Duration = Duration::from_secs(10);
|
const CLIENT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||||
|
|
||||||
/// Actor for subscription via websocket
|
/// Actor for subscription via websocket
|
||||||
pub struct WSSubscription<Query, Mutation, Subscription> {
|
pub struct WSSubscription<Query, Mutation, Subscription> {
|
||||||
schema: Option<Schema<Query, Mutation, Subscription>>,
|
schema: Schema<Query, Mutation, Subscription>,
|
||||||
|
protocol: WebSocketProtocols,
|
||||||
last_heartbeat: Instant,
|
last_heartbeat: Instant,
|
||||||
messages: Option<async_channel::Sender<Vec<u8>>>,
|
messages: Option<async_channel::Sender<Vec<u8>>>,
|
||||||
initializer: Option<Box<dyn FnOnce(serde_json::Value) -> Result<Data> + Send + Sync>>,
|
initializer: Option<Box<dyn FnOnce(serde_json::Value) -> Result<Data> + Send + Sync>>,
|
||||||
|
@ -26,15 +33,41 @@ where
|
||||||
Mutation: ObjectType + Send + Sync + 'static,
|
Mutation: ObjectType + Send + Sync + 'static,
|
||||||
Subscription: SubscriptionType + Send + Sync + 'static,
|
Subscription: SubscriptionType + Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
/// Create an actor for subscription connection via websocket.
|
/// Start an actor for subscription connection via websocket.
|
||||||
pub fn new(schema: Schema<Query, Mutation, Subscription>) -> Self {
|
pub fn start<T>(
|
||||||
Self {
|
schema: Schema<Query, Mutation, Subscription>,
|
||||||
schema: Some(schema),
|
request: &HttpRequest,
|
||||||
last_heartbeat: Instant::now(),
|
stream: T,
|
||||||
messages: None,
|
) -> Result<HttpResponse, Error>
|
||||||
initializer: None,
|
where
|
||||||
continuation: Vec::new(),
|
T: Stream<Item = Result<Bytes, PayloadError>> + 'static,
|
||||||
}
|
{
|
||||||
|
let protocol = match request
|
||||||
|
.headers()
|
||||||
|
.get("sec-websocket-protocol")
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.and_then(|value| WebSocketProtocols::from_str(value).ok())
|
||||||
|
{
|
||||||
|
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,
|
||||||
|
initializer: None,
|
||||||
|
continuation: Vec::new(),
|
||||||
|
},
|
||||||
|
&["graphql-transport-ws", "graphql-ws"],
|
||||||
|
request,
|
||||||
|
stream,
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set a context data initialization function.
|
/// Set a context data initialization function.
|
||||||
|
@ -71,13 +104,18 @@ where
|
||||||
|
|
||||||
let (tx, rx) = async_channel::unbounded();
|
let (tx, rx) = async_channel::unbounded();
|
||||||
|
|
||||||
WebSocket::with_data(self.schema.take().unwrap(), rx, self.initializer.take())
|
WebSocket::with_data(
|
||||||
.into_actor(self)
|
self.schema.clone(),
|
||||||
.map(|response, _act, ctx| {
|
rx,
|
||||||
ctx.text(response);
|
self.initializer.take(),
|
||||||
})
|
self.protocol,
|
||||||
.finish()
|
)
|
||||||
.spawn(ctx);
|
.into_actor(self)
|
||||||
|
.map(|response, _act, ctx| {
|
||||||
|
ctx.text(response);
|
||||||
|
})
|
||||||
|
.finish()
|
||||||
|
.spawn(ctx);
|
||||||
|
|
||||||
self.messages = Some(tx);
|
self.messages = Some(tx);
|
||||||
}
|
}
|
||||||
|
|
|
@ -1,3 +1,4 @@
|
||||||
|
use async_graphql::http::WebSocketProtocols;
|
||||||
use async_graphql::*;
|
use async_graphql::*;
|
||||||
use futures_util::stream::{Stream, StreamExt};
|
use futures_util::stream::{Stream, StreamExt};
|
||||||
|
|
||||||
|
@ -23,7 +24,7 @@ pub async fn test_subscription_ws_transport() {
|
||||||
|
|
||||||
let schema = Schema::new(QueryRoot, EmptyMutation, SubscriptionRoot);
|
let schema = Schema::new(QueryRoot, EmptyMutation, SubscriptionRoot);
|
||||||
let (tx, rx) = async_channel::unbounded();
|
let (tx, rx) = async_channel::unbounded();
|
||||||
let mut stream = http::WebSocket::new(schema, rx);
|
let mut stream = http::WebSocket::new(schema, rx, WebSocketProtocols::SubscriptionsTransportWS);
|
||||||
|
|
||||||
tx.send(
|
tx.send(
|
||||||
serde_json::to_string(&value!({
|
serde_json::to_string(&value!({
|
||||||
|
@ -115,6 +116,7 @@ pub async fn test_subscription_ws_transport_with_token() {
|
||||||
data.insert(Token(payload.token));
|
data.insert(Token(payload.token));
|
||||||
Ok(data)
|
Ok(data)
|
||||||
}),
|
}),
|
||||||
|
WebSocketProtocols::SubscriptionsTransportWS,
|
||||||
);
|
);
|
||||||
|
|
||||||
tx.send(
|
tx.send(
|
||||||
|
@ -204,7 +206,7 @@ pub async fn test_subscription_ws_transport_error() {
|
||||||
|
|
||||||
let schema = Schema::new(QueryRoot, EmptyMutation, SubscriptionRoot);
|
let schema = Schema::new(QueryRoot, EmptyMutation, SubscriptionRoot);
|
||||||
let (tx, rx) = async_channel::unbounded();
|
let (tx, rx) = async_channel::unbounded();
|
||||||
let mut stream = http::WebSocket::new(schema, rx);
|
let mut stream = http::WebSocket::new(schema, rx, WebSocketProtocols::SubscriptionsTransportWS);
|
||||||
|
|
||||||
tx.send(
|
tx.send(
|
||||||
serde_json::to_string(&value!({
|
serde_json::to_string(&value!({
|
||||||
|
@ -276,7 +278,7 @@ pub async fn test_query_over_websocket() {
|
||||||
|
|
||||||
let schema = Schema::new(QueryRoot, EmptyMutation, EmptySubscription);
|
let schema = Schema::new(QueryRoot, EmptyMutation, EmptySubscription);
|
||||||
let (tx, rx) = async_channel::unbounded();
|
let (tx, rx) = async_channel::unbounded();
|
||||||
let mut stream = http::WebSocket::new(schema, rx);
|
let mut stream = http::WebSocket::new(schema, rx, WebSocketProtocols::SubscriptionsTransportWS);
|
||||||
|
|
||||||
tx.send(
|
tx.send(
|
||||||
serde_json::to_string(&value!({
|
serde_json::to_string(&value!({
|
||||||
|
|
Loading…
Reference in New Issue
Block a user