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

73
src/subscription/ws_transport.rs
View File

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

12
tests/subscription_websocket.rs
View File

@ -282,10 +282,18 @@ pub async fn test_query_over_websocket() {
assert_eq!(
Some(serde_json::json!({
"type": "complete",
"type": "data",
"id": "1",
"payload": { "data": { "value": 999 } },
})),
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()
);
}