[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,15 +4,19 @@
 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':
-        print(color_message("You're about to commit or push to the main, are you sure this is what you want?", "red"))
+    if local_branch == get_default_branch():
+        print(
+            color_message(
+                f"You're about to commit or push to {get_default_branch()}, are you sure this is what you want?", "red"
+            )
+        )
         raise Exit(code=1)
 
     if re.fullmatch(r'^[0-9]+\.[0-9]+\.x$', local_branch):

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/gotest.py

@@ -514,12 +514,13 @@ def get_modified_packages(ctx, build_tags=None, lint=False) -> list[GoModule]:
             modules_to_test[best_module_path] = GoModule(best_module_path, test_targets=[relative_target])
 
     # Clean up duplicated paths to reduce Go test cmd length
+    default_modules = get_default_modules()
     for module in modules_to_test:
         modules_to_test[module].test_targets = clean_nested_paths(modules_to_test[module].test_targets)
         if (
             len(modules_to_test[module].test_targets) >= WINDOWS_MAX_PACKAGES_NUMBER
         ):  # With more packages we can reach the limit of the command line length on Windows
-            modules_to_test[module].test_targets = get_default_modules()[module].test_targets
+            modules_to_test[module].test_targets = default_modules[module].test_targets
 
     print("Running tests for the following modules:")
     for module in modules_to_test:
@@ -752,16 +753,17 @@ def format_packages(ctx: Context, impacted_packages: set[str], build_tags: list[
     packages = [f'{package.replace("github.com/DataDog/datadog-agent/", "./")}' for package in impacted_packages]
     modules_to_test = {}
 
+    default_modules = get_default_modules()
     for package in packages:
         module_path = get_go_module(package)
 
         # Check if the module is in the target list of the modules we want to test
-        if module_path not in get_default_modules() or not get_default_modules()[module_path].should_test():
+        if module_path not in default_modules or not default_modules[module_path].should_test():
             continue
 
         # Check if the package is in the target list of the module we want to test
         targeted = False
-        for target in get_default_modules()[module_path].test_targets:
+        for target in default_modules[module_path].test_targets:
             if normpath(os.path.join(module_path, target)) in package:
                 targeted = True
                 break
@@ -784,12 +786,13 @@ def format_packages(ctx: Context, impacted_packages: set[str], build_tags: list[
             modules_to_test[module_path] = GoModule(module_path, test_targets=[relative_target])
 
     # Clean up duplicated paths to reduce Go test cmd length
+    default_modules = get_default_modules()
     for module in modules_to_test:
         modules_to_test[module].test_targets = clean_nested_paths(modules_to_test[module].test_targets)
         if (
             len(modules_to_test[module].test_targets) >= WINDOWS_MAX_PACKAGES_NUMBER
         ):  # With more packages we can reach the limit of the command line length on Windows
-            modules_to_test[module].test_targets = get_default_modules()[module].test_targets
+            modules_to_test[module].test_targets = default_modules[module].test_targets
 
     module_to_remove = []
     # Clean up to avoid running tests on package with no Go files matching build tags

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