From 3d401c263f5b4ee534eacf967ac7234f7c4ee029 Mon Sep 17 00:00:00 2001
From: Adam Norberg <norberg@google.com>
Date: Fri, 21 Nov 2025 14:01:25 -0800
Subject: [PATCH] Caffeinate fetches on Mac.

Since fetches involve multiple subprocess calls, any of which can be
slow, the per-subprocess caffeination strategy does not seem suitable --
the Mac might sleep as soon as the wake lock is dropped, before it
starts a new one. This instead implements a context manager to allow
caffeinating a scope.

To allow flag control, caffeinate.scope takes an argument that decides
whether or not it should actually do anything useful; it looks silly,
but the alternative is to interfere with flag parsing more or to require
users to write separate codepaths to decide whether to enter the context
manager scope or not; the "use the context manager in a mode where it
does not do anything" prevents this.

Bug: 462507017
Change-Id: Icc5bb9cadda30b5a120f112b10bf96ffd3b6550f
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/7183647
Reviewed-by: Gavin Mak <gavinmak@google.com>
Commit-Queue: Adam Norberg <norberg@google.com>
---
 caffeinate.py       | 31 +++++++++++++++++++++++++++++++
 fetch.py            | 13 +++++++++++--
 tests/fetch_test.py |  6 +++++-
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/caffeinate.py b/caffeinate.py
index 2e309dfa6d..e59e06f0b0 100644
--- a/caffeinate.py
+++ b/caffeinate.py
@@ -2,6 +2,7 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import contextlib
 import os
 import subprocess
 import sys
@@ -26,3 +27,33 @@ def call(args, **call_kwargs):
         else:
             args = ['caffeinate'] + args
     return subprocess.call(args, **call_kwargs)
+
+
+@contextlib.contextmanager
+def scope(actually_caffeinate=True):
+    """Acts as a context manager keeping a Mac awake, unless flagged off.
+
+    If the process is not running on a Mac, or `actually_caffeinate` is falsey,
+    this acts as a context manager that does nothing. The `actually_caffeinate`
+    flag is provided so command line flags can control the caffeinate behavior
+    without requiring weird plumbing to use or not use the context manager.
+
+    If running on a Mac while actually_caffeinate is True (the default), this
+    runs `caffeinate` in a separate process, which is terminated when the
+    context manager exits.
+    """
+    if sys.platform != 'darwin' or not actually_caffeinate:
+        # Behave like a no-op context manager.
+        yield False
+        return
+
+    cmd = ['caffeinate', '-i', '-w', str(os.getpid())]
+
+    proc = subprocess.Popen(cmd,
+                            stdin=subprocess.DEVNULL,
+                            stdout=subprocess.DEVNULL,
+                            stderr=subprocess.DEVNULL)
+    try:
+        yield True
+    finally:
+        proc.terminate()
diff --git a/fetch.py b/fetch.py
index 99ba59d16a..9614b65527 100755
--- a/fetch.py
+++ b/fetch.py
@@ -24,6 +24,7 @@ import shlex
 import subprocess
 import sys
 
+import caffeinate
 import gclient_utils
 import git_common
 
@@ -215,6 +216,13 @@ def handle_args(argv):
         type=str,
         default=None,
         help='Protocol to use to fetch dependencies, defaults to https.')
+    parser.add_argument(
+        '--caffeinate',
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help=('On macOS, prevent idle sleep during the operation. Enabled by'
+              ' default. Use --no-caffeinate to disable. No effect on other'
+              ' platforms.'))
 
     parser.add_argument('config',
                         type=str,
@@ -306,8 +314,9 @@ def run(options, spec, root):
 
 def main():
     args = handle_args(sys.argv)
-    spec, root = run_config_fetch(args.config, args.props)
-    return run(args, spec, root)
+    with caffeinate.scope(actually_caffeinate=args.caffeinate):
+        spec, root = run_config_fetch(args.config, args.props)
+        return run(args, spec, root)
 
 
 if __name__ == '__main__':
diff --git a/tests/fetch_test.py b/tests/fetch_test.py
index 5790487714..bf9277a865 100755
--- a/tests/fetch_test.py
+++ b/tests/fetch_test.py
@@ -49,6 +49,7 @@ class TestUtilityFunctions(unittest.TestCase):
                                force=False,
                                config='foo',
                                protocol_override=None,
+                               caffeinate=True,
                                props=[]), response)
 
         response = fetch.handle_args([
@@ -63,11 +64,13 @@ class TestUtilityFunctions(unittest.TestCase):
                                force=True,
                                config='foo',
                                protocol_override='sso',
+                               caffeinate=True,
                                props=['--some-param=1', '--bar=2']), response)
 
         response = fetch.handle_args([
             'filename', '-n', '--dry-run', '--no-hooks', '--nohistory',
-            '--force', '-p', 'sso', 'foo', '--some-param=1', '--bar=2'
+            '--force', '--no-caffeinate', '-p', 'sso', 'foo', '--some-param=1',
+            '--bar=2'
         ])
         self.assertEqual(
             argparse.Namespace(dry_run=True,
@@ -76,6 +79,7 @@ class TestUtilityFunctions(unittest.TestCase):
                                force=True,
                                config='foo',
                                protocol_override='sso',
+                               caffeinate=False,
                                props=['--some-param=1', '--bar=2']), response)
 
     @mock.patch('os.path.exists', return_value=False)