bitfab 0.65.0 → 0.65.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: aa27d322488c12e530f14a5f85a7f9978208edf332ab91f7bf156f4ccb7b0b25
-  data.tar.gz: f8cadf7e8693aa155c254b7866b41f1781d5de4493dad8b0aad6643d285469a6
+  metadata.gz: 1ccbb4e0e912ec35fe388c4cdb2225c4fa3c4f5e020e87c7d52f81a3358907da
+  data.tar.gz: 4e674dc9c1fb2e1c2ef6f5d5b1cf21146904a0d85989d7e41f25ce6c0b04c740
 SHA512:
-  metadata.gz: b274fec88de7c05f84977a60cc146290e89e2e8778fec9a498967304cc2c46a7cb65efa27f65210ec76ba1a2b58dbd3738e4f969d6ac126a9c3a87b25f80c30c
-  data.tar.gz: 8ddf0c559dc1dd33139497e18acf48a065017a59ea05b81e708ee9d3aa4a8eae1b202fb1b5243e7027131b27bcaedd5a2c3065df2fa964793c388609f199e3ad
+  metadata.gz: a5f9344e7be229c578f2fb875de96aa11f49487b24775eec15955f32d18ab4552067842af6b4a8f12041f6140ee21bafbd3983630d0929e00a35236fb427e95f
+  data.tar.gz: 15bd45c45d9f6700edd7b96bbc792236441581763097307d13d4a14a76d0d8d733e5444c1586c9a24f6c990a8575e70851c37a80ff10b4f153bb84bb554e00c2

data/lib/bitfab/cloudReplay.py

@@ -116,6 +116,8 @@ PUSH_TRIGGER = re.compile(
     r"^[ \t]*\"?(?:on\"?[ \t]*:.*\bpush\b|push\"?[ \t]*:)", re.MULTILINE
 )
 DISPATCH_INPUTS = ("execution_id", "request", KEY_INPUT)
+PUSH_SUBJECT = "Bitfab replay started on push"
+PUSH_CONDITION = f"github.event_name != 'push' || !startsWith(github.ref_name, '{PREFIX}') || startsWith(github.event.head_commit.message, '{PUSH_SUBJECT}')"
 BRANCH_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._/-]*")
 ENV_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
 RESERVED_ENV_PREFIXES = ("GITHUB_", "RUNNER_", "ACTIONS_", "BITFAB_REPLAY_")
@@ -254,7 +256,8 @@ Set up once:
                [--runs-on LABEL] [--secret-prefix PREFIX] [--check COMMAND]
       Writes .github/workflows/bitfab-replay.yml, the only file cloud replay keeps
       in the repository, and lists what is left to do. On an existing workflow it
-      adds only the lines a newer SDK needs and lists what it changed.
+      adds only the lines a newer SDK needs, such as the push trigger that lets a
+      replay run without GitHub access, and lists what it changed.
   --cloud-update [same options] [--remove-old-key-secrets]
       The same command as --cloud-init, for bringing an existing setup up to date.
       It also lists BITFAB_REPLAY_KEY_* secrets an older SDK left in the
@@ -266,6 +269,16 @@ Set up once:
 GitHub access comes from the GitHub CLI when it is logged in, otherwise from GH_TOKEN
 or GITHUB_TOKEN, otherwise from the github.com credential git already stores. Without
 the GitHub CLI, --cloud-secrets lists the secrets to create by hand and where.
