third-party/watchman/src/watchman/integration/lib/TempDir.py

   1 # Copyright (c) Meta Platforms, Inc. and affiliates.
   2 #
   3 # This source code is licensed under the MIT license found in the
   4 # LICENSE file in the root directory of this source tree.
   5
   6
   7 import atexit
   8 import errno
   9 import os
  10 import shutil
  11 import sys
  12 import tempfile
  13 import time
  14
  15 import pywatchman
  16
  17 from . import path_utils as path
  18
  19
  20 global_temp_dir = None
  21
  22
  23 class TempDir(object):
  24     """
  25     This is a helper for locating a reasonable place for temporary files.
  26     When run in the watchman test suite, we compute this up-front and then
  27     store everything under that temporary directory.
  28     When run under the FB internal test runner, we infer a reasonable grouped
  29     location from the process group environmental variable exported by the
  30     test runner.
  31     """
  32
  33     def __init__(self, keepAtShutdown: bool = False) -> None:
  34         # We'll put all our temporary stuff under one dir so that we
  35         # can clean it all up at the end.
  36
  37         parent_dir = tempfile.gettempdir()
  38         prefix = "wat"
  39
  40         self.temp_dir = path.get_canonical_filesystem_path(
  41             tempfile.mkdtemp(dir=parent_dir, prefix=prefix)
  42         )
  43
  44         if os.name != "nt":
  45             # On some platforms, setting the setgid bit on a directory doesn't
  46             # work if the user isn't a member of the directory's group. Set the
  47             # group explicitly to avoid this.
  48             os.chown(self.temp_dir, -1, os.getegid())
  49             # Some environments have a weird umask that can leave state
  50             # directories too open and break tests.
  51             os.umask(0o022)
  52         # Redirect all temporary files to that location
  53         tempfile.tempdir = os.fsdecode(self.temp_dir)
  54
  55         self.keep = keepAtShutdown
  56
  57         def cleanup():
  58             if self.keep:
  59                 sys.stdout.write("Preserving output in %s\n" % self.temp_dir)
  60                 return
  61             self._retry_rmtree(self.temp_dir)
  62
  63         atexit.register(cleanup)
  64
  65     def get_dir(self):
  66         return self.temp_dir
  67
  68     def set_keep(self, value) -> None:
  69         self.keep = value
  70
  71     def _retry_rmtree(self, top) -> None:
  72         # Keep trying to remove it; on Windows it may take a few moments
  73         # for any outstanding locks/handles to be released
  74         for _ in range(1, 10):
  75             shutil.rmtree(top, onerror=_remove_readonly)
  76             if not os.path.isdir(top):
  77                 return
  78             sys.stdout.write("Waiting to remove temp data under %s\n" % top)
  79             time.sleep(0.2)
  80         sys.stdout.write("Failed to completely remove %s\n" % top)
  81
  82
  83 def _remove_readonly(func, path, exc_info) -> None:
  84     # If we encounter an EPERM or EACCESS error removing a file try making its parent
  85     # directory writable and then retry the removal.  This is necessary to clean up
  86     # eden mount point directories after the checkout is unmounted, as these directories
  87     # are made read-only by "eden clone"
  88     _ex_type, ex, _traceback = exc_info
  89     if not (
  90         isinstance(ex, EnvironmentError) and ex.errno in (errno.EACCES, errno.EPERM)
  91     ):
  92         # Just ignore other errors.  This will be retried by _retry_rmtree()
  93         return
  94
  95     try:
  96         parent_dir = os.path.dirname(path)
  97         os.chmod(parent_dir, 0o755)
  98         # func() is the function that failed.
  99         # This is usually os.unlink() or os.rmdir().
 100         func(path)
 101     except OSError:
 102         return
 103
 104
 105 def get_temp_dir(keep=None):
 106     global global_temp_dir
 107     if global_temp_dir:
 108         return global_temp_dir
 109     if keep is None:
 110         keep = os.environ.get("WATCHMAN_TEST_KEEP", "0") == "1"
 111     global_temp_dir = TempDir(keep)
 112     return global_temp_dir