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

Fix send data with 'complete' WebSocket message. #238

Browse Source
This commit is contained in:
Sunli 2020-08-14 15:21:10 +08:00
parent 702869a1f2
commit fe336dd4d0
3 changed files with 60 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
src/subscription/connection.rs
View File

@ -4,6 +4,7 @@ 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 std::collections::VecDeque;
use std::future::Future; use std::future::Future;
use std::pin::Pin; use std::pin::Pin;
@ -44,7 +45,7 @@ pub trait SubscriptionTransport: Send + Sync + Unpin + 'static {
schema: &Schema<Query, Mutation, Subscription>, schema: &Schema<Query, Mutation, Subscription>,
streams: &mut SubscriptionStreams, streams: &mut SubscriptionStreams,
data: Bytes, data: Bytes,
) -> std::result::Result<Option<Bytes>, Self::Error> ) -> std::result::Result<Option<Vec<Bytes>>, Self::Error>
where where
Query: ObjectType + Sync + Send + 'static, Query: ObjectType + Sync + Send + 'static,
Mutation: ObjectType + Sync + Send + 'static, Mutation: ObjectType + Sync + Send + 'static,
@ -78,6 +79,7 @@ where
rx_bytes, rx_bytes,
handle_request_fut: None, handle_request_fut: None,
waker: AtomicWaker::new(), waker: AtomicWaker::new(),
send_buf: VecDeque::new(),
}; };
while let Some(data) = inner_stream.next().await { while let Some(data) = inner_stream.next().await {
yield data; yield data;
@ -90,7 +92,7 @@ type HandleRequestBoxFut<'a, T> = Pin<
Box< Box<
dyn Future< dyn Future<
Output = ( Output = (
std::result::Result<Option<Bytes>, <T as SubscriptionTransport>::Error>, std::result::Result<Option<Vec<Bytes>>, <T as SubscriptionTransport>::Error>,
&'a mut T, &'a mut T,
&'a mut SubscriptionStreams, &'a mut SubscriptionStreams,
), ),
@ -108,6 +110,7 @@ struct SubscriptionStream<'a, Query, Mutation, Subscription, T: SubscriptionTran
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>,
} }
impl<'a, Query, Mutation, Subscription, T> Stream impl<'a, Query, Mutation, Subscription, T> Stream
@ -125,14 +128,21 @@ where
loop { loop {
// receive bytes // receive bytes
if let Some(bytes) = this.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(bytes), transport, streams)) => {
this.transport = Some(transport); this.transport = Some(transport);
this.streams = Some(streams); this.streams = Some(streams);
this.handle_request_fut = None; this.handle_request_fut = None;
if let Some(bytes) = bytes { if let Some(mut msgs) = bytes {
return Poll::Ready(Some(bytes)); if !msgs.is_empty() {
this.send_buf.extend(msgs.drain(1..));
return Poll::Ready(Some(msgs.remove(0)));
}
} }
continue; continue;
} }

45
src/subscription/ws_transport.rs
View File

@ -50,7 +50,7 @@ impl SubscriptionTransport for WebSocketTransport {
schema: &Schema<Query, Mutation, Subscription>, schema: &Schema<Query, Mutation, Subscription>,
streams: &mut SubscriptionStreams, streams: &mut SubscriptionStreams,
data: Bytes, data: Bytes,
) -> std::result::Result<Option<Bytes>, Self::Error> ) -> std::result::Result<Option<Vec<Bytes>>, Self::Error>
where where
Query: ObjectType + Sync + Send + 'static, Query: ObjectType + Sync + Send + 'static,
Mutation: ObjectType + Sync + Send + 'static, Mutation: ObjectType + Sync + Send + 'static,
@ -64,15 +64,13 @@ impl SubscriptionTransport for WebSocketTransport {
self.data = Arc::new(init_context_data(payload)?); self.data = Arc::new(init_context_data(payload)?);
} }
} }
Ok(Some( Ok(Some(vec![serde_json::to_vec(&OperationMessage {
serde_json::to_vec(&OperationMessage {
ty: "connection_ack".to_string(), ty: "connection_ack".to_string(),
id: None, id: None,
payload: None, payload: None,
}) })
.unwrap() .unwrap()
.into(), .into()]))
))
} }
"start" => { "start" => {
if let (Some(id), Some(payload)) = (msg.id, msg.payload) { if let (Some(id), Some(payload)) = (msg.id, msg.payload) {
@ -108,10 +106,10 @@ impl SubscriptionTransport for WebSocketTransport {
} }
match builder.execute(schema).await { match builder.execute(schema).await {
Ok(resp) => Ok(Some( Ok(resp) => Ok(Some(vec![
serde_json::to_vec(&OperationMessage { serde_json::to_vec(&OperationMessage {
ty: "complete".to_string(), ty: "data".to_string(),
id: Some(id), id: Some(id.clone()),
payload: Some( payload: Some(
serde_json::to_value(&GQLResponse(Ok(resp))) serde_json::to_value(&GQLResponse(Ok(resp)))
.unwrap(), .unwrap(),
@ -119,9 +117,16 @@ impl SubscriptionTransport for WebSocketTransport {
}) })
.unwrap() .unwrap()
.into(), .into(),
)),
Err(err) => Ok(Some(
serde_json::to_vec(&OperationMessage { serde_json::to_vec(&OperationMessage {
ty: "complete".to_string(),
id: Some(id),
payload: None,
})
.unwrap()
.into(),
])),
Err(err) => {
Ok(Some(vec![serde_json::to_vec(&OperationMessage {
ty: "error".to_string(), ty: "error".to_string(),
id: Some(id), id: Some(id),
payload: Some( payload: Some(
@ -129,21 +134,17 @@ impl SubscriptionTransport for WebSocketTransport {
), ),
}) })
.unwrap() .unwrap()
.into(), .into()]))
)),
} }
} }
Err(err) => Ok(Some( }
serde_json::to_vec(&OperationMessage { Err(err) => Ok(Some(vec![serde_json::to_vec(&OperationMessage {
ty: "error".to_string(), ty: "error".to_string(),
id: Some(id), id: Some(id),
payload: Some( payload: Some(serde_json::to_value(GQLError(&err)).unwrap()),
serde_json::to_value(GQLError(&err)).unwrap(),
),
}) })
.unwrap() .unwrap()
.into(), .into()])),
)),
} }
} else { } else {
Ok(None) Ok(None)
@ -157,15 +158,13 @@ impl SubscriptionTransport for WebSocketTransport {
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( return Ok(Some(vec![serde_json::to_vec(&OperationMessage {
serde_json::to_vec(&OperationMessage {
ty: "complete".to_string(), ty: "complete".to_string(),
id: Some(id), id: Some(id),
payload: None, payload: None,
}) })
.unwrap() .unwrap()
.into(), .into()]));
));
} }
} }
Ok(None) Ok(None)

12
tests/subscription_websocket.rs
View File

@ -282,10 +282,18 @@ pub async fn test_query_over_websocket() {
assert_eq!( assert_eq!(
Some(serde_json::json!({ Some(serde_json::json!({
"type": "complete", "type": "data",
"id": "1", "id": "1",
"payload": { "data": { "value": 999 } }, "payload": { "data": { "value": 999 } },
})), })),
serde_json::from_slice(&stream.next().await.unwrap()).unwrap() serde_json::from_slice(&stream.next().await.unwrap()).unwrap()
) );
assert_eq!(
Some(serde_json::json!({
"type": "complete",
"id": "1",
})),
serde_json::from_slice(&stream.next().await.unwrap()).unwrap()
);
} }