clankerdiff_client/

client.rs

use crate::{
    ClientError, ClientOptions, ClientState, ConnectionState, DiffReviewEvent, DiffScope,
    ReconnectPolicy, RemoteError, RepositoryAction, platform,
    protocol::{
        client::{ClientCommand, LocalClientTransport},
        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())
    }

    /// Waits until the published state satisfies `predicate`, or `timeout` elapses.
    ///
    /// # Errors
    ///
    /// Returns [`ClientError::Disconnected`] when the client stops publishing or
    /// the timeout elapses first.
    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,
                client: ClientState::default(),
                handshake: Handshake::default(),
                pending: None,
                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,
    client: ClientState,
    handshake: Handshake,
    pending: Option<Pending>,
    closed: bool,
}

#[derive(Default)]
struct Handshake {
    initialized: bool,
}

impl Handshake {
    fn reset(&mut self) {
        self.initialized = false;
    }

    fn check(&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.initialized, true) {
                    return Err(ClientError::Protocol("repeated initialization".to_owned()));
                }
                Ok(())
            }
            Event::Error(_) => Ok(()),
            _ if !self.initialized => {
                Err(ClientError::Protocol("expected initialization".to_owned()))
            }
            _ => Ok(()),
        }
    }
}

impl Worker {
    fn publish(&self) {
        self.state_tx.send_replace(Arc::new(self.client.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.client.set_connection(ConnectionState::Failed(
                    result.err().unwrap_or(ClientError::Disconnected),
                ));
                self.publish();
                return;
            }
            self.client.set_connection(ConnectionState::Connecting);
            self.publish();
            if !self.reconnect(&mut transport).await {
                break;
            }
        }
        self.client
            .set_connection(ConnectionState::Failed(ClientError::Disconnected));
        self.publish();
    }

    async fn connection<T: Transport>(
        &mut self,
        transport: &mut T,
        ready: &mut Option<Reply>,
    ) -> Result<(), ClientError> {
        self.handshake.reset();
        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.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.handshake.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> {
        self.handshake.check(event)?;
        if let Event::RequestResult(result) = event {
            if let Some(pending) = self.pending.take() {
                let _ = pending
                    .reply
                    .send(result.clone().map_err(ClientError::Remote));
            }
            return Ok(());
        }
        let next = self.client.clone().apply(event);
        if next != self.client {
            self.client = next;
            self.publish();
        }
        if let Event::Error(error) = event {
            return Err(ClientError::Remote(error.clone()));
        }
        Ok(())
    }

    async fn request<T: Transport>(
        &mut self,
        transport: &mut T,
        request: ClientCommand,
        reply: Reply,
    ) -> Result<(), ClientError> {
        if !self.handshake.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.pending = Some(Pending {
            reply,
            outcome_unknown_on_disconnect,
        });
        Ok(())
    }

    fn abandon_request(&mut self) {
        if let Some(pending) = self.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,
                ..
            }))
    )
}