[RFC] [ACIX-453] Implement Shared Agent 6 / 7 Tasks

tasks/__init__.py

@@ -62,6 +62,7 @@
     trace_agent,
     vim,
     vscode,
+    worktree,
 )
 from tasks.build_tags import audit_tag_impact, print_default_build_tags
 from tasks.components import lint_components, lint_fxutil_oneshot_test
@@ -210,6 +211,7 @@
 ns.add_collection(collector)
 ns.add_collection(invoke_unit_tests)
 ns.add_collection(debug)
+ns.add_collection(worktree)
 ns.configure(
     {
         "run": {

tasks/git.py

@@ -4,14 +4,14 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import color_message
-from tasks.libs.common.git import get_current_branch
+from tasks.libs.common.git import get_current_branch, get_default_branch
 
 
 @task
 def check_protected_branch(ctx):
     local_branch = get_current_branch(ctx)
 
-    if local_branch == 'main':
+    if local_branch == get_default_branch():
         print(color_message("You're about to commit or push to the main, are you sure this is what you want?", "red"))
         raise Exit(code=1)
 

tasks/github_tasks.py

@@ -19,8 +19,9 @@
     trigger_macos_workflow,
 )
 from tasks.libs.common.color import color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH, DEFAULT_INTEGRATIONS_CORE_BRANCH
+from tasks.libs.common.constants import DEFAULT_INTEGRATIONS_CORE_BRANCH
 from tasks.libs.common.datadog_api import create_gauge, send_event, send_metrics
+from tasks.libs.common.git import get_default_branch
 from tasks.libs.common.junit_upload_core import repack_macos_junit_tar
 from tasks.libs.common.utils import get_git_pretty_ref
 from tasks.libs.owners.linter import codeowner_has_orphans, directory_has_packages_without_owner
@@ -36,7 +37,7 @@ def concurrency_key():
     current_ref = get_git_pretty_ref()
 
     # We want workflows to run to completion on the default branch and release branches
-    if re.search(rf'^({DEFAULT_BRANCH}|\d+\.\d+\.x)$', current_ref):
+    if re.search(rf'^({get_default_branch()}|\d+\.\d+\.x)$', current_ref):
         return None
 
     return current_ref
@@ -68,7 +69,7 @@ def _trigger_macos_workflow(release, destination=None, retry_download=0, retry_i
 def trigger_macos(
     _,
     workflow_type="build",
-    datadog_agent_ref=DEFAULT_BRANCH,
+    datadog_agent_ref=None,
     release_version="nightly-a7",
     major_version="7",
     destination=".",
@@ -79,6 +80,13 @@ def trigger_macos(
     test_washer=False,
     integrations_core_ref=DEFAULT_INTEGRATIONS_CORE_BRANCH,
 ):
+    """
+    Args:
+        datadog_agent_ref: If None, will be the default branch.
+    """
+
+    datadog_agent_ref = datadog_agent_ref or get_default_branch()
+
     if workflow_type == "build":
         conclusion = _trigger_macos_workflow(
             # Provide the release version to be able to fetch the associated

tasks/libs/ciproviders/github_actions_tools.py

@@ -11,13 +11,13 @@
 
 from tasks.libs.ciproviders.github_api import GithubAPI
 from tasks.libs.common.color import color_message
-from tasks.libs.common.utils import DEFAULT_BRANCH
+from tasks.libs.common.git import get_default_branch
 
 
 def trigger_macos_workflow(
     workflow_name="macos.yaml",
     github_action_ref="master",
-    datadog_agent_ref=DEFAULT_BRANCH,
+    datadog_agent_ref=None,
     release_version=None,
     major_version=None,
     gitlab_pipeline_id=None,
@@ -31,6 +31,8 @@ def trigger_macos_workflow(
     """
     Trigger a workflow to build a MacOS Agent.
     """
+
+    datadog_agent_ref = datadog_agent_ref or get_default_branch()
     inputs = {}
 
     if datadog_agent_ref is not None:

tasks/libs/ciproviders/github_api.py

@@ -13,6 +13,7 @@
 
 from tasks.libs.common.color import color_message
 from tasks.libs.common.constants import GITHUB_REPO_NAME
+from tasks.libs.common.git import get_default_branch
 
 try:
     import semver
@@ -494,7 +495,7 @@ def create_release_pr(title, base_branch, target_branch, version, changelog_pr=F
     ]
 
     if changelog_pr:
-        labels.append("backport/main")
+        labels.append(f"backport/{get_default_branch()}")
 
     updated_pr = github.update_pr(
         pull_number=pr.number,

tasks/libs/ciproviders/gitlab_api.py

@@ -23,8 +23,7 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH
-from tasks.libs.common.git import get_common_ancestor, get_current_branch
+from tasks.libs.common.git import get_common_ancestor, get_current_branch, get_default_branch
 from tasks.libs.common.utils import retry_function
 
 BASE_URL = "https://gitlab.ddbuild.io"
@@ -1214,7 +1213,7 @@ def compute_gitlab_ci_config_diff(ctx, before: str, after: str):
     after_name = after or "local files"
 
     # The before commit is the LCA commit between before and after
-    before = before or DEFAULT_BRANCH
+    before = before or get_default_branch()
     before = get_common_ancestor(ctx, before, after or "HEAD")
 
     print(f'Getting after changes config ({color_message(after_name, Color.BOLD)})')

tasks/libs/common/constants.py

@@ -1,4 +1,3 @@
-DEFAULT_BRANCH = "main"
 DEFAULT_INTEGRATIONS_CORE_BRANCH = "master"
 GITHUB_ORG = "DataDog"
 REPO_NAME = "datadog-agent"

tasks/libs/common/git.py

@@ -5,10 +5,10 @@
 from contextlib import contextmanager
 from typing import TYPE_CHECKING
 
+from invoke import Context
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import DEFAULT_BRANCH
 from tasks.libs.common.user_interactions import yes_no_question
 
 if TYPE_CHECKING:
@@ -90,7 +90,9 @@ def get_file_modifications(
     return modifications
 
 
-def get_modified_files(ctx, base_branch="main") -> list[str]:
+def get_modified_files(ctx, base_branch=None) -> list[str]:
+    base_branch = base_branch or get_default_branch()
+
     return get_file_modifications(
         ctx, base_branch=base_branch, added=True, modified=True, only_names=True, no_renames=True
     )
@@ -100,7 +102,23 @@ def get_current_branch(ctx) -> str:
     return ctx.run("git rev-parse --abbrev-ref HEAD", hide=True).stdout.strip()
 
 
-def get_common_ancestor(ctx, branch, base=DEFAULT_BRANCH) -> str:
+def is_agent6(ctx) -> bool:
+    return get_current_branch(ctx).startswith("6.")
+
+
+def get_default_branch():
+    """Returns the default git branch given the current context (agent 6 / 7)."""
+
+    # We create a context to avoid passing context in each function
+    # This context is used to get the current branch so there is no side effect
+    ctx = Context()
+
+    return '6.53.x' if is_agent6(ctx) else 'main'
+
+
+def get_common_ancestor(ctx, branch, base=None) -> str:
+    base = base or get_default_branch()
+
     return ctx.run(f"git merge-base {branch} {base}", hide=True).stdout.strip()
 
 
@@ -132,15 +150,15 @@ def get_main_parent_commit(ctx) -> str:
     """
     Get the commit sha your current branch originated from
     """
-    return ctx.run("git merge-base HEAD origin/main", hide=True).stdout.strip()
+    return ctx.run(f"git merge-base HEAD origin/{get_default_branch()}", hide=True).stdout.strip()
 
 
 def check_base_branch(branch, release_version):
     """
     Checks if the given branch is either the default branch or the release branch associated
     with the given release version.
     """
-    return branch == DEFAULT_BRANCH or branch == release_version.branch()
+    return branch == get_default_branch() or branch == release_version.branch()
 
 
 def try_git_command(ctx, git_command):

tasks/libs/common/gomodules.py

@@ -7,13 +7,13 @@
 import sys
 from collections.abc import Callable
 from dataclasses import dataclass
-from functools import lru_cache
 from pathlib import Path
 from typing import ClassVar
 
 import yaml
 
 import tasks
+from tasks.libs.common.utils import agent_working_directory
 
 
 class ConfigDumper(yaml.SafeDumper):
@@ -305,14 +305,15 @@ def dependency_path(self, agent_version):
 AGENT_MODULE_PATH_PREFIX = "github.com/DataDog/datadog-agent/"
 
 
-@lru_cache
 def get_default_modules(base_dir: Path | None = None) -> dict[str, GoModule]:
     """Load the default modules from the modules.yml file.
 
     Args:
         base_dir: Root directory of the agent repository ('.' by default).
     """
 
+    base_dir = base_dir or agent_working_directory()
+
     return Configuration.from_file(base_dir).modules
 
 

tasks/libs/common/utils.py

@@ -24,8 +24,8 @@
 from invoke.exceptions import Exit
 
 from tasks.libs.common.color import Color, color_message
-from tasks.libs.common.constants import ALLOWED_REPO_ALL_BRANCHES, DEFAULT_BRANCH, REPO_PATH
-from tasks.libs.common.git import get_commit_sha
+from tasks.libs.common.constants import ALLOWED_REPO_ALL_BRANCHES, REPO_PATH
+from tasks.libs.common.git import get_commit_sha, get_default_branch
 from tasks.libs.owners.parsing import search_owners
 from tasks.libs.releasing.version import get_version
 from tasks.libs.types.arch import Arch
@@ -373,6 +373,7 @@ def get_version_ldflags(ctx, major_version='7', install_path=None):
     Compute the version from the git tags, and set the appropriate compiler
     flags
     """
+
     payload_v = get_payload_version()
     commit = get_commit_sha(ctx, short=True)
 
@@ -494,8 +495,8 @@ def environ(env):
 
 
 def is_pr_context(branch, pr_id, test_name):
-    if branch == DEFAULT_BRANCH:
-        print(f"Running on {DEFAULT_BRANCH}, skipping check for {test_name}.")
+    if branch == get_default_branch():
+        print(f"Running on {get_default_branch()}, skipping check for {test_name}.")
         return False
     if not pr_id:
         print(f"PR not found, skipping check for {test_name}.")
@@ -749,3 +750,11 @@ def get_metric_origin(origin_product, origin_sub_product, origin_product_detail,
     if origin_field:
         return {"origin": metric_origin}
     return metric_origin
+
+
+def agent_working_directory():
+    """Returns the working directory for the current context (agent 6 / 7)."""
+
+    from tasks.libs.common.worktree import LOCAL_DIRECTORY, WORKTREE_DIRECTORY, is_worktree
+
+    return WORKTREE_DIRECTORY if is_worktree() else LOCAL_DIRECTORY

tasks/libs/common/worktree.py

@@ -0,0 +1,106 @@
+"""Worktree utilities, used to execute tasks from this local repository (main) to a worktree with a different HEAD (e.g. 6.53.x).
+
+Common environment variables that can be used:
+- WORKTREE_NO_PULL: If set to any value, the worktree will not be pulled before running the command.
+"""
+
+import os
+from contextlib import contextmanager
+from pathlib import Path
+
+from invoke.exceptions import Exit
+
+from tasks.libs.common.color import Color, color_message
+
+WORKTREE_DIRECTORY = Path.cwd().parent / "datadog-agent-worktree"
+LOCAL_DIRECTORY = Path.cwd().resolve()
+
+
+def init_env(ctx, branch: str | None = None):
+    """Will prepare the environment for commands applying to a worktree.
+
+    To be used before each worktree section.
+    Will:
+    1. Add the agent worktree if not present.
+    2. Fetch the latest changes from the agent worktree.
+    """
+
+    if not WORKTREE_DIRECTORY.is_dir():
+        if not ctx.run(f"git worktree add '{WORKTREE_DIRECTORY}' origin/{branch or 'main'}", warn=True):
+            raise Exit(
+                f'{color_message("Error", Color.RED)}: Cannot initialize worktree environment. You might want to reset the worktree directory with `inv worktree.remove`',
+                code=1,
+            )
+
+    if branch:
+        worktree_branch = ctx.run(
+            f"git -C '{WORKTREE_DIRECTORY}' rev-parse --abbrev-ref HEAD", hide=True
+        ).stdout.strip()
+        if worktree_branch != branch:
+            ctx.run(f"git -C '{WORKTREE_DIRECTORY}' checkout '{branch}'", hide=True)
+
+        if not os.environ.get("AGENT_WORKTREE_NO_PULL"):
+            ctx.run(f"git -C '{WORKTREE_DIRECTORY}' pull", hide=True)
+
+
+def remove_env(ctx):
+    """Will remove the environment for commands applying to a worktree."""
+
+    ctx.run(f"git worktree remove -f '{WORKTREE_DIRECTORY}'", warn=True)
+
+
+def is_worktree():
+    """Will return True if the current environment is a worktree environment."""
+
+    return Path.cwd() != LOCAL_DIRECTORY
+
+
+def enter_env(ctx, branch: str, no_checkout=False):
+    """Enters the worktree environment."""
+
+    if not no_checkout:
+        init_env(ctx, branch)
+    else:
+        assert WORKTREE_DIRECTORY.is_dir(), "Worktree directory is not present and no_switch is set to True"
+
+    os.chdir(WORKTREE_DIRECTORY)
+
+
+def exit_env():
+    """Exits the worktree environment."""
+
+    os.chdir(LOCAL_DIRECTORY)
+
+
+@contextmanager
+def agent_context(ctx, branch: str | None, no_checkout=False):
+    """Applies code to the worktree environment if the version is not None.
+
+    Args:
+        branch: The branch to switch to.
+        no_checkout: If True, the branch will not be switched (no pull will be performed too).
+
+    Usage:
+        > with agent_context(ctx, branch):
+        >    ctx.run("head CHANGELOG.rst")  # Displays the changelog of the target branch
+    """
+
+    if branch is not None:
+        # Do not stack two environments
+        if is_worktree():
+            yield
+            return
+
+        try:
+            # Enter
+            enter_env(ctx, branch, no_checkout=no_checkout)
+
+            yield
+        finally:
+            # Exit
+            exit_env()
+    else:
+        # NOTE: This ensures that we don't push local context from a worktree context (context might be switched within inner functions)
+        assert not is_worktree(), 'Local context cannot be used within a worktree context'
+
+        yield