Fix send data with 'complete' WebSocket message. #238
This commit is contained in:
parent
702869a1f2
commit
fe336dd4d0
|
@ -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;
|
||||
}
|
||||
|
|
|
@ -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 {
|
||||
Ok(Some(vec![serde_json::to_vec(&OperationMessage {
|
||||
ty: "connection_ack".to_string(),
|
||||
id: None,
|
||||
payload: None,
|
||||
})
|
||||
.unwrap()
|
||||
.into(),
|
||||
))
|
||||
.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 {
|
||||
}
|
||||
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(),
|
||||
),
|
||||
payload: Some(serde_json::to_value(GQLError(&err)).unwrap()),
|
||||
})
|
||||
.unwrap()
|
||||
.into(),
|
||||
)),
|
||||
.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 {
|
||||
return Ok(Some(vec![serde_json::to_vec(&OperationMessage {
|
||||
ty: "complete".to_string(),
|
||||
id: Some(id),
|
||||
payload: None,
|
||||
})
|
||||
.unwrap()
|
||||
.into(),
|
||||
));
|
||||
.into()]));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
|
|
|
@ -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()
|
||||
);
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue
Block a user