netapp/src/client.rs

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::atomic::{self, AtomicU32};
use std::sync::{Arc, Mutex};

use arc_swap::ArcSwapOption;
use log::{debug, error, trace};

use tokio::net::TcpStream;
use tokio::select;
use tokio::sync::{mpsc, oneshot, watch};
use tokio_util::compat::*;

use futures::io::AsyncReadExt;

use async_trait::async_trait;

use kuska_handshake::async_std::{handshake_client, BoxStream};

use crate::endpoint::*;
use crate::error::*;
use crate::netapp::*;
use crate::proto::*;
use crate::util::*;

pub(crate) struct ClientConn {
	pub(crate) remote_addr: SocketAddr,
	pub(crate) peer_id: NodeID,

	query_send: ArcSwapOption<mpsc::UnboundedSender<(RequestID, RequestPriority, Vec<u8>)>>,

	next_query_number: AtomicU32,
	inflight: Mutex<HashMap<RequestID, oneshot::Sender<Vec<u8>>>>,
}

impl ClientConn {
	pub(crate) async fn init(
		netapp: Arc<NetApp>,
		socket: TcpStream,
		peer_id: NodeID,
	) -> Result<(), Error> {
		let remote_addr = socket.peer_addr()?;
		let mut socket = socket.compat();

		let handshake = handshake_client(
			&mut socket,
			netapp.netid.clone(),
			netapp.id,
			netapp.privkey.clone(),
			peer_id,
		)
		.await?;

		debug!(
			"Handshake complete (client) with {}@{}",
			hex::encode(&peer_id),
			remote_addr
		);

		let (read, write) = socket.split();

		let (read, write) =
			BoxStream::from_handshake(read, write, handshake, 0x8000).split_read_write();

		let (query_send, query_recv) = mpsc::unbounded_channel();

		let (stop_recv_loop, stop_recv_loop_recv) = watch::channel(false);

		let conn = Arc::new(ClientConn {
			remote_addr,
			peer_id,
			next_query_number: AtomicU32::from(RequestID::default()),
			query_send: ArcSwapOption::new(Some(Arc::new(query_send))),
			inflight: Mutex::new(HashMap::new()),
		});

		netapp.connected_as_client(peer_id, conn.clone());

		tokio::spawn(async move {
			let send_future = tokio::spawn(conn.clone().send_loop(query_recv, write));

			let conn2 = conn.clone();
			let recv_future = tokio::spawn(async move {
				select! {
					r = conn2.recv_loop(read) => r,
					_ = await_exit(stop_recv_loop_recv) => Ok(())
				}
			});

			send_future.await.log_err("ClientConn send_loop");

			// FIXME: should do here: wait for inflight requests to all have their response
			stop_recv_loop
				.send(true)
				.log_err("ClientConn send true to stop_recv_loop");

			recv_future.await.log_err("ClientConn recv_loop");

			// Make sure we don't wait on any more requests that won't
			// have a response
			conn.inflight.lock().unwrap().clear();

			netapp.disconnected_as_client(&peer_id, conn);
		});

		Ok(())
	}

	pub fn close(&self) {
		self.query_send.store(None);
	}

	pub(crate) async fn call<T>(
		self: Arc<Self>,
		rq: &T,
		path: &str,
		prio: RequestPriority,
	) -> Result<<T as Message>::Response, Error>
	where
		T: Message,
	{
		let query_send = self.query_send.load_full().ok_or(Error::ConnectionClosed)?;

		let id = self
			.next_query_number
			.fetch_add(1, atomic::Ordering::Relaxed);

		cfg_if::cfg_if! {
			if #[cfg(feature = "telemetry")] {
				use opentelemetry::{trace::{TraceId, TraceContextExt}, Context};
				let trace_id_int = Context::current()
					.span()
					.span_context()
					.trace_id();
				let trace_id = if trace_id_int == TraceId::INVALID {
					None
				} else {
					Some(trace_id_int.to_bytes().to_vec())
				};
			} else {
				let trace_id: Option<Vec<u8>> = None;
			}
		};

		// Encode request
		let mut bytes = vec![];

		bytes.extend_from_slice(&[prio, path.as_bytes().len() as u8]);
		bytes.extend_from_slice(path.as_bytes());

		if let Some(by) = trace_id {
			bytes.push(by.len() as u8);
			bytes.extend(by);
		} else {
			bytes.push(0);
		}

		bytes.extend_from_slice(&rmp_to_vec_all_named(rq)?[..]);

		// Send request through
		let (resp_send, resp_recv) = oneshot::channel();
		let old = self.inflight.lock().unwrap().insert(id, resp_send);
		if let Some(old_ch) = old {
			error!(
				"Too many inflight requests! RequestID collision. Interrupting previous request."
			);
			if old_ch.send(vec![]).is_err() {
				debug!("Could not send empty response to collisionned request, probably because request was interrupted. Dropping response.");
			}
		}

		trace!("request: query_send {}, {} bytes", id, bytes.len());
		query_send.send((id, prio, bytes))?;

		let resp = resp_recv.await?;
		if resp.is_empty() {
			return Err(Error::Message(
				"Response is 0 bytes, either a collision or a protocol error".into(),
			));
		}

		trace!("request response {}: ", id);

		let code = resp[0];
		if code == 0 {
			Ok(rmp_serde::decode::from_read_ref::<
				_,
				<T as Message>::Response,
			>(&resp[1..])?)
		} else {
			Err(Error::Remote(format!("Remote error code {}", code)))
		}
	}
}

impl SendLoop for ClientConn {}

#[async_trait]
impl RecvLoop for ClientConn {
	fn recv_handler(self: &Arc<Self>, id: RequestID, msg: Vec<u8>) {
		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() {
				debug!("Could not send request response, probably because request was interrupted. Dropping response.");
			}
		}
	}
}