anna/async-graphql
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Improve subscriptions performance.

Browse Source
This commit is contained in:
Sunli 2020-09-02 10:47:18 +08:00
parent 5dcef3e938
commit d2bb44798a
3 changed files with 97 additions and 85 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.toml
View File

@ -50,7 +50,7 @@ bson = { version = "1.0.0", optional = true }
uuid = { version = "0.8.1", features = ["v4", "serde"] } uuid = { version = "0.8.1", features = ["v4", "serde"] }
url = { version = "2.1.1", optional = true } url = { version = "2.1.1", optional = true }
chrono-tz = { version = "0.5.1", optional = true } chrono-tz = { version = "0.5.1", optional = true }
smallvec = "1.4.2"
[dev-dependencies] [dev-dependencies]
async-std = { version = "1.5.0", features = ["attributes"] } async-std = { version = "1.5.0", features = ["attributes"] }

41
src/subscription/connection.rs
View File

@ -4,7 +4,6 @@ use futures::channel::mpsc;
use futures::task::{AtomicWaker, Context, Poll}; use futures::task::{AtomicWaker, Context, Poll};
use futures::{Stream, StreamExt}; use futures::{Stream, StreamExt};
use slab::Slab; use slab::Slab;
use smallvec::SmallVec;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::future::Future; use std::future::Future;
use std::pin::Pin; use std::pin::Pin;
@ -45,8 +44,9 @@ pub trait ConnectionTransport: Send + Sync + Unpin + 'static {
&mut self, &mut self,
schema: &Schema<Query, Mutation, Subscription>, schema: &Schema<Query, Mutation, Subscription>,
streams: &mut SubscriptionStreams, streams: &mut SubscriptionStreams,
data: Bytes, request: Bytes,
) -> std::result::Result<Option<SmallVec<[Bytes; 4]>>, Self::Error> send_buf: &mut VecDeque<Bytes>,
) -> std::result::Result<(), Self::Error>
where where
Query: ObjectType + Sync + Send + 'static, Query: ObjectType + Sync + Send + 'static,
Mutation: ObjectType + Sync + Send + 'static, Mutation: ObjectType + Sync + Send + 'static,
@ -73,6 +73,7 @@ where
let mut streams = SubscriptionStreams { let mut streams = SubscriptionStreams {
streams: Default::default(), streams: Default::default(),
}; };
let mut send_buf = Default::default();
let mut inner_stream = SubscriptionStream { let mut inner_stream = SubscriptionStream {
schema: &schema, schema: &schema,
transport: Some(&mut transport), transport: Some(&mut transport),
@ -80,7 +81,7 @@ where
rx_bytes, rx_bytes,
handle_request_fut: None, handle_request_fut: None,
waker: AtomicWaker::new(), waker: AtomicWaker::new(),
send_buf: VecDeque::new(), send_buf: Some(&mut send_buf),
}; };
while let Some(data) = inner_stream.next().await { while let Some(data) = inner_stream.next().await {
yield data; yield data;
@ -93,12 +94,10 @@ type HandleRequestBoxFut<'a, T> = Pin<
Box< Box<
dyn Future< dyn Future<
Output = ( Output = (
std::result::Result< std::result::Result<(), <T as ConnectionTransport>::Error>,
Option<SmallVec<[Bytes; 4]>>,
<T as ConnectionTransport>::Error,
>,
&'a mut T, &'a mut T,
&'a mut SubscriptionStreams, &'a mut SubscriptionStreams,
&'a mut VecDeque<Bytes>,
), ),
> + Send > + Send
+ 'a, + 'a,
@ -114,7 +113,7 @@ struct SubscriptionStream<'a, Query, Mutation, Subscription, T: ConnectionTransp
rx_bytes: mpsc::UnboundedReceiver<Bytes>, rx_bytes: mpsc::UnboundedReceiver<Bytes>,
handle_request_fut: Option<HandleRequestBoxFut<'a, T>>, handle_request_fut: Option<HandleRequestBoxFut<'a, T>>,
waker: AtomicWaker, waker: AtomicWaker,
send_buf: VecDeque<Bytes>, send_buf: Option<&'a mut VecDeque<Bytes>>,
} }
impl<'a, Query, Mutation, Subscription, T> Stream impl<'a, Query, Mutation, Subscription, T> Stream
@ -132,25 +131,22 @@ where
loop { loop {
// receive bytes // receive bytes
if let Some(bytes) = this.send_buf.pop_front() { if let Some(send_buf) = &mut this.send_buf {
return Poll::Ready(Some(bytes)); if let Some(bytes) = send_buf.pop_front() {
return Poll::Ready(Some(bytes));
}
} }
if let Some(handle_request_fut) = &mut this.handle_request_fut { if let Some(handle_request_fut) = &mut this.handle_request_fut {
match handle_request_fut.as_mut().poll(cx) { match handle_request_fut.as_mut().poll(cx) {
Poll::Ready((Ok(bytes), transport, streams)) => { Poll::Ready((Ok(()), transport, streams, send_buf)) => {
this.transport = Some(transport); this.transport = Some(transport);
this.streams = Some(streams); this.streams = Some(streams);
this.send_buf = Some(send_buf);
this.handle_request_fut = None; this.handle_request_fut = None;
if let Some(mut msgs) = bytes {
if !msgs.is_empty() {
this.send_buf.extend(msgs.drain(1..));
return Poll::Ready(Some(msgs.remove(0)));
}
}
continue; continue;
} }
Poll::Ready((Err(_), _, _)) => return Poll::Ready(None), Poll::Ready((Err(_), _, _, _)) => return Poll::Ready(None),
Poll::Pending => {} Poll::Pending => {}
} }
} else { } else {
@ -159,9 +155,12 @@ where
let transport = this.transport.take().unwrap(); let transport = this.transport.take().unwrap();
let schema = this.schema; let schema = this.schema;
let streams = this.streams.take().unwrap(); let streams = this.streams.take().unwrap();
let send_buf = this.send_buf.take().unwrap();
this.handle_request_fut = Some(Box::pin(async move { this.handle_request_fut = Some(Box::pin(async move {
let res = transport.handle_request(schema, streams, data).await; let res = transport
(res, transport, streams) .handle_request(schema, streams, data, send_buf)
.await;
(res, transport, streams, send_buf)
})); }));
this.waker.wake(); this.waker.wake();
continue; continue;

139
src/subscription/ws_transport.rs
View File

@ -6,8 +6,7 @@ use crate::{
}; };
use bytes::Bytes; use bytes::Bytes;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use smallvec::{smallvec, SmallVec}; use std::collections::{HashMap, VecDeque};
use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
#[derive(Serialize, Deserialize)] #[derive(Serialize, Deserialize)]
@ -43,6 +42,12 @@ impl WebSocketTransport {
} }
} }
fn send_message<T: Serialize>(send_buf: &mut VecDeque<Bytes>, msg: &T) {
if let Ok(data) = serde_json::to_vec(msg) {
send_buf.push_back(data.into());
}
}
#[async_trait::async_trait] #[async_trait::async_trait]
impl ConnectionTransport for WebSocketTransport { impl ConnectionTransport for WebSocketTransport {
type Error = FieldError; type Error = FieldError;
@ -51,14 +56,15 @@ impl ConnectionTransport for WebSocketTransport {
&mut self, &mut self,
schema: &Schema<Query, Mutation, Subscription>, schema: &Schema<Query, Mutation, Subscription>,
streams: &mut SubscriptionStreams, streams: &mut SubscriptionStreams,
data: Bytes, request: Bytes,
) -> std::result::Result<Option<SmallVec<[Bytes; 4]>>, Self::Error> send_buf: &mut VecDeque<Bytes>,
) -> std::result::Result<(), Self::Error>
where where
Query: ObjectType + Sync + Send + 'static, Query: ObjectType + Sync + Send + 'static,
Mutation: ObjectType + Sync + Send + 'static, Mutation: ObjectType + Sync + Send + 'static,
Subscription: SubscriptionType + Sync + Send + 'static, Subscription: SubscriptionType + Sync + Send + 'static,
{ {
match serde_json::from_slice::<OperationMessage>(&data) { match serde_json::from_slice::<OperationMessage>(&request) {
Ok(msg) => match msg.ty.as_str() { Ok(msg) => match msg.ty.as_str() {
"connection_init" => { "connection_init" => {
if let Some(payload) = msg.payload { if let Some(payload) = msg.payload {
@ -66,13 +72,15 @@ impl ConnectionTransport for WebSocketTransport {
self.data = Arc::new(init_context_data(payload)?); self.data = Arc::new(init_context_data(payload)?);
} }
} }
Ok(Some(smallvec![serde_json::to_vec(&OperationMessage { send_message(
ty: "connection_ack".to_string(), send_buf,
id: None, &OperationMessage {
payload: None, ty: "connection_ack".to_string(),
}) id: None,
.unwrap() payload: None,
.into()])) },
);
Ok(())
} }
"start" => { "start" => {
if let (Some(id), Some(payload)) = (msg.id, msg.payload) { if let (Some(id), Some(payload)) = (msg.id, msg.payload) {
@ -95,7 +103,6 @@ impl ConnectionTransport for WebSocketTransport {
let stream_id = streams.add(stream); let stream_id = streams.add(stream);
self.id_to_sid.insert(id.clone(), stream_id); self.id_to_sid.insert(id.clone(), stream_id);
self.sid_to_id.insert(stream_id, id); self.sid_to_id.insert(stream_id, id);
Ok(None)
} }
Err(Error::Query { err, .. }) Err(Error::Query { err, .. })
if err == QueryError::NotSupported => if err == QueryError::NotSupported =>
@ -108,72 +115,78 @@ impl ConnectionTransport for WebSocketTransport {
} }
match builder.execute(schema).await { match builder.execute(schema).await {
Ok(resp) => Ok(Some(smallvec![ Ok(resp) => {
serde_json::to_vec(&OperationMessage { send_message(
ty: "data".to_string(), send_buf,
id: Some(id.clone()), &OperationMessage {
payload: Some( ty: "data".to_string(),
serde_json::to_value(&GQLResponse(Ok(resp))) id: Some(id.clone()),
payload: Some(
serde_json::to_value(&GQLResponse(Ok(
resp,
)))
.unwrap(), .unwrap(),
), ),
}) },
.unwrap() );
.into(),
serde_json::to_vec(&OperationMessage { send_message(
ty: "complete".to_string(), send_buf,
id: Some(id), &OperationMessage {
payload: None, ty: "complete".to_string(),
}) id: Some(id),
.unwrap() payload: None,
.into(), },
])), );
Err(err) => Ok(Some(smallvec![serde_json::to_vec( }
&OperationMessage { Err(err) => {
ty: "error".to_string(), send_message(
id: Some(id), send_buf,
payload: Some( &OperationMessage {
serde_json::to_value(GQLError(&err)).unwrap(), ty: "error".to_string(),
), id: Some(id),
} payload: Some(
) serde_json::to_value(GQLError(&err))
.unwrap() .unwrap(),
.into()])), ),
},
);
}
} }
} }
Err(err) => { Err(err) => {
Ok(Some(smallvec![serde_json::to_vec(&OperationMessage { send_message(
ty: "error".to_string(), send_buf,
id: Some(id), &OperationMessage {
payload: Some( ty: "error".to_string(),
serde_json::to_value(GQLError(&err)).unwrap() id: Some(id),
), payload: Some(
}) serde_json::to_value(GQLError(&err)).unwrap(),
.unwrap() ),
.into()])) },
);
} }
} }
} else {
Ok(None)
} }
} else {
Ok(None)
} }
Ok(())
} }
"stop" => { "stop" => {
if let Some(id) = msg.id { if let Some(id) = msg.id {
if let Some(sid) = self.id_to_sid.remove(&id) { if let Some(sid) = self.id_to_sid.remove(&id) {
self.sid_to_id.remove(&sid); self.sid_to_id.remove(&sid);
streams.remove(sid); streams.remove(sid);
return Ok(Some(smallvec![serde_json::to_vec(&OperationMessage { send_message(
ty: "complete".to_string(), send_buf,
id: Some(id), &OperationMessage {
payload: None, ty: "complete".to_string(),
}) id: Some(id),
.unwrap() payload: None,
.into()])); },
);
} }
} }
Ok(None) Ok(())
} }
"connection_terminate" => Err("connection_terminate".into()), "connection_terminate" => Err("connection_terminate".into()),
_ => Err("Unknown op".into()), _ => Err("Unknown op".into()),