src/third_party/chromite/lib/paygen/gslock.py

   1 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """This library can use Google Storage files as basis for locking.
   6
   7    This is mostly convenient because it works inter-server.
   8 """
   9
  10 from __future__ import print_function
  11
  12 import fixup_path
  13 fixup_path.FixupPath()
  14
  15 import datetime
  16 import logging
  17 import os
  18 import random
  19 import re
  20 import socket
  21
  22 from chromite.lib.paygen import gslib
  23 from chromite.lib.paygen import utils
  24
  25
  26 class LockProbeError(Exception):
  27   """Raised when there was an error probing a lock file."""
  28
  29
  30 class LockNotAcquired(Exception):
  31   """Raised when the lock is already held by another process."""
  32
  33
  34 class Lock(object):
  35   """This class manages a google storage file as a form of lock.
  36
  37   This class can be used in conjuction with a "with" clause to ensure
  38   the lock is released, or directly.
  39
  40     try:
  41       with gslock.Lock("gs://chromoes-releases/lock-file"):
  42         # Protected code
  43         ...
  44     except LockNotAcquired:
  45       # Error handling
  46       ...
  47
  48     lock = gslock.Lock("gs://chromoes-releases/lock-file")
  49     try:
  50       lock.Acquire()
  51     except LockNotAcquired:
  52       # Error handling
  53
  54     # Protected code
  55     ...
  56
  57     lock.Release()
  58
  59     Locking is strictly atomic, except when timeouts are involved.
  60
  61     It assumes that local server time is in sync with Google Storage server
  62     time.
  63   """
  64
  65   def __init__(self, gs_path, lock_timeout_mins=120, dry_run=False):
  66     """Initializer for the lock.
  67
  68     Args:
  69       gs_path:
  70         Path to the potential GS file we use for lock management.
  71       lock_timeout_mins:
  72         How long should an existing lock be considered valid? This timeout
  73         should be long enough that it's never hit unless a server is
  74         unexpectedly rebooted, lost network connectivity or had
  75         some other catastrophic error.
  76       dry_run: do nothing, always succeed
  77     """
  78     self._gs_path = gs_path
  79     self._timeout = datetime.timedelta(minutes=lock_timeout_mins)
  80     self._contents = repr((socket.gethostname(), os.getpid(), id(self),
  81                            random.random()))
  82     self._generation = 0
  83     self._dry_run = dry_run
  84
  85   def _LockExpired(self):
  86     """Check to see if an existing lock has timed out.
  87
  88     Returns:
  89       True if the lock is expired. False otherwise.
  90     """
  91     try:
  92       modified = self.LastModified()
  93     except LockProbeError:
  94       # If we couldn't figure out when the file was last modified, it might
  95       # have already been released. In any case, it's probably not safe to try
  96       # to clear the lock, so we'll return False here.
  97       return False
  98     return modified and datetime.datetime.utcnow() > modified + self._timeout
  99
 100   def _AcquireLock(self, filename, retries):
 101     """Attempt to acquire the lock.
 102
 103     Args:
 104       filename: local file to copy into the lock's contents.
 105       retries: How many times to retry to GS operation to fetch the lock.
 106
 107     Returns:
 108       Whether or not the lock was acquired.
 109     """
 110     try:
 111       res = gslib.RunGsutilCommand(['cp', '-v', filename, self._gs_path],
 112                                    generation=self._generation,
 113                                    redirect_stdout=True, redirect_stderr=True)
 114       m = re.search(r'%s#(\d+)' % self._gs_path, res.error)
 115       if m:
 116         error = None
 117         self._generation = int(m.group(1))
 118       else:
 119         error = 'No generation found.'
 120         self._generation = 0
 121     except gslib.GSLibError as ex:
 122       error = str(ex)
 123
 124     try:
 125       result = gslib.Cat(self._gs_path, generation=self._generation)
 126     except gslib.CatFail as ex:
 127       self._generation = 0
 128       if error:
 129         raise LockNotAcquired(error)
 130       raise LockNotAcquired(ex)
 131
 132     if result == self._contents:
 133       if error is not None:
 134         logging.warning('Lock at %s acquired despite copy error.',
 135                         self._gs_path)
 136     elif self._LockExpired() and retries >= 0:
 137       logging.warning('Timing out lock at %s.', self._gs_path)
 138       try:
 139         # Attempt to set our generation to whatever the current generation is.
 140         res = gslib.RunGsutilCommand(['stat', self._gs_path],
 141                                      redirect_stdout=True)
 142         m = re.search(r'Generation:\s*(\d+)', res.output)
 143         # Make sure the lock is still expired and hasn't been stolen by
 144         # someone else.
 145         if self._LockExpired():
 146           self._generation = int(m.group(1))
 147       except gslib.GSLibError as ex:
 148         logging.warning('Exception while stat-ing %s: %s', self._gs_path, ex)
 149         logging.warning('Lock may have been cleared by someone else')
 150
 151       self._AcquireLock(filename, retries - 1)
 152     else:
 153       self._generation = 0
 154       raise LockNotAcquired(result)
 155
 156   def LastModified(self):
 157     """Return the lock's last modification time.
 158
 159     If the lock is already acquired, uses in-memory values. Otherwise, probes
 160     the lock file.
 161
 162     Returns:
 163       The UTC time when the lock was last modified. None if corresponding
 164       attribute is missing.
 165
 166     Raises:
 167       LockProbeError: if a (non-acquired) lock is not present.
 168     """
 169     try:
 170       res = gslib.RunGsutilCommand(['stat', self._gs_path],
 171                                    redirect_stdout=True)
 172       m = re.search(r'Creation time:\s*(.*)', res.output)
 173       if not m:
 174         raise LockProbeError('Failed to extract creation time.')
 175       return datetime.datetime.strptime(m.group(1), '%a, %d %b %Y %H:%M:%S %Z')
 176     except gslib.GSLibError as ex:
 177       raise LockProbeError(ex)
 178
 179     return None
 180
 181   def Acquire(self):
 182     """Attempt to acquire the lock.
 183
 184     Will remove an existing lock if it has timed out.
 185
 186     Raises:
 187       LockNotAcquired if it is unable to get the lock.
 188     """
 189     if self._dry_run:
 190       return
 191
 192     with utils.CreateTempFileWithContents(self._contents) as tmp_file:
 193       self._AcquireLock(tmp_file.name, gslib.RETRY_ATTEMPTS)
 194
 195   def Release(self):
 196     """Release the lock."""
 197     if self._dry_run:
 198       return
 199
 200     try:
 201       gslib.Remove(self._gs_path, generation=self._generation,
 202                    ignore_no_match=True)
 203     except gslib.RemoveFail:
 204       if not self._LockExpired():
 205         raise
 206       logging.warning('Lock at %s expired and was stolen.', self._gs_path)
 207     self._generation = 0
 208
 209   def Renew(self):
 210     """Resets the timeout on a lock you are holding.
 211
 212     Raises:
 213       LockNotAcquired if it can't Renew the lock for any reason.
 214     """
 215     if self._dry_run:
 216       return
 217
 218     if int(self._generation) == 0:
 219       raise LockNotAcquired('Lock not held')
 220     self.Acquire()
 221
 222   def __enter__(self):
 223     """Support for entering a with clause."""
 224     self.Acquire()
 225     return self
 226
 227   def __exit__(self, _type, _value, _traceback):
 228     """Support for exiting a with clause."""
 229     self.Release()