mod connection_pool; use crate::api::{CloudflareIpCountryHeader, SystemIdHeader};use crate::{    AppState, Error, Result, auth,    db::{        self, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser, Database,        InviteMemberResult, MembershipUpdated, NotificationId, ProjectId, RejoinedProject,        RemoveChannelMemberResult, RespondToChannelInvite, RoomId, ServerId, SharedThreadId, User,        UserId,    },    executor::Executor,};use anyhow::{Context as _, anyhow, bail};use async_tungstenite::tungstenite::{    Message as TungsteniteMessage, protocol::CloseFrame as TungsteniteCloseFrame,};use axum::headers::UserAgent;use axum::{    Extension, Router, TypedHeader,    body::Body,    extract::{        ConnectInfo, WebSocketUpgrade,        ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage},    },    headers::{Header, HeaderName},    http::StatusCode,    middleware,    response::IntoResponse,    routing::get,};use collections::{HashMap, HashSet};pub use connection_pool::{ConnectionPool, ZedVersion};use core::fmt::{self, Debug, Formatter};use futures::TryFutureExt as _;use rpc::proto::split_repository_update;use tracing::Span;use util::paths::PathStyle; use futures::{    FutureExt, SinkExt, StreamExt, TryStreamExt, channel::oneshot, future::BoxFuture,    stream::FuturesUnordered,};use prometheus::{IntGauge, register_int_gauge};use rpc::{    Connection, ConnectionId, ErrorCode, ErrorCodeExt, ErrorExt, Peer, Receipt, TypedEnvelope,    proto::{        self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo,        RequestMessage, ShareProject, UpdateChannelBufferCollaborators,    },};use semver::Version;use std::{    any::TypeId,    future::Future,    marker::PhantomData,    mem,    net::SocketAddr,    ops::{Deref, DerefMut},    rc::Rc,    sync::{        Arc, OnceLock,        atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},    },    time::{Duration, Instant},};use tokio::sync::{Semaphore, watch};use tower::ServiceBuilder;use tracing::{    Instrument,    field::{self},    info_span, instrument,}; pub const RECONNECT_TIMEOUT: Duration = Duration::from_secs(30); // kubernetes gives terminated pods 10s to shutdown gracefully. After they're gone, we can clean up old resources.pub const CLEANUP_TIMEOUT: Duration = Duration::from_secs(15); const NOTIFICATION_COUNT_PER_PAGE: usize = 50;const MAX_CONCURRENT_CONNECTIONS: usize = 512; static CONCURRENT_CONNECTIONS: AtomicUsize = AtomicUsize::new(0); const TOTAL_DURATION_MS: &str = "total_duration_ms";const PROCESSING_DURATION_MS: &str = "processing_duration_ms";const QUEUE_DURATION_MS: &str = "queue_duration_ms";const HOST_WAITING_MS: &str = "host_waiting_ms"; type MessageHandler =    Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session, Span) -> BoxFuture<'static, ()>>; pub struct ConnectionGuard; impl ConnectionGuard {    pub fn try_acquire() -> Result<Self, ()> {        let current_connections = CONCURRENT_CONNECTIONS.fetch_add(1, SeqCst);        if current_connections >= MAX_CONCURRENT_CONNECTIONS {            CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst);            tracing::error!(                "too many concurrent connections: {}",                current_connections + 1            );            return Err(());        }        Ok(ConnectionGuard)    }} impl Drop for ConnectionGuard {    fn drop(&mut self) {        CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst);    }} struct Response<R> {    peer: Arc<Peer>,    receipt: Receipt<R>,    responded: Arc<AtomicBool>,} impl<R: RequestMessage> Response<R> {    fn send(self, payload: R::Response) -> Result<()> {        self.responded.store(true, SeqCst);        self.peer.respond(self.receipt, payload)?;        Ok(())    }} #[derive(Clone, Debug)]pub enum Principal {    User(User),} impl Principal {    fn update_span(&self, span: &tracing::Span) {        match &self {            Principal::User(user) => {                span.record("user_id", user.id.0);                span.record("login", &user.github_login);            }        }    }} #[derive(Clone)]struct MessageContext {    session: Session,    span: tracing::Span,} impl Deref for MessageContext {    type Target = Session;     fn deref(&self) -> &Self::Target {        &self.session    }} impl MessageContext {    pub fn forward_request<T: RequestMessage>(        &self,        receiver_id: ConnectionId,        request: T,    ) -> impl Future<Output = anyhow::Result<T::Response>> {        let request_start_time = Instant::now();        let span = self.span.clone();        tracing::info!("start forwarding request");        self.peer            .forward_request(self.connection_id, receiver_id, request)            .inspect(move |_| {                span.record(                    HOST_WAITING_MS,                    request_start_time.elapsed().as_micros() as f64 / 1000.0,                );            })            .inspect_err(|_| tracing::error!("error forwarding request"))            .inspect_ok(|_| tracing::info!("finished forwarding request"))    }} #[derive(Clone)]struct Session {    principal: Principal,    connection_id: ConnectionId,    db: Arc<tokio::sync::Mutex<DbHandle>>,    peer: Arc<Peer>,    connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,    app_state: Arc<AppState>,    /// The GeoIP country code for the user.    #[allow(unused)]    geoip_country_code: Option<String>,    #[allow(unused)]    system_id: Option<String>,    _executor: Executor,} impl Session {    async fn db(&self) -> tokio::sync::MutexGuard<'_, DbHandle> {        #[cfg(feature = "test-support")]        tokio::task::yield_now().await;        let guard = self.db.lock().await;        #[cfg(feature = "test-support")]        tokio::task::yield_now().await;        guard    }     async fn connection_pool(&self) -> ConnectionPoolGuard<'_> {        #[cfg(feature = "test-support")]        tokio::task::yield_now().await;        let guard = self.connection_pool.lock();        ConnectionPoolGuard {            guard,            _not_send: PhantomData,        }    }     #[expect(dead_code)]    fn is_staff(&self) -> bool {        match &self.principal {            Principal::User(user) => user.admin,        }    }     fn user_id(&self) -> UserId {        match &self.principal {            Principal::User(user) => user.id,        }    }} impl Debug for Session {    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {        let mut result = f.debug_struct("Session");        match &self.principal {            Principal::User(user) => {                result.field("user", &user.github_login);            }        }        result.field("connection_id", &self.connection_id).finish()    }} struct DbHandle(Arc<Database>); impl Deref for DbHandle {    type Target = Database;     fn deref(&self) -> &Self::Target {        self.0.as_ref()    }} pub struct Server {    id: parking_lot::Mutex<ServerId>,    peer: Arc<Peer>,    pub connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,    app_state: Arc<AppState>,    handlers: HashMap<TypeId, MessageHandler>,    teardown: watch::Sender<bool>,} struct ConnectionPoolGuard<'a> {    guard: parking_lot::MutexGuard<'a, ConnectionPool>,    _not_send: PhantomData<Rc<()>>,} impl Server {    pub fn new(id: ServerId, app_state: Arc<AppState>) -> Arc<Self> {        let mut server = Self {            id: parking_lot::Mutex::new(id),            peer: Peer::new(id.0 as u32),            app_state,            connection_pool: Default::default(),            handlers: Default::default(),            teardown: watch::channel(false).0,        };         server            .add_request_handler(ping)            .add_request_handler(create_room)            .add_request_handler(join_room)            .add_request_handler(rejoin_room)            .add_request_handler(leave_room)            .add_request_handler(set_room_participant_role)            .add_request_handler(call)            .add_request_handler(cancel_call)            .add_message_handler(decline_call)            .add_request_handler(update_participant_location)            .add_request_handler(share_project)            .add_message_handler(unshare_project)            .add_request_handler(join_project)            .add_message_handler(leave_project)            .add_request_handler(update_project)            .add_request_handler(update_worktree)            .add_request_handler(update_repository)            .add_request_handler(remove_repository)            .add_message_handler(start_language_server)            .add_message_handler(update_language_server)            .add_message_handler(update_diagnostic_summary)            .add_message_handler(update_worktree_settings)            .add_request_handler(forward_read_only_project_request::<proto::FindSearchCandidates>)            .add_request_handler(forward_read_only_project_request::<proto::GetDocumentHighlights>)            .add_request_handler(forward_read_only_project_request::<proto::GetDocumentSymbols>)            .add_request_handler(forward_read_only_project_request::<proto::GetProjectSymbols>)            .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>)            .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>)            .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>)            .add_request_handler(forward_read_only_project_request::<proto::ResolveInlayHint>)            .add_request_handler(forward_read_only_project_request::<proto::GetColorPresentation>)            .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>)            .add_request_handler(forward_read_only_project_request::<proto::OpenImageByPath>)            .add_request_handler(forward_read_only_project_request::<proto::DownloadFileByPath>)            .add_request_handler(forward_read_only_project_request::<proto::GitGetBranches>)            .add_request_handler(forward_read_only_project_request::<proto::GetDefaultBranch>)            .add_request_handler(forward_read_only_project_request::<proto::OpenUnstagedDiff>)            .add_request_handler(forward_read_only_project_request::<proto::OpenUncommittedDiff>)            .add_request_handler(forward_read_only_project_request::<proto::LspExtExpandMacro>)            .add_request_handler(forward_read_only_project_request::<proto::LspExtOpenDocs>)            .add_request_handler(forward_mutating_project_request::<proto::LspExtRunnables>)            .add_request_handler(                forward_read_only_project_request::<proto::LspExtSwitchSourceHeader>,            )            .add_request_handler(forward_read_only_project_request::<proto::LspExtGoToParentModule>)            .add_request_handler(forward_read_only_project_request::<proto::LspExtCancelFlycheck>)            .add_request_handler(forward_read_only_project_request::<proto::LspExtRunFlycheck>)            .add_request_handler(forward_read_only_project_request::<proto::LspExtClearFlycheck>)            .add_request_handler(                forward_mutating_project_request::<proto::RegisterBufferWithLanguageServers>,            )            .add_request_handler(forward_mutating_project_request::<proto::UpdateGitBranch>)            .add_request_handler(forward_mutating_project_request::<proto::GetCompletions>)            .add_request_handler(                forward_mutating_project_request::<proto::ApplyCompletionAdditionalEdits>,            )            .add_request_handler(forward_mutating_project_request::<proto::OpenNewBuffer>)            .add_request_handler(                forward_mutating_project_request::<proto::ResolveCompletionDocumentation>,            )            .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeAction>)            .add_request_handler(forward_mutating_project_request::<proto::PrepareRename>)            .add_request_handler(forward_mutating_project_request::<proto::PerformRename>)            .add_request_handler(forward_mutating_project_request::<proto::ReloadBuffers>)            .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeActionKind>)            .add_request_handler(forward_mutating_project_request::<proto::FormatBuffers>)            .add_request_handler(forward_mutating_project_request::<proto::CreateProjectEntry>)            .add_request_handler(forward_mutating_project_request::<proto::RenameProjectEntry>)            .add_request_handler(forward_mutating_project_request::<proto::CopyProjectEntry>)            .add_request_handler(forward_mutating_project_request::<proto::DeleteProjectEntry>)            .add_request_handler(forward_mutating_project_request::<proto::ExpandProjectEntry>)            .add_request_handler(                forward_mutating_project_request::<proto::ExpandAllForProjectEntry>,            )            .add_request_handler(forward_mutating_project_request::<proto::OnTypeFormatting>)            .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>)            .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>)            .add_request_handler(lsp_query)            .add_message_handler(broadcast_project_message_from_host::<proto::LspQueryResponse>)            .add_request_handler(forward_mutating_project_request::<proto::RestartLanguageServers>)            .add_request_handler(forward_mutating_project_request::<proto::StopLanguageServers>)            .add_request_handler(forward_mutating_project_request::<proto::LinkedEditingRange>)            .add_message_handler(create_buffer_for_peer)            .add_message_handler(create_image_for_peer)            .add_request_handler(update_buffer)            .add_message_handler(broadcast_project_message_from_host::<proto::RefreshInlayHints>)            .add_message_handler(                broadcast_project_message_from_host::<proto::RefreshSemanticTokens>,            )            .add_message_handler(broadcast_project_message_from_host::<proto::RefreshCodeLens>)            .add_message_handler(broadcast_project_message_from_host::<proto::UpdateBufferFile>)            .add_message_handler(broadcast_project_message_from_host::<proto::BufferReloaded>)            .add_message_handler(broadcast_project_message_from_host::<proto::BufferSaved>)            .add_message_handler(broadcast_project_message_from_host::<proto::UpdateDiffBases>)            .add_message_handler(                broadcast_project_message_from_host::<proto::PullWorkspaceDiagnostics>,            )            .add_request_handler(get_users)            .add_request_handler(fuzzy_search_users)            .add_request_handler(request_contact)            .add_request_handler(remove_contact)            .add_request_handler(respond_to_contact_request)            .add_message_handler(subscribe_to_channels)            .add_request_handler(create_channel)            .add_request_handler(delete_channel)            .add_request_handler(invite_channel_member)            .add_request_handler(remove_channel_member)            .add_request_handler(set_channel_member_role)            .add_request_handler(set_channel_visibility)            .add_request_handler(rename_channel)            .add_request_handler(join_channel_buffer)            .add_request_handler(leave_channel_buffer)            .add_message_handler(update_channel_buffer)            .add_request_handler(rejoin_channel_buffers)            .add_request_handler(get_channel_members)            .add_request_handler(respond_to_channel_invite)            .add_request_handler(join_channel)            .add_request_handler(join_channel_chat)            .add_message_handler(leave_channel_chat)            .add_request_handler(send_channel_message)            .add_request_handler(remove_channel_message)            .add_request_handler(update_channel_message)            .add_request_handler(get_channel_messages)            .add_request_handler(get_channel_messages_by_id)            .add_request_handler(get_notifications)            .add_request_handler(mark_notification_as_read)            .add_request_handler(move_channel)            .add_request_handler(reorder_channel)            .add_request_handler(follow)            .add_message_handler(unfollow)            .add_message_handler(update_followers)            .add_message_handler(acknowledge_channel_message)            .add_message_handler(acknowledge_buffer_version)            .add_request_handler(forward_mutating_project_request::<proto::Stage>)            .add_request_handler(forward_mutating_project_request::<proto::Unstage>)            .add_request_handler(forward_mutating_project_request::<proto::Stash>)            .add_request_handler(forward_mutating_project_request::<proto::StashPop>)            .add_request_handler(forward_mutating_project_request::<proto::StashDrop>)            .add_request_handler(forward_mutating_project_request::<proto::Commit>)            .add_request_handler(forward_mutating_project_request::<proto::RunGitHook>)            .add_request_handler(forward_mutating_project_request::<proto::GitInit>)            .add_request_handler(forward_read_only_project_request::<proto::GetRemotes>)            .add_request_handler(forward_read_only_project_request::<proto::GitShow>)            .add_request_handler(forward_read_only_project_request::<proto::LoadCommitDiff>)            .add_request_handler(forward_read_only_project_request::<proto::GitReset>)            .add_request_handler(forward_read_only_project_request::<proto::GitCheckoutFiles>)            .add_request_handler(forward_mutating_project_request::<proto::SetIndexText>)            .add_request_handler(forward_mutating_project_request::<proto::ToggleBreakpoint>)            .add_message_handler(broadcast_project_message_from_host::<proto::BreakpointsForFile>)            .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>)            .add_request_handler(forward_mutating_project_request::<proto::GitDiff>)            .add_request_handler(forward_mutating_project_request::<proto::GetTreeDiff>)            .add_request_handler(forward_mutating_project_request::<proto::GetBlobContent>)            .add_request_handler(forward_mutating_project_request::<proto::GitCreateBranch>)            .add_request_handler(forward_mutating_project_request::<proto::GitChangeBranch>)            .add_request_handler(forward_mutating_project_request::<proto::GitCreateRemote>)            .add_request_handler(forward_mutating_project_request::<proto::GitRemoveRemote>)            .add_request_handler(forward_read_only_project_request::<proto::GitGetWorktrees>)            .add_request_handler(forward_read_only_project_request::<proto::GitGetHeadSha>)            .add_request_handler(forward_mutating_project_request::<proto::GitCreateWorktree>)            .add_request_handler(disallow_guest_request::<proto::GitRemoveWorktree>)            .add_request_handler(disallow_guest_request::<proto::GitRenameWorktree>)            .add_request_handler(forward_mutating_project_request::<proto::CheckForPushedCommits>)            .add_request_handler(forward_mutating_project_request::<proto::ToggleLspLogs>)            .add_message_handler(broadcast_project_message_from_host::<proto::LanguageServerLog>)            .add_request_handler(share_agent_thread)            .add_request_handler(get_shared_agent_thread)            .add_request_handler(forward_project_search_chunk);         Arc::new(server)    }     pub async fn start(&self) -> Result<()> {        let server_id = *self.id.lock();        let app_state = self.app_state.clone();        let peer = self.peer.clone();        let timeout = self.app_state.executor.sleep(CLEANUP_TIMEOUT);        let pool = self.connection_pool.clone();        let livekit_client = self.app_state.livekit_client.clone();         let span = info_span!("start server");        self.app_state.executor.spawn_detached(            async move {                tracing::info!("waiting for cleanup timeout");                timeout.await;                tracing::info!("cleanup timeout expired, retrieving stale rooms");                 app_state                    .db                    .delete_stale_channel_chat_participants(                        &app_state.config.zed_environment,                        server_id,                    )                    .await                    .trace_err();                 if let Some((room_ids, channel_ids)) = app_state                    .db                    .stale_server_resource_ids(&app_state.config.zed_environment, server_id)                    .await                    .trace_err()                {                    tracing::info!(stale_room_count = room_ids.len(), "retrieved stale rooms");                    tracing::info!(                        stale_channel_buffer_count = channel_ids.len(),                        "retrieved stale channel buffers"                    );                     for channel_id in channel_ids {                        if let Some(refreshed_channel_buffer) = app_state                            .db                            .clear_stale_channel_buffer_collaborators(channel_id, server_id)                            .await                            .trace_err()                        {                            for connection_id in refreshed_channel_buffer.connection_ids {                                peer.send(                                    connection_id,                                    proto::UpdateChannelBufferCollaborators {                                        channel_id: channel_id.to_proto(),                                        collaborators: refreshed_channel_buffer

Lines 1–500 of 4076. Load more