commit 649bcdd896dd69c395e678506ab216aa496e560b
Author: Balazs Lecz <leczb@google.com>
Date:   Fri Mar 26 17:01:06 2010 +0000

    Manage the assignment of uids from the uid pool
    
    Signed-off-by: Balazs Lecz <leczb@google.com>
    Reviewed-by: Guido Trotter <ultrotter@google.com>

diff --git a/lib/uidpool.py b/lib/uidpool.py
index 85e664e..d216ec5 100644
--- a/lib/uidpool.py
+++ b/lib/uidpool.py
@@ -29,6 +29,11 @@ from the pool.
 
 """
 
+import errno
+import logging
+import os
+import random
+
 from ganeti import errors
 from ganeti import constants
 from ganeti import utils
@@ -172,3 +177,154 @@ def ExpandUidPool(uid_pool):
   for lower, higher in uid_pool:
     uids.update(range(lower, higher + 1))
   return list(uids)
+
+
+def _IsUidUsed(uid):
+  """Check if there is any process in the system running with the given user-id
+
+  """
+  pgrep_command = [constants.PGREP, "-u", uid]
+  result = utils.RunCmd(pgrep_command)
+
+  if result.exit_code == 0:
+    return True
+  elif result.exit_code == 1:
+    return False
+  else:
+    raise errors.CommandError("Running pgrep failed. exit code: %s"
+                              % result.exit_code)
+
+
+class LockedUid(object):
+  """Class representing a locked user-id in the uid-pool.
+
+  This binds together a userid and a lock.
+
+  """
+  def __init__(self, uid, lock):
+    """Constructor
+
+    @param uid: a user-id
+    @param lock: a utils.FileLock object
+
+    """
+    self._uid = uid
+    self._lock = lock
+
+  def Unlock(self):
+    # Release the exclusive lock and close the filedescriptor
+    self._lock.Close()
+
+  def __str__(self):
+    return "%s" % self._uid
+
+
+def RequestUnusedUid(all_uids):
+  """Tries to find an unused uid from the uid-pool, locks it and returns it.
+
+  Usage pattern:
+
+    1) When starting a process
+
+      from ganeti import ssconf
+      from ganeti import uidpool
+
+      # Get list of all user-ids in the uid-pool from ssconf
+      ss = ssconf.SimpleStore()
+      uid_pool = uidpool.ParseUidPool(ss.GetUidPool(), separator="\n")
+      all_uids = set(uidpool.ExpandUidPool(uid_pool))
+
+      uid = uidpool.RequestUnusedUid(all_uids)
+      try:
+        <start a process with the UID>
+        # Once the process is started, we can release the file lock
+        uid.Unlock()
+      except ..., err:
+        # Return the UID to the pool
+        uidpool.ReleaseUid(uid)
+
+    2) Stopping a process
+
+      from ganeti import uidpool
+
+      uid = <get the UID the process is running under>
+      <stop the process>
+      uidpool.ReleaseUid(uid)
+
+  @param all_uids: a set containing all the user-ids in the user-id pool
+  @return: a LockedUid object representing the unused uid. It's the caller's
+           responsibility to unlock the uid once an instance is started with
+           this uid.
+
+  """
+  # Create the lock dir if it's not yet present
+  try:
+    utils.EnsureDirs([(constants.UIDPOOL_LOCKDIR, 0755)])
+  except errors.GenericError, err:
+    raise errors.LockError("Failed to create user-id pool lock dir: %s" % err)
+
+  # Get list of currently used uids from the filesystem
+  try:
+    taken_uids = set(os.listdir(constants.UIDPOOL_LOCKDIR))
+    # Filter out spurious entries from the directory listing
+    taken_uids = all_uids.intersection(taken_uids)
+  except OSError, err:
+    raise errors.LockError("Failed to get list of used user-ids: %s" % err)
+
+  # Remove the list of used uids from the list of all uids
+  unused_uids = list(all_uids - taken_uids)
+  if not unused_uids:
+    logging.info("All user-ids in the uid-pool are marked 'taken'")
+
+  # Randomize the order of the unused user-id list
+  random.shuffle(unused_uids)
+
+  # Randomize the order of the unused user-id list
+  taken_uids = list(taken_uids)
+  random.shuffle(taken_uids)
+
+  for uid in (unused_uids + taken_uids):
+    try:
+      # Create the lock file
+      # Note: we don't care if it exists. Only the fact that we can
+      # (or can't) lock it later is what matters.
+      uid_path = utils.PathJoin(constants.UIDPOOL_LOCKDIR, str(uid))
+      lock = utils.FileLock.Open(uid_path)
+    except OSError, err:
+      raise errors.LockError("Failed to create lockfile for user-id %s: %s"
+                             % (uid, err))
+    try:
+      # Try acquiring an exclusive lock on the lock file
+      lock.Exclusive()
+      # Check if there is any process running with this user-id
+      if _IsUidUsed(uid):
+        logging.debug("There is already a process running under"
+                      " user-id %s", uid)
+        lock.Unlock()
+        continue
+      return LockedUid(uid, lock)
+    except IOError, err:
+      if err.errno == errno.EAGAIN:
+        # The file is already locked, let's skip it and try another unused uid
+        logging.debug("Lockfile for user-id is already locked %s: %s", uid, err)
+        continue
+    except errors.LockError, err:
+      # There was an unexpected error while trying to lock the file
+      logging.error("Failed to lock the lockfile for user-id %s: %s", uid, err)
+      raise
+
+  raise errors.LockError("Failed to find an unused user-id")
+
+
+def ReleaseUid(uid):
+  """This should be called when the given user-id is no longer in use.
+
+  """
+  # Make sure we release the exclusive lock, if there is any
+  uid.Unlock()
+  try:
+    uid_path = utils.PathJoin(constants.UIDPOOL_LOCKDIR, str(uid))
+    os.remove(uid_path)
+  except OSError, err:
+    raise errors.LockError("Failed to remove user-id lockfile"
+                           " for user-id %s: %s" % (uid, err))
