use crate::{
ClientError, ClientOptions, ClientState, ConnectionState, DiffReviewEvent, DiffScope,
DiffSnapshot, ReconnectPolicy, RemoteError, RepositoryAction, platform,
protocol::{
client::{ClientCommand, LocalClientTransport, capabilities},
server::ServerEvent,
shared::{Event, LIVE_PROTOCOL_VERSION, RemoteErrorCode},
},
transport::{ClientMessageTransport, Decoded, Transport},
};
#[cfg(feature = "websocket")]
use crate::{ConnectionHeader, transport::Reconnecting};
use async_channel::{Receiver, Sender, unbounded};
use futures_util::FutureExt;
use std::{sync::Arc, time::Duration};
use tokio::sync::{oneshot, watch};
const IDLE_TIMEOUT: Duration = Duration::from_secs(45);
type Reply = oneshot::Sender<Result<(), ClientError>>;
type Ready = oneshot::Receiver<Result<(), ClientError>>;
#[derive(Clone)]
pub struct DiffClient {
commands: Sender<Command>,
state: watch::Receiver<Arc<ClientState>>,
}
pub struct ClientSubscription {
state: watch::Receiver<Arc<ClientState>>,
}
impl ClientSubscription {
#[must_use]
pub fn latest(&self) -> Arc<ClientState> {
self.state.borrow().clone()
}
pub async fn changed(&mut self) -> Result<Arc<ClientState>, ClientError> {
self.state
.changed()
.await
.map_err(|_| ClientError::Disconnected)?;
Ok(self.state.borrow_and_update().clone())
}
pub async fn wait_until(
&mut self,
timeout: Duration,
mut predicate: impl FnMut(&ClientState) -> bool,
) -> Result<Arc<ClientState>, ClientError> {
let deadline = platform::sleep(timeout).fuse();
futures_util::pin_mut!(deadline);
loop {
let state = self.latest();
if predicate(&state) {
return Ok(state);
}
let changed = self.changed().fuse();
futures_util::pin_mut!(changed);
futures_util::select! {
state = changed => { state?; }
() = deadline => return Err(ClientError::Disconnected),
}
}
}
}
impl DiffClient {
#[cfg(feature = "websocket")]
pub async fn connect(url: &str, options: ClientOptions) -> Result<Self, ClientError> {
Self::connect_with_headers(url, options, Vec::new()).await
}
#[cfg(feature = "websocket")]
pub async fn connect_with_headers(
url: &str,
options: ClientOptions,
headers: Vec<ConnectionHeader>,
) -> Result<Self, ClientError> {
let transport = Reconnecting::connect(url, headers).await?;
Self::start(transport, options).await
}
pub async fn from_transport(
transport: LocalClientTransport,
options: ClientOptions,
) -> Result<Self, ClientError> {
Self::start(transport, options).await
}
pub async fn from_messages(
transport: impl ClientMessageTransport,
options: ClientOptions,
) -> Result<Self, ClientError> {
Self::start(Decoded::new(transport), options).await
}
#[must_use]
pub fn spawn(transport: impl ClientMessageTransport, options: ClientOptions) -> Self {
Self::start_worker(Decoded::new(transport), options).0
}
async fn start(transport: impl Transport, options: ClientOptions) -> Result<Self, ClientError> {
let (client, ready) = Self::start_worker(transport, options);
ready.await.map_err(|_| ClientError::Disconnected)??;
Ok(client)
}
fn start_worker(transport: impl Transport, options: ClientOptions) -> (Self, Ready) {
let (commands, rx) = unbounded();
let (state_tx, state) = watch::channel(Arc::new(ClientState::default()));
let (ready_tx, ready_rx) = oneshot::channel();
platform::spawn(async move {
Worker {
commands: rx,
state_tx,
options,
state: WorkerState::default(),
connection: ConnectionState::Connecting,
closed: false,
}
.run(transport, ready_tx)
.await;
});
(Self { commands, state }, ready_rx)
}
#[must_use]
pub fn state(&self) -> Arc<ClientState> {
self.state.borrow().clone()
}
#[must_use]
pub fn subscribe(&self) -> ClientSubscription {
ClientSubscription {
state: self.state.clone(),
}
}
pub async fn set_scope(&self, scope: DiffScope) -> Result<(), ClientError> {
self.request(ClientCommand::SetScope(scope)).await
}
pub async fn apply(&self, action: RepositoryAction) -> Result<(), ClientError> {
self.request(ClientCommand::Apply(action)).await
}
pub async fn refresh(&self) -> Result<(), ClientError> {
self.request(ClientCommand::Refresh).await
}
pub async fn handle(&self, event: DiffReviewEvent) -> Result<(), ClientError> {
match self.dispatch(event)? {
Some(reply) => reply.await.map_err(|_| ClientError::Disconnected)?,
None => Ok(()),
}
}
pub fn dispatch(
&self,
event: DiffReviewEvent,
) -> Result<Option<oneshot::Receiver<Result<(), ClientError>>>, ClientError> {
let request = match event {
DiffReviewEvent::RepositoryAction(action) => ClientCommand::Apply(action),
DiffReviewEvent::SetScope(scope) => ClientCommand::SetScope(scope),
DiffReviewEvent::Refresh => ClientCommand::Refresh,
DiffReviewEvent::SubmitReview(submission) => ClientCommand::Submit(submission),
DiffReviewEvent::Cancel => ClientCommand::Cancel,
DiffReviewEvent::CopyFormattedReview(_) => return Ok(None),
};
let (reply, receiver) = oneshot::channel();
self.commands
.try_send(Command::Request(request, reply))
.map_err(|_| ClientError::Disconnected)?;
Ok(Some(receiver))
}
pub async fn close(self) -> Result<(), ClientError> {
self.command(Command::Close).await
}
async fn request(&self, request: ClientCommand) -> Result<(), ClientError> {
self.command(|reply| Command::Request(request, reply)).await
}
async fn command(&self, build: impl FnOnce(Reply) -> Command) -> Result<(), ClientError> {
let (reply, rx) = oneshot::channel();
self.commands
.send(build(reply))
.await
.map_err(|_| ClientError::Disconnected)?;
rx.await.map_err(|_| ClientError::Disconnected)?
}
}
enum Command {
Request(ClientCommand, Reply),
Close(Reply),
}
enum ConnectionInput {
Command(Option<Command>),
Event(Result<ServerEvent, ClientError>),
Idle,
}
struct Pending {
reply: Reply,
outcome_unknown_on_disconnect: bool,
}
struct Worker {
commands: Receiver<Command>,
state_tx: watch::Sender<Arc<ClientState>>,
options: ClientOptions,
state: WorkerState,
connection: ConnectionState,
closed: bool,
}
#[derive(Default)]
struct WorkerState {
snapshot: Option<Arc<DiffSnapshot>>,
error: Option<RemoteError>,
pending: Option<Pending>,
initialized: bool,
}
impl Worker {
fn publish(&self) {
let connected = matches!(self.connection, ConnectionState::Connected);
self.state_tx.send_replace(Arc::new(ClientState {
capabilities: capabilities(connected, self.state.snapshot.is_some()),
snapshot: self.state.snapshot.clone(),
connection: self.connection.clone(),
error: self.state.error.clone(),
}));
}
async fn run<T: Transport>(mut self, mut transport: T, ready: Reply) {
let mut ready = Some(ready);
loop {
let result = self.connection(&mut transport, &mut ready).await;
transport.close().await;
self.abandon_request();
let retry = !self.closed
&& ready.is_none()
&& transport.reconnects()
&& matches!(self.options.reconnect, ReconnectPolicy::Retry)
&& !is_terminal(&result);
if let Some(ready) = ready.take() {
let _ = ready.send(result.clone());
}
if !retry {
self.connection =
ConnectionState::Failed(result.err().unwrap_or(ClientError::Disconnected));
self.publish();
return;
}
self.connection = ConnectionState::Connecting;
self.publish();
if !self.reconnect(&mut transport).await {
break;
}
}
self.connection = ConnectionState::Failed(ClientError::Disconnected);
self.publish();
}
async fn connection<T: Transport>(
&mut self,
transport: &mut T,
ready: &mut Option<Reply>,
) -> Result<(), ClientError> {
self.state.initialized = false;
transport
.send(ClientCommand::Initialize {
protocol_version: LIVE_PROTOCOL_VERSION,
scope: self.options.scope,
})
.await
.map_err(|_| ClientError::Disconnected)?;
let idle = platform::sleep(IDLE_TIMEOUT).fuse();
futures_util::pin_mut!(idle);
loop {
let input = {
let command = async {
if self.state.pending.is_none() {
self.commands.recv().await.ok()
} else {
futures_util::future::pending().await
}
}
.fuse();
let event = transport.recv().fuse();
futures_util::pin_mut!(command, event);
futures_util::select! {
command = command => ConnectionInput::Command(command),
event = event => ConnectionInput::Event(event),
() = idle => ConnectionInput::Idle,
}
};
match input {
ConnectionInput::Command(None) => {
self.closed = true;
return Ok(());
}
ConnectionInput::Command(Some(Command::Close(reply))) => {
self.closed = true;
let _ = reply.send(Ok(()));
return Ok(());
}
ConnectionInput::Command(Some(Command::Request(request, reply))) => {
self.request(transport, request, reply).await?;
}
ConnectionInput::Event(event) => {
self.event(event?)?;
if self.state.initialized
&& let Some(ready) = ready.take()
{
let _ = ready.send(Ok(()));
}
}
ConnectionInput::Idle => return Err(ClientError::Disconnected),
}
idle.set(platform::sleep(IDLE_TIMEOUT).fuse());
}
}
fn event(&mut self, event: ServerEvent) -> Result<(), ClientError> {
match event {
Event::Initialize {
protocol_version, ..
} => {
if protocol_version != LIVE_PROTOCOL_VERSION {
return Err(ClientError::Remote(RemoteError::new(
RemoteErrorCode::UnsupportedVersion,
"unsupported live protocol version",
)));
}
if std::mem::replace(&mut self.state.initialized, true) {
return Err(ClientError::Protocol("repeated initialization".to_owned()));
}
self.state.error = None;
self.publish();
Ok(())
}
Event::Error(error) => Err(ClientError::Remote(error)),
_ if !self.state.initialized => {
Err(ClientError::Protocol("expected initialization".to_owned()))
}
Event::Document(snapshot) => {
self.state.snapshot = Some(snapshot);
self.connection = ConnectionState::Connected;
self.publish();
Ok(())
}
Event::RequestResult(result) => {
if let Some(pending) = self.state.pending.take() {
let _ = pending.reply.send(result.map_err(ClientError::Remote));
}
Ok(())
}
Event::Health { error } => {
if self.state.error != error {
self.state.error = error;
self.publish();
}
Ok(())
}
}
}
async fn request<T: Transport>(
&mut self,
transport: &mut T,
request: ClientCommand,
reply: Reply,
) -> Result<(), ClientError> {
if !self.state.initialized {
let _ = reply.send(Err(ClientError::Disconnected));
return Ok(());
}
let outcome_unknown_on_disconnect = matches!(
request,
ClientCommand::Apply(_) | ClientCommand::Submit(_) | ClientCommand::Cancel
);
if transport.send(request).await.is_err() {
let _ = reply.send(Err(ClientError::Disconnected));
return Err(ClientError::Disconnected);
}
self.state.pending = Some(Pending {
reply,
outcome_unknown_on_disconnect,
});
Ok(())
}
fn abandon_request(&mut self) {
if let Some(pending) = self.state.pending.take() {
let _ = pending
.reply
.send(Err(if pending.outcome_unknown_on_disconnect {
ClientError::OutcomeUnknown
} else {
ClientError::Disconnected
}));
}
}
async fn reconnect<T: Transport>(&mut self, transport: &mut T) -> bool {
let mut delay = 250;
loop {
let wait = delay;
let connecting = async {
platform::sleep(platform::reconnect_delay(wait)).await;
transport.reconnect().await
}
.fuse();
let command = self.commands.recv().fuse();
futures_util::pin_mut!(connecting, command);
futures_util::select! {
next = connecting => match next {
Ok(()) => return true,
Err(_) => delay = (delay * 2).min(5000),
},
command = command => match command {
Ok(Command::Request(_, reply)) => { let _ = reply.send(Err(ClientError::Disconnected)); }
Ok(Command::Close(reply)) => { self.closed = true; let _ = reply.send(Ok(())); return false; }
Err(_) => { self.closed = true; return false; }
},
}
}
}
}
fn is_terminal(result: &Result<(), ClientError>) -> bool {
matches!(
result,
Err(ClientError::Protocol(_)
| ClientError::Remote(RemoteError {
code: RemoteErrorCode::UnsupportedVersion | RemoteErrorCode::Protocol,
..
}))
)
}