perf(track): bulk metadata upserts + drain pacing + single-pass scrub

fleet/track/api.py

@@ -34,6 +34,7 @@
 
 DEFAULT_TIMEOUT = 30.0
 SERVER_UPLOAD_URL_BATCH_CAP = 100  # /v1/track/upload-urls returns 400 above this.
+SERVER_BULK_UPSERT_BATCH_CAP = 100  # /v1/track/sessions/bulk; chunk to match.
 
 
 AuthInfo = Union[str, Tuple[str, str]]
@@ -56,6 +57,22 @@ class TrackTextMatch:
     negate: bool = False
 
 
+@dataclass(frozen=True)
+class BulkSessionUpsert:
+    """One item in a /v1/track/sessions/bulk request.
+
+    Mirrors the per-arg shape of `upsert_session`. `content_codec`,
+    `raw_bytes`, `stored_bytes` are only sent when this row carries
+    a fresh upload (i.e. include_content_metadata=True equivalent).
+    """
+
+    path: str
+    session: Any
+    content_codec: Optional[str] = None
+    raw_bytes: Optional[int] = None
+    stored_bytes: Optional[int] = None
+
+
 @dataclass(frozen=True)
 class TrackSessionSearchRequest:
     """Structured body for `POST /v1/track/sessions/search`."""
