1675 lines · 57.8 KB · Rust
RawDownload
#[cfg(any(test, feature = "test-support"))]pub mod test; pub mod telemetry;pub mod user; use anyhow::{anyhow, Context as _, Result};use async_recursion::async_recursion;use async_tungstenite::tungstenite::{    error::Error as WebsocketError,    http::{Request, StatusCode},};use futures::{    channel::oneshot, future::LocalBoxFuture, AsyncReadExt, FutureExt, SinkExt, StreamExt,    TryFutureExt as _, TryStreamExt,};use gpui::{    actions, serde_json, AnyModel, AnyWeakModel, AppContext, AsyncAppContext, Model,    SemanticVersion, Task, WeakModel,};use lazy_static::lazy_static;use parking_lot::RwLock;use postage::watch;use rand::prelude::*;use rpc::proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, PeerId, RequestMessage};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use settings::Settings;use std::{    any::TypeId,    collections::HashMap,    convert::TryFrom,    fmt::Write as _,    future::Future,    marker::PhantomData,    path::PathBuf,    sync::{atomic::AtomicU64, Arc, Weak},    time::{Duration, Instant},};use telemetry::Telemetry;use thiserror::Error;use url::Url;use util::channel::ReleaseChannel;use util::http::HttpClient;use util::{ResultExt, TryFutureExt}; pub use rpc::*;pub use telemetry::ClickhouseEvent;pub use user::*; lazy_static! {    pub static ref ZED_SERVER_URL: String =        std::env::var("ZED_SERVER_URL").unwrap_or_else(|_| "https://zed.dev".to_string());    pub static ref IMPERSONATE_LOGIN: Option<String> = std::env::var("ZED_IMPERSONATE")        .ok()        .and_then(|s| if s.is_empty() { None } else { Some(s) });    pub static ref ADMIN_API_TOKEN: Option<String> = std::env::var("ZED_ADMIN_API_TOKEN")        .ok()        .and_then(|s| if s.is_empty() { None } else { Some(s) });    pub static ref ZED_APP_VERSION: Option<SemanticVersion> = std::env::var("ZED_APP_VERSION")        .ok()        .and_then(|v| v.parse().ok());    pub static ref ZED_APP_PATH: Option<PathBuf> =        std::env::var("ZED_APP_PATH").ok().map(PathBuf::from);    pub static ref ZED_ALWAYS_ACTIVE: bool =        std::env::var("ZED_ALWAYS_ACTIVE").map_or(false, |e| e.len() > 0);} pub const ZED_SECRET_CLIENT_TOKEN: &str = "618033988749894";pub const INITIAL_RECONNECTION_DELAY: Duration = Duration::from_millis(100);pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(5); actions!(client, [SignIn, SignOut, Reconnect]); pub fn init_settings(cx: &mut AppContext) {    TelemetrySettings::register(cx);} pub fn init(client: &Arc<Client>, cx: &mut AppContext) {    init_settings(cx);     let client = Arc::downgrade(client);    cx.on_action({        let client = client.clone();        move |_: &SignIn, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(                    |cx| async move { client.authenticate_and_connect(true, &cx).log_err().await },                )                .detach();            }        }    });     cx.on_action({        let client = client.clone();        move |_: &SignOut, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(|cx| async move {                    client.disconnect(&cx);                })                .detach();            }        }    });     cx.on_action({        let client = client.clone();        move |_: &Reconnect, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(|cx| async move {                    client.reconnect(&cx);                })                .detach();            }        }    });} pub struct Client {    id: AtomicU64,    peer: Arc<Peer>,    http: Arc<dyn HttpClient>,    telemetry: Arc<Telemetry>,    state: RwLock<ClientState>,     #[allow(clippy::type_complexity)]    #[cfg(any(test, feature = "test-support"))]    authenticate: RwLock<        Option<Box<dyn 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>>>,    >,     #[allow(clippy::type_complexity)]    #[cfg(any(test, feature = "test-support"))]    establish_connection: RwLock<        Option<            Box<                dyn 'static                    + Send                    + Sync                    + Fn(                        &Credentials,                        &AsyncAppContext,                    ) -> Task<Result<Connection, EstablishConnectionError>>,            >,        >,    >,} #[derive(Error, Debug)]pub enum EstablishConnectionError {    #[error("upgrade required")]    UpgradeRequired,    #[error("unauthorized")]    Unauthorized,    #[error("{0}")]    Other(#[from] anyhow::Error),    #[error("{0}")]    Http(#[from] util::http::Error),    #[error("{0}")]    Io(#[from] std::io::Error),    #[error("{0}")]    Websocket(#[from] async_tungstenite::tungstenite::http::Error),} impl From<WebsocketError> for EstablishConnectionError {    fn from(error: WebsocketError) -> Self {        if let WebsocketError::Http(response) = &error {            match response.status() {                StatusCode::UNAUTHORIZED => return EstablishConnectionError::Unauthorized,                StatusCode::UPGRADE_REQUIRED => return EstablishConnectionError::UpgradeRequired,                _ => {}            }        }        EstablishConnectionError::Other(error.into())    }} impl EstablishConnectionError {    pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {        Self::Other(error.into())    }} #[derive(Copy, Clone, Debug, PartialEq)]pub enum Status {    SignedOut,    UpgradeRequired,    Authenticating,    Connecting,    ConnectionError,    Connected {        peer_id: PeerId,        connection_id: ConnectionId,    },    ConnectionLost,    Reauthenticating,    Reconnecting,    ReconnectionError {        next_reconnection: Instant,    },} impl Status {    pub fn is_connected(&self) -> bool {        matches!(self, Self::Connected { .. })    }     pub fn is_signed_out(&self) -> bool {        matches!(self, Self::SignedOut | Self::UpgradeRequired)    }} struct ClientState {    credentials: Option<Credentials>,    status: (watch::Sender<Status>, watch::Receiver<Status>),    entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,    _reconnect_task: Option<Task<()>>,    reconnect_interval: Duration,    entities_by_type_and_remote_id: HashMap<(TypeId, u64), WeakSubscriber>,    models_by_message_type: HashMap<TypeId, AnyWeakModel>,    entity_types_by_message_type: HashMap<TypeId, TypeId>,    #[allow(clippy::type_complexity)]    message_handlers: HashMap<        TypeId,        Arc<            dyn Send                + Sync                + Fn(                    AnyModel,                    Box<dyn AnyTypedEnvelope>,                    &Arc<Client>,                    AsyncAppContext,                ) -> LocalBoxFuture<'static, Result<()>>,        >,    >,} enum WeakSubscriber {    Entity { handle: AnyWeakModel },    Pending(Vec<Box<dyn AnyTypedEnvelope>>),} #[derive(Clone, Debug)]pub struct Credentials {    pub user_id: u64,    pub access_token: String,} impl Default for ClientState {    fn default() -> Self {        Self {            credentials: None,            status: watch::channel_with(Status::SignedOut),            entity_id_extractors: Default::default(),            _reconnect_task: None,            reconnect_interval: Duration::from_secs(5),            models_by_message_type: Default::default(),            entities_by_type_and_remote_id: Default::default(),            entity_types_by_message_type: Default::default(),            message_handlers: Default::default(),        }    }} pub enum Subscription {    Entity {        client: Weak<Client>,        id: (TypeId, u64),    },    Message {        client: Weak<Client>,        id: TypeId,    },} impl Drop for Subscription {    fn drop(&mut self) {        match self {            Subscription::Entity { client, id } => {                if let Some(client) = client.upgrade() {                    let mut state = client.state.write();                    let _ = state.entities_by_type_and_remote_id.remove(id);                }            }            Subscription::Message { client, id } => {                if let Some(client) = client.upgrade() {                    let mut state = client.state.write();                    let _ = state.entity_types_by_message_type.remove(id);                    let _ = state.message_handlers.remove(id);                }            }        }    }} pub struct PendingEntitySubscription<T: 'static> {    client: Arc<Client>,    remote_id: u64,    _entity_type: PhantomData<T>,    consumed: bool,} impl<T: 'static> PendingEntitySubscription<T> {    pub fn set_model(mut self, model: &Model<T>, cx: &mut AsyncAppContext) -> Subscription {        self.consumed = true;        let mut state = self.client.state.write();        let id = (TypeId::of::<T>(), self.remote_id);        let Some(WeakSubscriber::Pending(messages)) =            state.entities_by_type_and_remote_id.remove(&id)        else {            unreachable!()        };         state.entities_by_type_and_remote_id.insert(            id,            WeakSubscriber::Entity {                handle: model.downgrade().into(),            },        );        drop(state);        for message in messages {            self.client.handle_message(message, cx);        }        Subscription::Entity {            client: Arc::downgrade(&self.client),            id,        }    }} impl<T: 'static> Drop for PendingEntitySubscription<T> {    fn drop(&mut self) {        if !self.consumed {            let mut state = self.client.state.write();            if let Some(WeakSubscriber::Pending(messages)) = state                .entities_by_type_and_remote_id                .remove(&(TypeId::of::<T>(), self.remote_id))            {                for message in messages {                    log::info!("unhandled message {}", message.payload_type_name());                }            }        }    }} #[derive(Copy, Clone)]pub struct TelemetrySettings {    pub diagnostics: bool,    pub metrics: bool,} #[derive(Default, Clone, Serialize, Deserialize, JsonSchema)]pub struct TelemetrySettingsContent {    pub diagnostics: Option<bool>,    pub metrics: Option<bool>,} impl settings::Settings for TelemetrySettings {    const KEY: Option<&'static str> = Some("telemetry");     type FileContent = TelemetrySettingsContent;     fn load(        default_value: &Self::FileContent,        user_values: &[&Self::FileContent],        _: &mut AppContext,    ) -> Result<Self> {        Ok(Self {            diagnostics: user_values.first().and_then(|v| v.diagnostics).unwrap_or(                default_value                    .diagnostics                    .ok_or_else(Self::missing_default)?,            ),            metrics: user_values                .first()                .and_then(|v| v.metrics)                .unwrap_or(default_value.metrics.ok_or_else(Self::missing_default)?),        })    }} impl Client {    pub fn new(http: Arc<dyn HttpClient>, cx: &mut AppContext) -> Arc<Self> {        Arc::new(Self {            id: AtomicU64::new(0),            peer: Peer::new(0),            telemetry: Telemetry::new(http.clone(), cx),            http,            state: Default::default(),             #[cfg(any(test, feature = "test-support"))]            authenticate: Default::default(),            #[cfg(any(test, feature = "test-support"))]            establish_connection: Default::default(),        })    }     pub fn id(&self) -> u64 {        self.id.load(std::sync::atomic::Ordering::SeqCst)    }     pub fn http_client(&self) -> Arc<dyn HttpClient> {        self.http.clone()    }     pub fn set_id(&self, id: u64) -> &Self {        self.id.store(id, std::sync::atomic::Ordering::SeqCst);        self    }     #[cfg(any(test, feature = "test-support"))]    pub fn teardown(&self) {        let mut state = self.state.write();        state._reconnect_task.take();        state.message_handlers.clear();        state.models_by_message_type.clear();        state.entities_by_type_and_remote_id.clear();        state.entity_id_extractors.clear();        self.peer.teardown();    }     #[cfg(any(test, feature = "test-support"))]    pub fn override_authenticate<F>(&self, authenticate: F) -> &Self    where        F: 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>,    {        *self.authenticate.write() = Some(Box::new(authenticate));        self    }     #[cfg(any(test, feature = "test-support"))]    pub fn override_establish_connection<F>(&self, connect: F) -> &Self    where        F: 'static            + Send            + Sync            + Fn(&Credentials, &AsyncAppContext) -> Task<Result<Connection, EstablishConnectionError>>,    {        *self.establish_connection.write() = Some(Box::new(connect));        self    }     pub fn user_id(&self) -> Option<u64> {        self.state            .read()            .credentials            .as_ref()            .map(|credentials| credentials.user_id)    }     pub fn peer_id(&self) -> Option<PeerId> {        if let Status::Connected { peer_id, .. } = &*self.status().borrow() {            Some(*peer_id)        } else {            None        }    }     pub fn status(&self) -> watch::Receiver<Status> {        self.state.read().status.1.clone()    }     fn set_status(self: &Arc<Self>, status: Status, cx: &AsyncAppContext) {        log::info!("set status on client {}: {:?}", self.id(), status);        let mut state = self.state.write();        *state.status.0.borrow_mut() = status;         match status {            Status::Connected { .. } => {                state._reconnect_task = None;            }            Status::ConnectionLost => {                let this = self.clone();                let reconnect_interval = state.reconnect_interval;                state._reconnect_task = Some(cx.spawn(move |cx| async move {                    #[cfg(any(test, feature = "test-support"))]                    let mut rng = StdRng::seed_from_u64(0);                    #[cfg(not(any(test, feature = "test-support")))]                    let mut rng = StdRng::from_entropy();                     let mut delay = INITIAL_RECONNECTION_DELAY;                    while let Err(error) = this.authenticate_and_connect(true, &cx).await {                        log::error!("failed to connect {}", error);                        if matches!(*this.status().borrow(), Status::ConnectionError) {                            this.set_status(                                Status::ReconnectionError {                                    next_reconnection: Instant::now() + delay,                                },                                &cx,                            );                            cx.background_executor().timer(delay).await;                            delay = delay                                .mul_f32(rng.gen_range(1.0..=2.0))                                .min(reconnect_interval);                        } else {                            break;                        }                    }

Lines 1–500 of 1675.