import os import shutil import traceback from datetime import datetime, timedelta import git import pygit2 from . import log from .base import ProcessName from .env import Environment from .lock import MutGuard, SyncLock, mut from .repos import pygit2_get, wrapper_get from git.repo.base import Repo from pygit2 import Worktree as PyGit2Worktree from pygit2.repository import Repository from typing import Any, Iterator, Optional, Tuple env = Environment() logger = log.get_logger(__name__) def cleanup(git_gecko: Repo, git_wpt: Repo) -> None: for repo in [git_gecko, git_wpt]: pygit2_repo = pygit2_get(repo) cleanup_repo(repo, pygit2_repo, get_max_worktree_count(repo)) def cleanup_repo(repo: Repo, pygit2_repo: Repository, max_count: Optional[int] = None) -> None: # TODO: Always cleanup repos where the sync is finished prune_worktrees(pygit2_repo) unprunable = [] maybe_prunable = [] prunable = [] now = datetime.now() for worktree in worktrees(pygit2_repo): if not os.path.exists(worktree.path): worktree.prune(True) continue process_name = ProcessName.from_tuple(worktree.name.split("-")) worktree_data = (datetime.fromtimestamp(os.stat(worktree.path).st_mtime), process_name, worktree) if process_name is None: logger.warning("Worktree doesn't correspond to a sync %s" % worktree.path) unprunable.append(worktree_data) continue pygit2_repo = pygit2.Repository(worktree.path) head_branch = pygit2_repo.head.name if not head_branch or head_branch == "HEAD": logger.warning("No branch associated with worktree %s" % worktree.path) maybe_prunable.append(worktree_data) continue if head_branch.startswith("refs/heads/"): head_branch = head_branch[len("refs/heads/"):] branch_process_name = ProcessName.from_path(head_branch) if branch_process_name is None: logger.warning("No sync head associated with worktree %s" % worktree.path) maybe_prunable.append(worktree_data) continue if branch_process_name != process_name: logger.warning("Head branch doesn't match worktree %s" % worktree.path) maybe_prunable.append(worktree_data) continue prunable.append(worktree_data) if max_count and len(unprunable) > max_count: logger.error("Unable to cleanup worktrees, because there are too many unprunable worktrees") if not max_count: delete_count = 0 else: delete_count = max(len(unprunable) + len(prunable) + len(maybe_prunable) - max_count, 0) prunable.sort() maybe_prunable.sort() for time, process_name, worktree in (prunable + maybe_prunable): assert process_name is not None if time < (now - timedelta(days=2)): logger.info("Removing worktree without recent activity %s" % worktree.path) delete_worktree(repo, process_name, worktree) delete_count -= 1 elif delete_count > 0: logger.info("Removing LRU worktree %s" % worktree.path) delete_worktree(repo, process_name, worktree) delete_count -= 1 else: break def delete_worktree(repo: Repo, process_name: ProcessName, worktree: PyGit2Worktree) -> None: assert worktree.path.startswith(os.path.join(env.config["root"], env.config["paths"]["worktrees"])) with SyncLock.for_process(process_name): try: logger.info("Deleting path %s" % worktree.path) shutil.rmtree(worktree.path) except Exception: logger.warning("Failed to remove worktree %s:%s" % (worktree.path, traceback.format_exc())) else: logger.info(f"Removed worktree {worktree.path}") worktree.prune(True) wrapper = wrapper_get(repo) assert wrapper is not None wrapper.after_worktree_delete(worktree.path) def worktrees(pygit2_repo: Repository) -> Iterator[PyGit2Worktree]: for name in pygit2_repo.list_worktrees(): yield pygit2_repo.lookup_worktree(name) def prune_worktrees(pygit2_repo: Repository) -> None: for worktree in worktrees(pygit2_repo): # For some reason libgit2 thinks worktrees are not prunable when their # working dir is gone if worktree.is_prunable or not os.path.exists(worktree.path): logger.info("Deleting worktree at path %s" % worktree.path) worktree.prune(True) def get_max_worktree_count(repo: Repo) -> Optional[Any]: repo_wrapper = wrapper_get(repo) if not repo_wrapper: return None repo_name = repo_wrapper.name max_count = env.config[repo_name]["worktree"]["max-count"] if not max_count: return None max_count = int(max_count) if max_count <= 0: return None return max_count class Worktree: """Wrapper for accessing a git worktree for a specific process. To access the worktree call .get() """ def __init__(self, repo: Repo, process_name: ProcessName) -> None: self.repo = repo self.pygit2_repo = pygit2_get(repo) self._worktree: Optional[Repo] = None self.process_name = process_name self.worktree_name = "-".join(str(item) for item in self.process_name.as_tuple()) working_dir = repo.working_dir assert working_dir is not None self.path = os.path.join(env.config["root"], env.config["paths"]["worktrees"], os.path.basename(working_dir), process_name.subtype, process_name.obj_id) self._lock = None def as_mut(self, lock: SyncLock) -> MutGuard: return MutGuard(lock, self) @property def lock_key(self) -> Tuple[str, str]: return (self.process_name.subtype, self.process_name.obj_id) @mut() def get(self) -> Repo: """Return the worktree. On first access, the worktree is reset to the current HEAD. Subsequent access doesn't perform the same check, so it's possible to retain state within a specific process.""" # TODO: We can get the worktree to only checkout the paths we actually # need. # To do this we have to # * Enable sparse checkouts by setting core.sparseCheckouts # * Add the worktree with --no-checkout # * Add the list of paths to check out under $REPO/worktrees/info/sparse-checkout # * Go to the worktree and check it out if self._worktree is None: all_worktrees = {item.name: item for item in worktrees(self.pygit2_repo)} count = len(all_worktrees) max_count = get_max_worktree_count(self.repo) if max_count and count >= max_count: cleanup_repo(self.repo, self.pygit2_repo, max_count - 1) path_exists = os.path.exists(self.path) if self.worktree_name in all_worktrees and not path_exists: prune_worktrees(self.pygit2_repo) del all_worktrees[self.worktree_name] if self.worktree_name not in all_worktrees: if path_exists: logger.warning("Found existing content in worktree path %s, removing" % self.path) shutil.rmtree(self.path) logger.info(f"Creating worktree {self.worktree_name} at {self.path}") if not os.path.exists(os.path.dirname(self.path)): os.makedirs(os.path.dirname(self.path)) worktree = self.pygit2_repo.add_worktree(self.worktree_name, os.path.abspath(self.path), self.pygit2_repo.lookup_reference( "refs/heads/%s" % self.process_name)) wrapper = wrapper_get(self.repo) assert wrapper is not None wrapper.after_worktree_create(self.path) else: worktree = self.pygit2_repo.lookup_worktree(self.worktree_name) assert os.path.exists(self.path) assert worktree.path == self.path self._worktree = git.Repo(self.path) # TODO: In general the worktree should be on the right branch, but it would # be good to check. In the specific case of landing, we move the wpt worktree # around various commits, so it isn't necessarily on the correct branch assert self._worktree is not None return self._worktree @mut() def delete(self) -> None: if not os.path.exists(self.path): return try: worktree = self.pygit2_repo.lookup_worktree(self.worktree_name) except Exception: worktree = None if worktree is None: for worktree in worktrees(self.pygit2_repo): if worktree.path == self.path: break else: # No worktree found return assert worktree.path == self.path delete_worktree(self.repo, self.process_name, worktree)