@@ -84,6 +101,26 @@ def __init__(
         self.detail = detail
 
 
+class BulkUpsertPartialFailure(TrackAPIError):
+    """Raised when one chunk of `upsert_sessions_bulk` fails after earlier
+    chunks succeeded. `unsent_items` is the failing chunk plus everything
+    after it; everything before it is already committed server-side. Lets
+    the caller re-enqueue only what hasn't been sent, avoiding repeat
+    sends of already-committed rows on the next flush.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        unsent_items: list["BulkSessionUpsert"],
+        status_code: Optional[int] = None,
+        detail: Any = None,
+    ) -> None:
+        super().__init__(message, status_code=status_code, detail=detail)
+        self.unsent_items = unsent_items
+
+
 def _default_auth_provider() -> Optional[AuthInfo]:
     """Production auth: prefer stored `flt login` creds, then FLEET_API_KEY.
 
@@ -223,6 +260,56 @@ def upsert_session(
         )
         _raise(resp)
 
+    def upsert_sessions_bulk(
+        self,
+        *,
+        device_id: str,
+        items: list["BulkSessionUpsert"],
+    ) -> None:
+        """Bulk-register metadata for many sessions in one request.
+
+        Server reuses the single-row upsert translation logic per item, so
+        path → s3_key conversion and validation behave identically. Empty
+        list is a no-op. Chunks at SERVER_BULK_UPSERT_BATCH_CAP to match
+        the server cap.
+
+        On partial failure (some chunk succeeds, a later one doesn't),
+        raises `BulkUpsertPartialFailure` carrying the unsent tail —
+        already-committed earlier chunks are left server-side and the
+        caller should only retry the unsent portion.
+        """
+        if not items:
+            return
+        for chunk_start in range(0, len(items), SERVER_BULK_UPSERT_BATCH_CAP):
+            chunk = items[chunk_start : chunk_start + SERVER_BULK_UPSERT_BATCH_CAP]
+            body = {
+                "device_id": device_id,
+                "items": [_bulk_item_payload(item) for item in chunk],
+            }
+            try:
+                resp = self._client.post(
+                    "/v1/track/sessions/bulk",
+                    json=body,
+                    headers=self._headers(),
+                )
+                _raise(resp)
+            except TrackAPIError as e:
+                unsent = items[chunk_start:]
+                raise BulkUpsertPartialFailure(
+                    str(e),
+                    unsent_items=unsent,
+                    status_code=e.status_code,
+                    detail=e.detail,
+                ) from e
+            except Exception as e:
+                # Network errors etc. — same partial-failure semantics: this
+                # chunk and everything after it didn't reach the server.
+                unsent = items[chunk_start:]
+                raise BulkUpsertPartialFailure(
+                    str(e),
+                    unsent_items=unsent,
+                ) from e
+
     def list_sessions(
         self,
         *,
@@ -395,6 +482,20 @@ def _session_payload(session: Any) -> dict[str, Any]:
     raise TypeError(f"Unsupported session payload type: {type(session)!r}")
 
 
+def _bulk_item_payload(item: "BulkSessionUpsert") -> dict[str, Any]:
+    out: dict[str, Any] = {
+        "path": item.path,
+        "session": _session_payload(item.session),
+    }
+    if item.content_codec is not None:
+        out["content_codec"] = item.content_codec
+    if item.raw_bytes is not None:
+        out["raw_bytes"] = item.raw_bytes
+    if item.stored_bytes is not None:
+        out["stored_bytes"] = item.stored_bytes
+    return out
+
+
 def _json_body(body: Any) -> dict[str, Any]:
     if is_dataclass(body):
         return asdict(body)

fleet/track/daemon.py

@@ -25,7 +25,12 @@
 import uuid
 from typing import TYPE_CHECKING, Optional
 
-from .api import TrackAPIClient, TrackAPIError
+from .api import (
+    BulkSessionUpsert,
+    BulkUpsertPartialFailure,
+    TrackAPIClient,
+    TrackAPIError,
+)
 from .blocklist import TrackBlocklist
 from .drainer import QueueDrainer
 from .merkle import HashCache, MerkleTree
@@ -141,6 +146,11 @@ def __init__(
         # in the orchestrator metadata index.
         self._metadata_indexed: dict[str, str] = {}
         self._metadata_lock = threading.Lock()
+        # Workers append confirmed uploads here; the main loop flushes via
+        # the bulk endpoint each iteration. Buffering halves orchestrator
+        # round trips and avoids workers blocking on per-file POSTs.
+        self._metadata_buffer: list[tuple[str, str, BulkSessionUpsert]] = []
+        self._metadata_buffer_lock = threading.Lock()
 
     # ------------------------------------------------------------------ #
     # Public entry points                                                  #
@@ -187,6 +197,7 @@ def run_once(self, *, device_id: Optional[str] = None) -> "ReconcileResult":  #
 
         # Wait for in-flight uploads.
         self._pool.drain(timeout=60)
+        self._flush_metadata_buffer()
         if self._manifest_dirty:
             self._upload_manifest()
             self._manifest_dirty = False
@@ -248,6 +259,7 @@ def run(self) -> None:
                 last_queue_reset = now
 
             self._drain_queue()
+            self._flush_metadata_buffer()
             # Persist manifest opportunistically when the queue is idle, but
             # also at least every MANIFEST_FLUSH_INTERVAL seconds so a daemon
             # that stays continuously busy still publishes progress. The
@@ -269,6 +281,7 @@ def run(self) -> None:
         if self._pool:
             self._pool.drain(timeout=60)
             self._pool.shutdown()
+        self._flush_metadata_buffer()
         if self._manifest_dirty:
             self._upload_manifest()
         self._queue.close()
@@ -375,10 +388,18 @@ def _upload_manifest(self) -> None:
     # ------------------------------------------------------------------ #
 
     def _drain_queue(self) -> None:
-        """Delegate to QueueDrainer; one pass."""
+        """Drain pending work into the upload pool until the queue is empty.
+
+        Tight-loops drain_once so dispatch isn't capped by the main loop's
+        10s sleep. The thread pool's worker count is the real concurrency
+        ceiling; this just keeps it fed.
+        """
         if self._drainer is None:
             return
-        self._drainer.drain_once(self._device_id)
+        while not self._stop.is_set():
+            result = self._drainer.drain_once(self._device_id)
+            if result.claimed == 0:
+                break
 
     # ------------------------------------------------------------------ #
     # Upload callbacks                                                     #
@@ -431,12 +452,13 @@ def _upsert_session_metadata(
         *,
         upload_payload: UploadPayload | None = None,
     ) -> None:
-        """Best-effort metadata index update for a confirmed S3 object.
+        """Buffer a metadata upsert for the next bulk flush.
 
-        The v1 syncer's correctness still comes from S3 bytes + manifest. The
-        metadata index is a read-side accelerator for listing/resume, so a
-        transient failure here should be retried on a later reconcile rather
-        than marking the file upload failed.
+        Workers call this synchronously from upload-completion callbacks;
+        the actual HTTP roundtrip happens later in `_flush_metadata_buffer`,
+        which the daemon main loop invokes after each drain pass and on
+        graceful shutdown. Buffering halves orchestrator round trips and
+        keeps workers off the network for metadata writes.
 
         `upload_payload` is only present immediately after this process uploads
         the file. For files merely confirmed by the remote manifest, omit
@@ -453,27 +475,72 @@ def _upsert_session_metadata(
         if session is None:
             return
 
-        kwargs = {"include_content_metadata": False}
         if upload_payload is not None:
-            kwargs = {
-                "content_codec": upload_payload.content_codec,
-                "raw_bytes": upload_payload.raw_bytes,
-                "stored_bytes": upload_payload.stored_bytes,
-            }
+            item = BulkSessionUpsert(
+                path=rel_path,
+                session=session,
+                content_codec=upload_payload.content_codec,
+                raw_bytes=upload_payload.raw_bytes,
+                stored_bytes=upload_payload.stored_bytes,
+            )
+        else:
+            item = BulkSessionUpsert(path=rel_path, session=session)
+
+        with self._metadata_buffer_lock:
+            # Drop any prior buffered entry for this path; the latest sha wins.
+            self._metadata_buffer = [
+                (p, s, i) for p, s, i in self._metadata_buffer if p != rel_path
+            ]
+            self._metadata_buffer.append((rel_path, sha256, item))
+
+    def _flush_metadata_buffer(self) -> None:
+        """Send buffered metadata upserts to the orchestrator in one bulk call.
+
+        Best-effort: a transient failure leaves only the unsent entries in
+        the buffer for the next flush. The bulk client raises
+        `BulkUpsertPartialFailure` when a later chunk fails after earlier
+        chunks succeeded — we use that to avoid re-sending items the server
+        already has. If the daemon dies, the next reconcile re-queues
+        anything missing, so we tolerate buffer loss.
+        """
+        with self._metadata_buffer_lock:
+            if not self._metadata_buffer:
+                return
+            pending = self._metadata_buffer
+            self._metadata_buffer = []
 
         try:
-            self._api.upsert_session(
+            self._api.upsert_sessions_bulk(
                 device_id=self._device_id,
-                path=rel_path,
-                session=session,
-                **kwargs,
+                items=[item for _, _, item in pending],
+            )
+            sent = pending
+            unsent: list[tuple[str, str, BulkSessionUpsert]] = []
+        except BulkUpsertPartialFailure as e:
+            unsent_set = {id(it) for it in e.unsent_items}
+            sent = [t for t in pending if id(t[2]) not in unsent_set]
+            unsent = [t for t in pending if id(t[2]) in unsent_set]
+            log.warning(
+                "bulk metadata upsert partial failure: %d sent, %d unsent: %s",
+                len(sent),
+                len(unsent),
+                e,
             )
         except Exception as e:
-            log.warning("metadata upsert failed %s: %s", rel_path, e)
-            return
-
-        with self._metadata_lock:
-            self._metadata_indexed[rel_path] = sha256
+            log.warning("bulk metadata upsert failed (%d items): %s", len(pending), e)
+            sent = []
+            unsent = pending
+
+        if unsent:
+            with self._metadata_buffer_lock:
+                # Put unsent items back at the front so they're retried first;
+                # newer items appended during the flush stay in order.
+                self._metadata_buffer = unsent + self._metadata_buffer
+
+        if sent:
+            with self._metadata_lock:
+                for rel_path, sha256, _ in sent:
+                    self._metadata_indexed[rel_path] = sha256
 
     # ------------------------------------------------------------------ #
     # Status                                                               #

fleet/track/drainer.py

@@ -26,7 +26,7 @@
 
 # Server caps /v1/track/upload-urls at 100 paths per request. Match it
 # here to avoid a 400 if a single drain claims more.
-DEFAULT_BATCH_SIZE = 32
+DEFAULT_BATCH_SIZE = 100
 
 
 @dataclass(frozen=True)

fleet/track/scrubber.py

@@ -121,7 +121,8 @@ def scrub(text: str, rules: Iterable[Rule] = DEFAULT_RULES) -> ScrubResult:
 
     Hits are recorded against the *original* text's line numbers so that
     `flt track inspect` can point at the line the user can find in their
-    on-disk session file.
+    on-disk session file. The upload path uses `scrub_text` instead since
+    it doesn't need hits — see scrub_bytes.
     """
     rule_list: Sequence[Rule] = tuple(rules)
     hits: list[Hit] = []
@@ -133,21 +134,29 @@ def scrub(text: str, rules: Iterable[Rule] = DEFAULT_RULES) -> ScrubResult:
             line = text.count("\n", 0, m.start()) + 1
             hits.append(Hit(rule=rule.name, line=line, matched=m.group(0)))
 
-    # Second pass: actual substitution. We re-run regexes here rather than
-    # building offsets, because subs in earlier rules can change later
-    # rules' match positions (e.g. a long secret becoming "[REDACTED]"
-    # could expose a substring that looks like another secret).
-    for rule in rule_list:
-        text = rule.pattern.sub(rule.replacement, text)
+    return ScrubResult(text=scrub_text(text, rule_list), hits=tuple(hits))
+
 
-    return ScrubResult(text=text, hits=tuple(hits))
+def scrub_text(text: str, rules: Iterable[Rule] = DEFAULT_RULES) -> str:
+    """Apply substitution rules in order; return scrubbed text only.
+
+    Half the work of `scrub` — skips the hit-enumeration pass. Used by
+    the upload path, which discards hits anyway. Re-runs regexes for
+    substitution rather than building offsets, because subs in earlier
+    rules can change later rules' match positions (e.g. a long secret
+    becoming "[REDACTED]" could expose a substring that looks like
+    another secret).
+    """
+    for rule in rules:
+        text = rule.pattern.sub(rule.replacement, text)
+    return text
 
 
 def scrub_bytes(data: bytes, rules: Iterable[Rule] = DEFAULT_RULES) -> bytes:
     """Scrub raw bytes (decoded as UTF-8, re-encoded). Backward-compatible
     wrapper for the upload path that just wants the scrubbed payload."""
     try:
         text = data.decode("utf-8", errors="replace")
-        return scrub(text, rules).text.encode("utf-8")
+        return scrub_text(text, rules).encode("utf-8")
     except Exception:
         return data