chromium_depot_tools/lockfile.py

# Copyright 2020 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Exclusive filelocking for all supported platforms."""

from __future__ import print_function

import contextlib
import logging
import os
import sys
import time


class LockError(Exception):
  pass


if sys.platform.startswith('win'):
  # Windows implementation
  import win32imports

  BYTES_TO_LOCK = 1

  def _open_file(lockfile):
    return win32imports.Handle(
        win32imports.CreateFileW(
            lockfile,  # lpFileName
            win32imports.GENERIC_WRITE,  # dwDesiredAccess
            0,  # dwShareMode=prevent others from opening file
            None,  # lpSecurityAttributes
            win32imports.CREATE_ALWAYS,  # dwCreationDisposition
            win32imports.FILE_ATTRIBUTE_NORMAL,  # dwFlagsAndAttributes
            None  # hTemplateFile
        ))

  def _close_file(handle):
    # CloseHandle releases lock too.
    win32imports.CloseHandle(handle)

  def _lock_file(handle):
    ret = win32imports.LockFileEx(
        handle,  # hFile
        win32imports.LOCKFILE_FAIL_IMMEDIATELY
        | win32imports.LOCKFILE_EXCLUSIVE_LOCK,  # dwFlags
        0,  #dwReserved
        BYTES_TO_LOCK,  # nNumberOfBytesToLockLow
        0,  # nNumberOfBytesToLockHigh
        win32imports.Overlapped()  # lpOverlapped
    )
    # LockFileEx returns result as bool, which is converted into an integer
    # (1 == successful; 0 == not successful)
    if ret == 0:
      error_code = win32imports.GetLastError()
      raise OSError('Failed to lock handle (error code: %d).' % error_code)
else:
  # Unix implementation
  import fcntl

  def _open_file(lockfile):
    open_flags = (os.O_CREAT | os.O_WRONLY)
    return os.open(lockfile, open_flags, 0o644)

  def _close_file(fd):
    os.close(fd)

  def _lock_file(fd):
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)


def _try_lock(lockfile):
  f = _open_file(lockfile)
  try:
    _lock_file(f)
  except Exception:
    _close_file(f)
    raise
  return lambda: _close_file(f)


def _lock(path, timeout=0):
  """_lock returns function to release the lock if locking was successful.

  _lock also implements simple retry logic."""
  elapsed = 0
  while True:
    try:
      return _try_lock(path + '.locked')
    except (OSError, IOError) as e:
      if elapsed < timeout:
        sleep_time = min(10, timeout - elapsed)
        logging.info(
            'Could not create git cache lockfile; '
            'will retry after sleep(%d).', sleep_time)
        elapsed += sleep_time
        time.sleep(sleep_time)
        continue
      raise LockError("Error locking %s (err: %s)" % (path, str(e)))


@contextlib.contextmanager
def lock(path, timeout=0):
  """Get exclusive lock to path.

  Usage:
    import lockfile
    with lockfile.lock(path, timeout):
      # Do something
      pass

   """
  release_fn = _lock(path, timeout)
  try:
    yield
  finally:
    release_fn()