+
+Without GitHub access, or with access that cannot read or push to the repository, a
+replay still runs when git can push the snapshot branch and the workflow also starts
+on a push to a bitfab-replay/ branch, as one written or updated by --cloud-init does.
+This command then warns, prints where to watch the run, and returns once the branch
+is pushed, with no updates and exit code 0. The run uses the workflow's BITFAB_API_KEY
+secret, and --cloud-full-output and --cloud-allow-file are refused, since only GitHub
+access can confirm the repository is private. Delete the snapshot branch once the run
+finishes. With access, a replay starts through the GitHub API as before, and its
+snapshot commit asks GitHub not to start any workflow on the push.
 """
 
 
@@ -281,6 +294,10 @@ class DispatchRejected(ValueError):
         self.record = record
 
 
+class GithubUnreachable(ValueError):
+    pass
+
+
 class KeepRedirect(urllib.request.HTTPRedirectHandler):
     def redirect_request(self, *args, **kwargs):
         return None
@@ -939,7 +956,15 @@ def sensitive(path):
     )
 
 
-def snapshot(root, workflow, files, request, *, preview=False, allowed=()):
+def snapshot_subject(execution_id, on_push):
+    if on_push:
+        return f"{PUSH_SUBJECT} {execution_id}"
+    return f"Bitfab replay {execution_id} [skip actions]"
+
+
+def snapshot(
+    root, workflow, files, request, *, preview=False, allowed=(), on_push=False
+):
     if git(root, "ls-files", "-u"):
         raise ValueError("Resolve merge conflicts before snapshotting")
     head = git(root, "rev-parse", "HEAD")
@@ -1014,7 +1039,7 @@ def snapshot(root, workflow, files, request, *, preview=False, allowed=()):
             tree,
             "-p",
             head,
-            input=f"Bitfab replay {request['id']}\n\n{json.dumps(request)}\n",
+            input=f"{snapshot_subject(request['id'], on_push)}\n\n{json.dumps(request)}\n",
         )
         return {**summary, "sha": sha}
 
@@ -1050,9 +1075,8 @@ def execution_lock(directory, execution_id):
 def find_run(record):
     if record.get("snapshot") is False:
         return find_branch_run(record)
-    query = urlencode(
-        {"event": "workflow_dispatch", "branch": record["branch"], "per_page": 100}
-    )
+    event = "push" if record.get("startedOnPush") else "workflow_dispatch"
+    query = urlencode({"event": event, "branch": record["branch"], "per_page": 100})
     runs = api(
         record["repository"], f"actions/workflows/{record['workflow']}/runs?{query}"
     )
@@ -1274,15 +1298,19 @@ def http_status(error):
 
 
 def preflight(repo, workflow, *, pushes=True, full_output=False, allows_files=False):
-    access = github_access()
+    try:
+        access = github_access()
+    except ValueError as error:
+        raise GithubUnreachable(str(error)) from error
     try:
         details = api(repo, "")
     except RuntimeError as error:
-        raise ValueError(
+        refused = getattr(error, "http_status", None) in (401, 403, 404)
+        raise (GithubUnreachable if refused else ValueError)(
             f"GitHub access from {access['source']} cannot read {repo}{http_status(error)}; use an account or token with write access to it"
         ) from error
     if pushes and (details or {}).get("permissions", {}).get("push") is False:
-        raise ValueError(
+        raise GithubUnreachable(
             f"GitHub access from {access['source']} cannot push to {repo}; use an account or token with write access to it"
         )
     if full_output and (details or {}).get("private") is False:
@@ -1344,6 +1372,12 @@ def submit(root, repo, workflow, parsed, path):
         record = json.loads(path.read_text())
         if record["request"] != request or record["repository"].lower() != repo.lower():
             raise ValueError("Execution ID already belongs to a different request")
+        if record.get("startedOnPush") and record["state"] != "prepared":
+            try:
+                github_access()
+            except ValueError as error:
+                record["unfollowed"] = str(error)
+                return record
         status(record, fetch_result=False)
         if record["state"] == "prepared":
             raise ValueError(
@@ -1351,17 +1385,24 @@ def submit(root, repo, workflow, parsed, path):
             )
         refuse_rejected(record)
         return record
-    preflight(
-        repo,
-        workflow,
-        full_output=request.get("fullOutput") is True,
-        allows_files=bool(allowed),
-    )
+    unfollowed = None
+    try:
+        preflight(
+            repo,
+            workflow,
+            full_output=request.get("fullOutput") is True,
+            allows_files=bool(allowed),
+        )
+    except GithubUnreachable as error:
+        unfollowed = str(error)
+        refuse_unfollowed(root, repo, workflow, request, allowed, unfollowed)
     if dry_run(request["args"]):
         progress(
             "--dry-run dispatches a GitHub run that checks secrets and resolves the traces without replaying; it receives the workflow's secrets"
         )
-    source = snapshot(root, workflow, files, request, allowed=allowed)
+    source = snapshot(
+        root, workflow, files, request, allowed=allowed, on_push=bool(unfollowed)
+    )
     describe_snapshot(source)
     record = {
         "id": execution_id,
@@ -1370,6 +1411,7 @@ def submit(root, repo, workflow, parsed, path):
         "branch": PREFIX + execution_id,
         "request": request,
         "state": "prepared",
+        **({"startedOnPush": True} if unfollowed else {}),
         **source,
     }
     save(path, record)
@@ -1389,6 +1431,9 @@ def submit(root, repo, workflow, parsed, path):
         ) from error
     record["state"] = "dispatch_unknown"
     save(path, record)
+    if unfollowed:
+        record["unfollowed"] = unfollowed
+        return record
     text = within(root, f"{WORKFLOW_DIRECTORY}/{workflow}").read_text()
     declared = declared_inputs(text)
     key_inputs = issue_key_token(root, record, text)
@@ -1400,6 +1445,62 @@ def submit(root, repo, workflow, parsed, path):
     )
 
 
+def push_trigger(lines):
+    for index, line in enumerate(lines):
+        if re.fullmatch(r'"?on"?[ \t]*:[ \t]*(#.*)?', line):
+            end = block_end(lines, index, 0)
+            step = child_indent(lines, index, end, "  ")
+            row = next(
+                (
+                    row
+                    for row in range(index + 1, end)
+                    if re.fullmatch(rf'{step}"?push"?[ \t]*:.*', lines[row])
+                ),
+                None,
+            )
+            if row is None:
+                return None
+            return lines[row : block_end(lines, row, len(step))]
+        if re.match(r'"?on"?[ \t]*:', line):
+            return [line] if re.search(r"\bpush\b", line) else None
+    return None
+
+
+def starts_on_push(text):
+    trigger = push_trigger(text.splitlines()) or []
+    return any(PREFIX in line for line in trigger) and not any(
+        "branches-ignore" in line for line in trigger
+    )
+
+
+def refuse_unfollowed(root, repo, workflow, request, allowed, reason):
+    path = f"{WORKFLOW_DIRECTORY}/{workflow}"
+    if not starts_on_push(within(root, path).read_text()):
+        raise ValueError(
+            f"{reason}. Without it a replay can still run, with no updates here, once {path} also starts on a push to a {PREFIX} branch: run bitfab-replay --cloud-update and commit the workflow it updates"
+        )
+    if request.get("fullOutput") is True:
+        raise ValueError(
+            f"--cloud-full-output writes every trace's input and output into the Actions job log, so it needs GitHub access that confirms {repo} is private: {reason}"
+        )
+    if allowed:
+        raise ValueError(
+            f"{ALLOW_FLAG} uploads a credential-like file to a snapshot branch, so it needs GitHub access that confirms {repo} is private: {reason}"
+        )
+
+
+def report_unfollowed(record, reason):
+    progress(
+        f"Warning: the replay started on GitHub, but this command gets no updates on it: {reason}"
+    )
+    progress(
+        f"It runs with the workflow's BITFAB_API_KEY secret, since replaying as you needs GitHub access. Watch it at https://github.com/{record['repository']}/actions/workflows/{record['workflow']}; its experiment appears in Bitfab when it finishes"
+    )
+    progress(
+        f"Once it finishes, delete its snapshot branch with git push origin --delete {record['branch']}, or with --cloud-cleanup {record['id']} from a machine with GitHub access"
+    )
+
+
 def local_client(root, request):
     try:
         output = command(
@@ -1736,6 +1837,11 @@ def run_cli(argv):
                 raise
             if record.get("preview"):
                 return record, None
+            reason = record.pop("unfollowed", None)
+            if reason:
+                report_unfollowed(record, reason)
+                save(path, record)
+                return record, None
         else:
             if not path.exists():
                 raise ValueError(
@@ -2791,7 +2897,7 @@ def push_triggers(root, workflow):
 
 
 def workflow_document(steps, run, directory, env, *, runs_on, environment):
-    job = {"runs-on": runs_on}
+    job = {"if": PUSH_CONDITION, "runs-on": runs_on}
     if environment:
         job["environment"] = environment
     job["steps"] = [
@@ -2816,7 +2922,8 @@ def workflow_document(steps, run, directory, env, *, runs_on, environment):
                     "request": {"type": "string", "required": False},
                     KEY_INPUT: {"type": "string", "required": False},
                 }
-            }
+            },
+            "push": {"branches": [PREFIX + "**"]},
         },
         "permissions": {"contents": "read", "id-token": "write"},
         "jobs": {"replay": job},
@@ -2830,10 +2937,18 @@ def indentation(line):
 def block_end(lines, start, indent):
     end = start + 1
     while end < len(lines) and (
-        not lines[end].strip() or indentation(lines[end]) > indent
+        not lines[end].strip()
+        or lines[end].lstrip().startswith("#")
+        or indentation(lines[end]) > indent
     ):
         end += 1
-    while end > start + 1 and not lines[end - 1].strip():
+    while end > start + 1 and (
+        not lines[end - 1].strip()
+        or (
+            lines[end - 1].lstrip().startswith("#")
+            and indentation(lines[end - 1]) <= indent
+        )
+    ):
         end -= 1
     return end
 
@@ -2938,19 +3053,67 @@ def add_id_token_permission(lines):
         if not re.fullmatch(r'"?permissions"?[ \t]*:[ \t]*(#.*)?', line):
             continue
         end = block_end(lines, index, 0)
-        child = next(
-            (
-                lines[row][: indentation(lines[row])]
-                for row in range(index + 1, end)
-                if lines[row].strip()
-            ),
-            "  ",
-        )
-        lines.insert(end, f"{child}id-token: write")
+        lines.insert(end, f"{child_indent(lines, index, end, '  ')}id-token: write")
         return True
     return False
 
 
+def holds_yaml(line):
+    return bool(line.strip()) and not line.lstrip().startswith("#")
+
+
+def child_indent(lines, index, end, fallback):
+    return next(
+        (
+            lines[row][: indentation(lines[row])]
+            for row in range(index + 1, end)
+            if holds_yaml(lines[row])
+        ),
+        fallback,
+    )
+
+
+def add_push_trigger(lines):
+    for index, line in enumerate(lines):
+        if not re.fullmatch(r'"?on"?[ \t]*:[ \t]*(#.*)?', line):
+            continue
+        end = block_end(lines, index, 0)
+        step = child_indent(lines, index, end, "  ")
+        lines[end:end] = [
+            f"{step}push:",
+            f"{step * 2}branches:",
+            f"{step * 3}- {PREFIX}**",
+        ]
+        return True
+    return False
+
+
+def add_push_condition(lines):
+    jobs = next(
+        (
+            index
+            for index, line in enumerate(lines)
+            if re.fullmatch(r'"?jobs"?[ \t]*:[ \t]*(#.*)?', line)
+        ),
+        None,
+    )
+    if jobs is None:
+        return False
+    end = block_end(lines, jobs, 0)
+    job = next((row for row in range(jobs + 1, end) if holds_yaml(lines[row])), None)
+    if job is None or not lines[job].rstrip().endswith(":"):
+        return False
+    job_end = block_end(lines, job, indentation(lines[job]))
+    step = child_indent(lines, job, job_end, " " * (indentation(lines[job]) + 2))
+    if any(
+        re.fullmatch(rf'{step}"?if"?[ \t]*:.*', lines[row])
+        for row in range(job + 1, job_end)
+    ):
+        return False
+    lines.insert(job + 1, f"{step}if: {json.dumps(PUSH_CONDITION)}")
+    return True
+
+
 def secret_scope(environment):
     return f"environments/{environment}/" if environment else "actions/"
 
@@ -3041,6 +3204,20 @@ def update_workflow(root, path, *, remove_old_key_secrets=False):
             by_hand.append(
                 f"{name} (type: string, required: false) under on.workflow_dispatch.inputs"
             )
+    if not starts_on_push(text):
+        if push_trigger(lines) is None and add_push_trigger(lines):
+            updated.append(
+                f"Started the workflow on a push to a {PREFIX} branch too, so a replay runs without GitHub access, with no updates on it"
+            )
+        else:
+            by_hand.append(f"{PREFIX}** in on.push.branches")
+    if push_trigger(lines) is not None and PUSH_SUBJECT not in text:
+        if add_push_condition(lines):
+            updated.append(
+                "Skipped the push run of a replay started through the GitHub API, so each replay runs once"
+            )
+        else:
+            by_hand.append(f"if: {json.dumps(PUSH_CONDITION)} on the replay job")
     if not re.search(r'^"?run-name"?[ \t]*:', text, re.MULTILINE):
         add_run_name(lines)
         updated.append("Named each run after its execution")
@@ -3095,7 +3272,7 @@ def update_workflow(root, path, *, remove_old_key_secrets=False):
 def initialize(argv, name="--cloud-init"):
     parser = argparse.ArgumentParser(
         prog=f"bitfab-replay {name}",
-        description="Write .github/workflows/bitfab-replay.yml, the only file cloud replay keeps in the repository, and list what is left to do. It detects the runtime, package manager, install command, and replay command from the directory you run it in. On an existing workflow it only adds what a newer SDK needs (the execution_id and key_id inputs, the run name, and the id-token: write permission), removes the BITFAB_REPLAY_BITFAB_API_KEY line older SDKs wrote, lists what it changed, and lists the BITFAB_REPLAY_KEY_* secrets older SDKs left in the repository; edit it directly for anything else.",
+        description="Write .github/workflows/bitfab-replay.yml, the only file cloud replay keeps in the repository, and list what is left to do. It detects the runtime, package manager, install command, and replay command from the directory you run it in. On an existing workflow it only adds what a newer SDK needs (the execution_id and key_id inputs, the push trigger and the condition that skips a duplicate push run, the run name, and the id-token: write permission), removes the BITFAB_REPLAY_BITFAB_API_KEY line older SDKs wrote, lists what it changed, and lists the BITFAB_REPLAY_KEY_* secrets older SDKs left in the repository; edit it directly for anything else.",
     )
     parser.add_argument(
         "--secret",

data/lib/bitfab/datasets.rb

@@ -17,6 +17,9 @@ module Bitfab
   # trace function. Experiments replay against it and its graders score its
   # members. Every method returns the parsed JSON response as a Hash with
   # string keys matching the HTTP API ("traceCount", "addedTraceIds", ...).
+  # Every dataset summary carries "golden", true when the dataset is golden
+  # for the whole organization and listed first in the Bitfab app, and false
+  # otherwise, including when an older server leaves the field out.
   class Datasets
     DEFAULT_RERUN_TIMEOUT_SECONDS = 90
     RERUN_POLL_INTERVAL_SECONDS = 1.0
@@ -29,11 +32,14 @@ module Bitfab
 
     # Create a dataset, or update the one already named this way under the
     # same trace function. The result's "created" reports which happened. A
-    # nil description leaves an existing description untouched.
-    def save(trace_function_key:, name:, description: nil)
+    # nil description leaves an existing description untouched. golden: true
+    # marks the dataset as golden for the whole organization so the Bitfab app
+    # lists it first, false unmarks it, and nil leaves it as it is.
+    def save(trace_function_key:, name:, description: nil, golden: nil)
       payload = {"traceFunctionKey" => trace_function_key, "name" => name}
       payload["description"] = description unless description.nil?
-      @http_client.request("/api/sdk/datasets", payload)
+      payload["golden"] = golden unless golden.nil?
+      with_golden_default(@http_client.request("/api/sdk/datasets", payload))
     end
 
     # List datasets, scoped to one trace function when given and
@@ -44,7 +50,7 @@ module Bitfab
       datasets = []
       loop do
         result = @http_client.get("/api/sdk/datasets?#{URI.encode_www_form(query)}")
-        datasets.concat(result["datasets"])
+        datasets.concat(result["datasets"].map { |summary| summary_with_golden_default(summary) })
         cursor = result["nextCursor"]
         return datasets if cursor.nil?
 
@@ -55,22 +61,24 @@ module Bitfab
     # Fetch one dataset by id. A dataset outside this organization raises
     # Net::HTTPError with a 404.
     def get(dataset_id)
-      @http_client.get(dataset_path(dataset_id))["dataset"]
+      summary_with_golden_default(@http_client.get(dataset_path(dataset_id))["dataset"])
     end
 
     # Rename the dataset or change its description, and return the updated
     # dataset. It keeps its id, traces, and graders. A nil argument leaves that
-    # field untouched, and an empty description clears it. Raises ArgumentError
-    # when both are nil. A blank name, or one another dataset under the same
+    # field untouched, and an empty description clears it. golden: true marks
+    # the dataset as golden for the whole organization and false unmarks it.
+    # Raises ArgumentError when name, description, and golden are all nil. A blank name, or one another dataset under the same
     # trace function already uses, raises Net::HTTPError with a 400. A dataset
     # outside this organization raises Net::HTTPError with a 404.
-    def update(dataset_id, name: nil, description: nil)
-      raise ArgumentError, "Provide name and/or description to update." if name.nil? && description.nil?
+    def update(dataset_id, name: nil, description: nil, golden: nil)
+      raise ArgumentError, "Provide name, description, and/or golden to update." if name.nil? && description.nil? && golden.nil?
 
       payload = {}
       payload["name"] = name unless name.nil?
       payload["description"] = description unless description.nil?
-      @http_client.request(dataset_path(dataset_id), payload, method: "PATCH")["dataset"]
+      payload["golden"] = golden unless golden.nil?
+      summary_with_golden_default(@http_client.request(dataset_path(dataset_id), payload, method: "PATCH")["dataset"])
     end
 
     # The ids of every trace in the dataset, the same membership a replay with
@@ -92,24 +100,24 @@ module Bitfab
     # organization or under another trace function are reported in
     # "skippedTraceIds" rather than failing the call.
     def add_traces(dataset_id, trace_ids)
-      @http_client.request(dataset_path(dataset_id, "/traces"), {"traceIds" => trace_ids})
+      with_golden_default(@http_client.request(dataset_path(dataset_id, "/traces"), {"traceIds" => trace_ids}))
     end
 
     # Remove traces from the dataset. The traces themselves are never deleted.
     def remove_traces(dataset_id, trace_ids)
-      @http_client.request(dataset_path(dataset_id, "/removeTraces"), {"traceIds" => trace_ids})
+      with_golden_default(@http_client.request(dataset_path(dataset_id, "/removeTraces"), {"traceIds" => trace_ids}))
     end
 
     # Assign graders to the dataset (1 to 100 ids per call). Graders outside
     # the organization or under another trace function are reported in
     # "skippedGraderIds" rather than failing the call.
     def add_graders(dataset_id, grader_ids)
-      @http_client.request(dataset_path(dataset_id, "/graders"), {"graderIds" => grader_ids})
+      with_golden_default(@http_client.request(dataset_path(dataset_id, "/graders"), {"graderIds" => grader_ids}))
     end
 
     # Unassign graders from the dataset.
     def remove_graders(dataset_id, grader_ids)
-      @http_client.request(dataset_path(dataset_id, "/removeGraders"), {"graderIds" => grader_ids})
+      with_golden_default(@http_client.request(dataset_path(dataset_id, "/removeGraders"), {"graderIds" => grader_ids}))
     end
 
     # Re-run graders over every trace in the dataset. Defaults to every
@@ -140,6 +148,18 @@ module Bitfab
 
     private
 
+    def with_golden_default(result)
+      return result unless result["dataset"].is_a?(Hash)
+
+      result.merge("dataset" => summary_with_golden_default(result["dataset"]))
+    end
+
+    def summary_with_golden_default(summary)
+      return summary unless summary.is_a?(Hash)
+
+      {"golden" => false}.merge(summary)
+    end
+
     def settle_grader_rerun(dataset_id, run, timeout)
       case run["status"]
       when "completed"

data/lib/bitfab/subtree.rb

@@ -796,6 +796,8 @@ module Bitfab
       end
 
       def skip_boundary_call?(event)
+        return false if sdk_frame?(event.path)
+
         boundary = @stacks[Fiber.current.object_id]&.last
         if boundary.is_a?(Boundary) && boundary.skip_next_call &&
             boundary.receiver.equal?(event.self) && boundary.method_name == called_method_id(event)
@@ -819,6 +821,13 @@ module Bitfab
         true
       end
 
+      def sdk_frame?(path)
+        @sdk_frames ||= {}
+        @sdk_frames.fetch(path) do
+          @sdk_frames[path] = Subtree.realpath(path)&.start_with?(SDK_PREFIX) || false
+        end
+      end
+
       def passthrough_wrapper?(parameters)
         kinds = parameters.map(&:first)
         kinds.any? { |kind| kind == :rest || kind == :keyrest } &&

data/lib/bitfab/traceable.rb

@@ -108,7 +108,7 @@ module Bitfab
               client: target_client, trace_function_key:, config: capture_config,
               max_depth:, max_captured_subtree_spans:, include_wrappers:, mock_on_replay_default:,
               root_receiver: self, root_method_name: method_name
-            ) { super(*args, **kwargs, &block) }
+            ) { Traceable.send(:as_root_call, self, method_name) { super(*args, **kwargs, &block) } }
           end
           call = Subtree::ConfiguredCall.new(receiver: self, method_name:, name: span_name,
             type:, args:, kwargs:, nested_trace_key: trace_function_key)
@@ -122,12 +122,14 @@ module Bitfab
               client: target_client, trace_function_key:, config: capture_config,
               max_depth:, max_captured_subtree_spans:, include_wrappers:, mock_on_replay_default:,
               root_receiver: self, root_method_name: method_name
-            ) { super(*args, **kwargs, &block) }
+            ) { Traceable.send(:as_root_call, self, method_name) { super(*args, **kwargs, &block) } }
           end
         end
       end
       wrapper.instance_variable_set(:@bitfab_span_method, method_name_str)
       wrapper.instance_variable_set(:@bitfab_span_key, trace_function_key)
+      wrapper.instance_variable_set(:@bitfab_trace_method, method_name_str)
+      wrapper.instance_variable_set(:@bitfab_trace_client, bound_client)
       preserve_visibility(klass, wrapper, method_name)
       klass.prepend(wrapper)
     end
@@ -142,9 +144,12 @@ module Bitfab
       bound_client = client
       wrapper = Module.new do
         define_method(method_name) do |*args, **kwargs, &block|
+          next super(*args, **kwargs, &block) if Traceable.send(:own_span_call?, self, method_name)
+
           Subtree.assert_surface!(:trace)
           session = Subtree.current_session
-          if !session || (bound_client && !session.client.equal?(bound_client))
+          if !session || (bound_client && !session.client.equal?(bound_client)) ||
+              Traceable.send(:root_handles_call?, self, method_name, wrapper)
             next super(*args, **kwargs, &block)
           end
           unless capture
@@ -176,6 +181,49 @@ module Bitfab
       klass.prepend(wrapper)
     end
 
+    ROOT_CALL_KEY = :bitfab_root_call
+
+    def self.as_root_call(receiver, method_name)
+      previous = Thread.current[ROOT_CALL_KEY]
+      Thread.current[ROOT_CALL_KEY] = [receiver, method_name.to_sym]
+      yield
+    ensure
+      Thread.current[ROOT_CALL_KEY] = previous
+    end
+    private_class_method :as_root_call
+
+    def self.root_handles_call?(receiver, method_name, node_wrapper)
+      marker = Thread.current[ROOT_CALL_KEY]
+      if marker && marker[0].equal?(receiver) && marker[1] == method_name.to_sym
+        Thread.current[ROOT_CALL_KEY] = nil
+        return true
+      end
+
+      trace_wrapper = inner_trace_wrapper(receiver, method_name, node_wrapper)
+      return false unless trace_wrapper
+
+      client = trace_wrapper.instance_variable_get(:@bitfab_trace_client) || Bitfab.client_or_nil
+      (client&.send(:should_record?) && !Subtree.nested_root_suppressed?) ? true : false
+    end
+    private_class_method :root_handles_call?
+
+    def self.own_span_call?(receiver, method_name)
+      span = SpanContext.current
+      return false unless span && span[:surface] == :span
+
+      span[:explicit_span_receiver].equal?(receiver) && span[:explicit_span_method_name] == method_name
+    end
+    private_class_method :own_span_call?
+
+    def self.inner_trace_wrapper(receiver, method_name, node_wrapper)
+      ancestors = receiver.is_a?(Module) ? receiver.singleton_class.ancestors : receiver.class.ancestors
+      name = method_name.to_s
+      ancestors.drop_while { |mod| !mod.equal?(node_wrapper) }.find do |mod|
+        mod.instance_variable_defined?(:@bitfab_trace_method) && mod.instance_variable_get(:@bitfab_trace_method) == name
+      end
+    end
+    private_class_method :inner_trace_wrapper
+
     def self.source_of(klass, method_name)
       path, line = klass.instance_method(method_name).source_location || [nil, nil]
       [path, line, method_name]
@@ -345,15 +393,14 @@ module Bitfab
 
       def method_added(method_name)
         super
+        if defined?(@_bitfab_pending_nodes) && @_bitfab_pending_nodes&.key?(method_name)
+          _bitfab_node_method(method_name, **@_bitfab_pending_nodes.delete(method_name))
+        end
         if defined?(@_bitfab_pending_spans) && @_bitfab_pending_spans&.key?(method_name)
-          config = @_bitfab_pending_spans.delete(method_name)
-          _bitfab_wrap_method(method_name, **config)
-        elsif defined?(@_bitfab_pending_traces) && @_bitfab_pending_traces&.key?(method_name)
-          config = @_bitfab_pending_traces.delete(method_name)
-          _bitfab_trace_method(method_name, **config)
-        elsif defined?(@_bitfab_pending_nodes) && @_bitfab_pending_nodes&.key?(method_name)
-          config = @_bitfab_pending_nodes.delete(method_name)
-          _bitfab_node_method(method_name, **config)
+          _bitfab_wrap_method(method_name, **@_bitfab_pending_spans.delete(method_name))
+        end
+        if defined?(@_bitfab_pending_traces) && @_bitfab_pending_traces&.key?(method_name)
+          _bitfab_trace_method(method_name, **@_bitfab_pending_traces.delete(method_name))
         end
       end
 

data/lib/bitfab/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Bitfab
-  VERSION = "0.65.0"
+  VERSION = "0.65.3"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.65.0
+  version: 0.65.3
 platform: ruby
 authors:
 - Harvest Team