1458 lines · 52.4 KB · Rust
RawDownload
use std::{    path::{Path, PathBuf},    sync::Arc,}; use anyhow::{Context as _, Result, anyhow};use gpui::{App, AsyncApp, Entity, Task};use project::{    LocalProjectFlags, Project, WorktreeId,    git_store::{Repository, resolve_git_worktree_to_main_repo, worktrees_directory_for_repo},    project_settings::ProjectSettings,};use remote::{RemoteConnectionOptions, same_remote_connection_identity};use settings::Settings;use util::{ResultExt, paths::PathStyle};use workspace::{AppState, MultiWorkspace, Workspace}; use crate::thread_metadata_store::{ArchivedGitWorktree, ThreadId, ThreadMetadataStore}; /// The plan for archiving a single git worktree root.////// A thread can have multiple folder paths open, so there may be multiple/// `RootPlan`s per archival operation. Each one captures everything needed to/// persist the worktree's git state and then remove it from disk.////// All fields are gathered synchronously by [`build_root_plan`] while the/// worktree is still loaded in open projects. This is important because/// workspace removal tears down project and repository entities, making/// them unavailable for the later async persist/remove steps.#[derive(Clone)]pub struct RootPlan {    /// Absolute path of the git worktree on disk.    pub root_path: PathBuf,    /// Absolute path to the main git repository this worktree is linked to.    /// Used both for creating a git ref to prevent GC of WIP commits during    /// [`persist_worktree_state`], and for `git worktree remove` during    /// [`remove_root`].    pub main_repo_path: PathBuf,    /// Every open `Project` that has this worktree loaded, so they can all    /// call `remove_worktree` and release it during [`remove_root`].    /// Multiple projects can reference the same path when the user has the    /// worktree open in more than one workspace.    pub affected_projects: Vec<AffectedProject>,    /// The `Repository` entity for this linked worktree, used to run git    /// commands (create WIP commits, stage files, reset) during    /// [`persist_worktree_state`].    pub worktree_repo: Entity<Repository>,    /// The branch the worktree was on, so it can be restored later.    /// `None` if the worktree was in detached HEAD state.    pub branch_name: Option<String>,    /// Remote connection options for the project that owns this worktree,    /// used to create temporary remote projects when the main repo isn't    /// loaded in any open workspace.    pub remote_connection: Option<RemoteConnectionOptions>,} /// A `Project` that references a worktree being archived, paired with the/// `WorktreeId` it uses for that worktree.////// The same worktree path can appear in multiple open workspaces/projects/// (e.g. when the user has two windows open that both include the same/// linked worktree). Each one needs to call `remove_worktree` and wait for/// the release during [`remove_root`], otherwise the project would still/// hold a reference to the directory and `git worktree remove` would fail.#[derive(Clone)]pub struct AffectedProject {    pub project: Entity<Project>,    pub worktree_id: WorktreeId,} fn archived_worktree_ref_name(id: i64) -> String {    format!("refs/archived-worktrees/{}", id)} /// Resolves the Zed-managed worktrees base directory for a given repo.////// This intentionally reads the *global* `git.worktree_directory` setting/// rather than any project-local override, because Zed always uses the/// global value when creating worktrees and the archive check must match.fn worktrees_base_for_repo(    main_repo_path: &Path,    path_style: PathStyle,    cx: &App,) -> Option<PathBuf> {    let setting = &ProjectSettings::get_global(cx).git.worktree_directory;    worktrees_directory_for_repo(main_repo_path, setting, path_style).log_err()} /// Builds a [`RootPlan`] for archiving the git worktree at `path`.////// This is a synchronous planning step that must run *before* any workspace/// removal, because it needs live project and repository entities that are/// torn down when a workspace is removed. It does three things:////// 1. Finds every `Project` across all open workspaces that has this///    worktree loaded (`affected_projects`)./// 2. Looks for a `Repository` entity whose snapshot identifies this path///    as a linked worktree (`worktree_repo`), which is needed for the git///    operations in [`persist_worktree_state`]./// 3. Determines the `main_repo_path` — the parent repo that owns this///    linked worktree — needed for both git ref creation and///    `git worktree remove`.////// Returns `None` if the path is not a linked worktree (main worktrees/// cannot be archived to disk) or if no open project has it loaded.pub fn build_root_plan(    path: &Path,    remote_connection: Option<&RemoteConnectionOptions>,    workspaces: &[Entity<Workspace>],    cx: &App,) -> Option<RootPlan> {    let path = path.to_path_buf();     let matches_target_connection = |project: &Entity<Project>, cx: &App| {        same_remote_connection_identity(            project.read(cx).remote_connection_options(cx).as_ref(),            remote_connection,        )    };     let affected_projects = workspaces        .iter()        .filter_map(|workspace| {            let project = workspace.read(cx).project().clone();            if !matches_target_connection(&project, cx) {                return None;            }            let worktree = project                .read(cx)                .visible_worktrees(cx)                .find(|worktree| worktree.read(cx).abs_path().as_ref() == path.as_path())?;            let worktree_id = worktree.read(cx).id();            Some(AffectedProject {                project,                worktree_id,            })        })        .collect::<Vec<_>>();     if affected_projects.is_empty() {        return None;    }     let linked_repo = workspaces        .iter()        .filter(|workspace| matches_target_connection(workspace.read(cx).project(), cx))        .flat_map(|workspace| {            workspace                .read(cx)                .project()                .read(cx)                .repositories(cx)                .values()                .cloned()                .collect::<Vec<_>>()        })        .find_map(|repo| {            let snapshot = repo.read(cx).snapshot();            (snapshot.is_linked_worktree()                && snapshot.work_directory_abs_path.as_ref() == path.as_path())            .then_some((snapshot, repo))        });     // Only linked worktrees can be archived to disk via `git worktree remove`.    // Main worktrees must be left alone — git refuses to remove them.    let (linked_snapshot, repo) = linked_repo?;    let main_repo_path = linked_snapshot.main_worktree_abs_path()?.to_path_buf();     // Only archive worktrees that live inside the Zed-managed worktrees    // directory (configured via `git.worktree_directory`). Worktrees the    // user created outside that directory should be left untouched.    let worktrees_base = worktrees_base_for_repo(&main_repo_path, linked_snapshot.path_style, cx)?;    if !path.starts_with(&worktrees_base) {        return None;    }     let branch_name = linked_snapshot        .branch        .as_ref()        .map(|branch| branch.name().to_string());     Some(RootPlan {        root_path: path,        main_repo_path,        affected_projects,        worktree_repo: repo,        branch_name,        remote_connection: remote_connection.cloned(),    })} /// Removes a worktree from all affected projects and deletes it from disk/// via `git worktree remove`.////// This is the destructive counterpart to [`persist_worktree_state`]. It/// first detaches the worktree from every [`AffectedProject`], waits for/// each project to fully release it, then asks the main repository to/// delete the worktree directory. If the git removal fails, the worktree/// is re-added to each project via [`rollback_root`].pub async fn remove_root(root: RootPlan, cx: &mut AsyncApp) -> Result<()> {    let release_tasks: Vec<_> = root        .affected_projects        .iter()        .map(|affected| {            let project = affected.project.clone();            let worktree_id = affected.worktree_id;            project.update(cx, |project, cx| {                let wait = project.wait_for_worktree_release(worktree_id, cx);                project.remove_worktree(worktree_id, cx);                wait            })        })        .collect();     if let Err(error) = remove_root_after_worktree_removal(&root, release_tasks, cx).await {        rollback_root(&root, cx).await;        return Err(error);    }     Ok(())} async fn remove_root_after_worktree_removal(    root: &RootPlan,    release_tasks: Vec<Task<Result<()>>>,    cx: &mut AsyncApp,) -> Result<()> {    for task in release_tasks {        if let Err(error) = task.await {            log::error!("Failed waiting for worktree release: {error:#}");        }    }     let (repo, project) =        find_or_create_repository(&root.main_repo_path, root.remote_connection.as_ref(), cx)            .await?;     // `Repository::remove_worktree` with `force = true` deletes the working    // directory before running `git worktree remove --force`, so there's no    // need to touch the filesystem here. For remote projects that cleanup    // runs on the headless server via the `GitRemoveWorktree` RPC, which is    // the only code path with access to the remote machine's filesystem.    let receiver = repo.update(cx, |repo: &mut Repository, _cx| {        repo.remove_worktree(root.root_path.clone(), true)    });    let result = receiver        .await        .map_err(|_| anyhow!("git worktree metadata cleanup was canceled"))?;    // `project` may be a live workspace project or a temporary one created    // by `find_or_create_repository`. In the temporary case we must keep it    // alive until the repo removes the worktree    drop(project);    result.context("git worktree metadata cleanup failed")?;    Ok(())} /// Finds a live `Repository` entity for the given path, or creates a temporary/// project to obtain one.////// `Repository` entities can only be obtained through a `Project` because/// `GitStore` (which creates and manages `Repository` entities) is owned by/// `Project`. When no open workspace contains the repo we need, we spin up a/// headless project just to get a `Repository` handle. For local paths this is/// a `Project::local`; for remote paths we build a `Project::remote` through/// the connection pool (reusing the existing SSH transport), which requires/// the caller to pass the matching `RemoteConnectionOptions` so we only match/// and fall back onto projects that share the same remote identity. The/// caller keeps the returned `Entity<Project>` alive for the duration of the/// git operations, then drops it.////// Future improvement: decoupling `GitStore` from `Project` so that/// `Repository` entities can be created standalone would eliminate this/// temporary-project workaround.async fn find_or_create_repository(    repo_path: &Path,    remote_connection: Option<&RemoteConnectionOptions>,    cx: &mut AsyncApp,) -> Result<(Entity<Repository>, Entity<Project>)> {    let repo_path_owned = repo_path.to_path_buf();    let remote_connection_owned = remote_connection.cloned();     // First, try to find a live repository in any open workspace whose    // remote connection matches (so a local `/project` and a remote    // `/project` are not confused).    let live_repo = cx.update(|cx| {        all_open_workspaces(cx)            .into_iter()            .filter_map(|workspace| {                let project = workspace.read(cx).project().clone();                let project_connection = project.read(cx).remote_connection_options(cx);                if !same_remote_connection_identity(                    project_connection.as_ref(),                    remote_connection_owned.as_ref(),                ) {                    return None;                }                Some((                    project                        .read(cx)                        .repositories(cx)                        .values()                        .find(|repo| {                            repo.read(cx).snapshot().work_directory_abs_path.as_ref()                                == repo_path_owned.as_path()                        })                        .cloned()?,                    project.clone(),                ))            })            .next()    });     if let Some((repo, project)) = live_repo {        return Ok((repo, project));    }     let app_state =        current_app_state(cx).context("no app state available for temporary project")?;     // For remote paths, create a fresh RemoteClient through the connection    // pool (reusing the existing SSH transport) and build a temporary    // remote project. Each RemoteClient gets its own server-side headless    // project, so there are no RPC routing conflicts with other projects.    let temp_project = if let Some(connection) = remote_connection_owned {        let remote_client = cx            .update(|cx| {                if !remote::has_active_connection(&connection, cx) {                    anyhow::bail!("cannot open repository on disconnected remote machine");                }                Ok(remote_connection::connect_reusing_pool(connection, cx))            })?            .await?            .context("remote connection was canceled")?;         cx.update(|cx| {            Project::remote(                remote_client,                app_state.client.clone(),                app_state.node_runtime.clone(),                app_state.user_store.clone(),                app_state.languages.clone(),                app_state.fs.clone(),                false,                cx,            )        })    } else {        cx.update(|cx| {            Project::local(                app_state.client.clone(),                app_state.node_runtime.clone(),                app_state.user_store.clone(),                app_state.languages.clone(),                app_state.fs.clone(),                None,                LocalProjectFlags::default(),                cx,            )        })    };     let repo_path_for_worktree = repo_path.to_path_buf();    let create_worktree = temp_project.update(cx, |project, cx| {        project.create_worktree(repo_path_for_worktree, true, cx)    });    let _worktree = create_worktree.await?;    let initial_scan = temp_project.read_with(cx, |project, cx| project.wait_for_initial_scan(cx));    initial_scan.await;     let repo_path_for_find = repo_path.to_path_buf();    let repo = temp_project        .update(cx, |project, cx| {            project                .repositories(cx)                .values()                .find(|repo| {                    repo.read(cx).snapshot().work_directory_abs_path.as_ref()                        == repo_path_for_find.as_path()                })                .cloned()        })        .context("failed to resolve temporary repository handle")?;     let barrier = repo.update(cx, |repo: &mut Repository, _cx| repo.barrier());    barrier        .await        .map_err(|_| anyhow!("temporary repository barrier canceled"))?;    Ok((repo, temp_project))} /// Re-adds the worktree to every affected project after a failed/// [`remove_root`].async fn rollback_root(root: &RootPlan, cx: &mut AsyncApp) {    for affected in &root.affected_projects {        let task = affected.project.update(cx, |project, cx| {            project.create_worktree(root.root_path.clone(), true, cx)        });        task.await.log_err();    }} /// Saves the worktree's full git state so it can be restored later.////// This creates two detached commits (via [`create_archive_checkpoint`] on/// the `GitRepository` trait) that capture the staged and unstaged state/// without moving any branch ref. The commits are:///   - "WIP staged": a tree matching the current index, parented on HEAD///   - "WIP unstaged": a tree with all files (including untracked),///     parented on the staged commit////// After creating the commits, this function:///   1. Records the commit SHAs, branch name, and paths in a DB record.///   2. Links every thread referencing this worktree to that record.///   3. Creates a git ref on the main repo to prevent GC of the commits.////// On success, returns the archived worktree DB row ID for rollback.pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Result<i64> {    let worktree_repo = root.worktree_repo.clone();     let original_commit_hash = worktree_repo        .update(cx, |repo, _cx| repo.head_sha())        .await        .map_err(|_| anyhow!("head_sha canceled"))?        .context("failed to read original HEAD SHA")?        .context("HEAD SHA is None")?;     // Create two detached WIP commits without moving the branch.    let checkpoint_rx = worktree_repo.update(cx, |repo, _cx| repo.create_archive_checkpoint());    let (staged_commit_hash, unstaged_commit_hash) = checkpoint_rx        .await        .map_err(|_| anyhow!("create_archive_checkpoint canceled"))?        .context("failed to create archive checkpoint")?;     // Create DB record    let store = cx.update(|cx| ThreadMetadataStore::global(cx));    let worktree_path_str = root.root_path.to_string_lossy().to_string();    let main_repo_path_str = root.main_repo_path.to_string_lossy().to_string();    let branch_name = root.branch_name.clone().or_else(|| {        worktree_repo.read_with(cx, |repo, _cx| {            repo.snapshot()                .branch                .as_ref()                .map(|branch| branch.name().to_string())        })    });     let db_result = store        .read_with(cx, |store, cx| {            store.create_archived_worktree(                worktree_path_str.clone(),                main_repo_path_str.clone(),                branch_name.clone(),                staged_commit_hash.clone(),                unstaged_commit_hash.clone(),                original_commit_hash.clone(),                cx,            )        })        .await        .context("failed to create archived worktree DB record");    let archived_worktree_id = match db_result {        Ok(id) => id,        Err(error) => {            return Err(error);        }    };     // Link all threads on this worktree to the archived record    let thread_ids: Vec<ThreadId> = store.read_with(cx, |store, _cx| {        store            .entries()            .filter(|thread| {                thread                    .folder_paths()                    .paths()                    .iter()                    .any(|p| p.as_path() == root.root_path)            })            .map(|thread| thread.thread_id)            .collect()    });     for thread_id in &thread_ids {        let link_result = store            .read_with(cx, |store, cx| {                store.link_thread_to_archived_worktree(*thread_id, archived_worktree_id, cx)            })            .await;        if let Err(error) = link_result {            if let Err(delete_error) = store                .read_with(cx, |store, cx| {                    store.delete_archived_worktree(archived_worktree_id, cx)                })                .await            {                log::error!(                    "Failed to delete archived worktree DB record during link rollback: \                     {delete_error:#}"                );            }

Lines 1–500 of 1458.