aboutsummaryrefslogtreecommitdiff
path: root/src/client.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/client.rs')
-rw-r--r--src/client.rs37
1 files changed, 23 insertions, 14 deletions
diff --git a/src/client.rs b/src/client.rs
index 8227e8f..bce7aca 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -37,10 +37,10 @@ pub(crate) struct ClientConn {
pub(crate) remote_addr: SocketAddr,
pub(crate) peer_id: NodeID,
- query_send: ArcSwapOption<mpsc::UnboundedSender<(RequestID, RequestPriority, Vec<u8>)>>,
+ query_send: ArcSwapOption<mpsc::UnboundedSender<(RequestID, RequestPriority, Data)>>,
next_query_number: AtomicU32,
- inflight: Mutex<HashMap<RequestID, oneshot::Sender<Vec<u8>>>>,
+ inflight: Mutex<HashMap<RequestID, oneshot::Sender<(Vec<u8>, AssociatedStream)>>>,
}
impl ClientConn {
@@ -148,9 +148,11 @@ impl ClientConn {
{
let query_send = self.query_send.load_full().ok_or(Error::ConnectionClosed)?;
+ // increment by 2; even are direct data; odd are associated stream
let id = self
.next_query_number
- .fetch_add(1, atomic::Ordering::Relaxed);
+ .fetch_add(2, atomic::Ordering::Relaxed);
+ let stream_id = id + 1;
cfg_if::cfg_if! {
if #[cfg(feature = "telemetry")] {
@@ -166,7 +168,7 @@ impl ClientConn {
};
// Encode request
- let body = rmp_to_vec_all_named(rq.borrow())?;
+ let (body, stream) = rmp_to_vec_all_named(rq.borrow())?;
drop(rq);
let request = QueryMessage {
@@ -185,7 +187,10 @@ impl ClientConn {
error!(
"Too many inflight requests! RequestID collision. Interrupting previous request."
);
- if old_ch.send(vec![]).is_err() {
+ if old_ch
+ .send((vec![], Box::pin(futures::stream::empty())))
+ .is_err()
+ {
debug!("Could not send empty response to collisionned request, probably because request was interrupted. Dropping response.");
}
}
@@ -195,15 +200,20 @@ impl ClientConn {
#[cfg(feature = "telemetry")]
span.set_attribute(KeyValue::new("len_query", bytes.len() as i64));
- query_send.send((id, prio, bytes))?;
+ query_send.send((id, prio, Data::Full(bytes)))?;
+ if let Some(stream) = stream {
+ query_send.send((stream_id, prio | PRIO_SECONDARY, Data::Streaming(stream)))?;
+ } else {
+ query_send.send((stream_id, prio, Data::Full(Vec::new())))?;
+ }
cfg_if::cfg_if! {
if #[cfg(feature = "telemetry")] {
- let resp = resp_recv
+ let (resp, stream) = resp_recv
.with_context(Context::current_with_span(span))
.await?;
} else {
- let resp = resp_recv.await?;
+ let (resp, stream) = resp_recv.await?;
}
}
@@ -217,10 +227,9 @@ impl ClientConn {
let code = resp[0];
if code == 0 {
- Ok(rmp_serde::decode::from_read_ref::<
- _,
- <T as Message>::Response,
- >(&resp[1..])?)
+ let mut deser = rmp_serde::decode::Deserializer::from_read_ref(&resp[1..]);
+ let res = T::Response::deserialize_msg(&mut deser, stream).await?;
+ Ok(res)
} else {
let msg = String::from_utf8(resp[1..].to_vec()).unwrap_or_default();
Err(Error::Remote(code, msg))
@@ -232,12 +241,12 @@ impl SendLoop for ClientConn {}
#[async_trait]
impl RecvLoop for ClientConn {
- fn recv_handler(self: &Arc<Self>, id: RequestID, msg: Vec<u8>) {
+ fn recv_handler(self: &Arc<Self>, id: RequestID, msg: Vec<u8>, stream: AssociatedStream) {
trace!("ClientConn recv_handler {} ({} bytes)", id, msg.len());
let mut inflight = self.inflight.lock().unwrap();
if let Some(ch) = inflight.remove(&id) {
- if ch.send(msg).is_err() {
+ if ch.send((msg, stream)).is_err() {
debug!("Could not send request response, probably because request was interrupted. Dropping response.");
}
}