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

@@ -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()

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,6 +314,7 @@ def run(options, spec, root):
 
 def main():
     args = handle_args(sys.argv)
+    with caffeinate.scope(actually_caffeinate=args.caffeinate):
         spec, root = run_config_fetch(args.config, args.props)
         return run(args, spec, root)
 

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)