Add a test opcode that sleeps for a given duration
[ganeti-local] / lib / utils.py
index d964f86..6cbd268 100644 (file)
@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#
 #
 
 # Copyright (C) 2006, 2007 Google Inc.
@@ -20,6 +20,7 @@
 
 
 """Ganeti small utilities
+
 """
 
 
@@ -33,13 +34,23 @@ import socket
 import tempfile
 import shutil
 import errno
+import pwd
+import itertools
+import select
+import fcntl
+
+from cStringIO import StringIO
 
 from ganeti import logger
 from ganeti import errors
+from ganeti import constants
+
 
 _locksheld = []
 _re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
 
+debug = False
+
 class RunResult(object):
   """Simple class for holding the result of running external programs.
 
@@ -74,6 +85,10 @@ class RunResult(object):
     else:
       self.fail_reason = "unable to determine termination reason"
 
+    if debug and self.failed:
+      logger.Debug("Command '%s' failed (%s); output: %s" %
+                   (self.cmd, self.fail_reason, self.output))
+
   def _GetOutput(self):
     """Returns the combined stdout and stderr for easier usage.
 
@@ -85,7 +100,7 @@ class RunResult(object):
 
 def _GetLockFile(subsystem):
   """Compute the file name for a given lock name."""
-  return "/var/lock/ganeti_lock_%s" % subsystem
+  return "%s/ganeti_lock_%s" % (constants.LOCK_DIR, subsystem)
 
 
 def Lock(name, max_retries=None, debug=False):
@@ -211,6 +226,7 @@ def RunCmd(cmd):
     shell = True
   env = os.environ.copy()
   env["LC_ALL"] = "C"
+  poller = select.poll()
   child = subprocess.Popen(cmd, shell=shell,
                            stderr=subprocess.PIPE,
                            stdout=subprocess.PIPE,
@@ -218,8 +234,35 @@ def RunCmd(cmd):
                            close_fds=True, env=env)
 
   child.stdin.close()
-  out = child.stdout.read()
-  err = child.stderr.read()
+  poller.register(child.stdout, select.POLLIN)
+  poller.register(child.stderr, select.POLLIN)
+  out = StringIO()
+  err = StringIO()
+  fdmap = {
+    child.stdout.fileno(): (out, child.stdout),
+    child.stderr.fileno(): (err, child.stderr),
+    }
+  for fd in fdmap:
+    status = fcntl.fcntl(fd, fcntl.F_GETFL)
+    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
+
+  while fdmap:
+    for fd, event in poller.poll():
+      if event & select.POLLIN or event & select.POLLPRI:
+        data = fdmap[fd][1].read()
+        # no data from read signifies EOF (the same as POLLHUP)
+        if not data:
+          poller.unregister(fd)
+          del fdmap[fd]
+          continue
+        fdmap[fd][0].write(data)
+      if (event & select.POLLNVAL or event & select.POLLHUP or
+          event & select.POLLERR):
+        poller.unregister(fd)
+        del fdmap[fd]
+
+  out = out.getvalue()
+  err = err.getvalue()
 
   status = child.wait()
   if status >= 0:
@@ -414,6 +457,12 @@ class HostInfo:
     self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
     self.ip = self.ipaddrs[0]
 
+  def ShortName(self):
+    """Returns the hostname without domain.
+
+    """
+    return self.name.split('.')[0]
+
   @staticmethod
   def SysName():
     """Return the current system's name.
@@ -721,21 +770,98 @@ def RemoveAuthorizedKey(file_name, key):
   key_fields = key.split()
 
   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
-  out = os.fdopen(fd, 'w')
   try:
-    f = open(file_name, 'r')
+    out = os.fdopen(fd, 'w')
     try:
-      for line in f:
-        # Ignore whitespace changes while comparing lines
-        if line.split() != key_fields:
+      f = open(file_name, 'r')
+      try:
+        for line in f:
+          # Ignore whitespace changes while comparing lines
+          if line.split() != key_fields:
+            out.write(line)
+
+        out.flush()
+        os.rename(tmpname, file_name)
+      finally:
+        f.close()
+    finally:
+      out.close()
+  except:
+    RemoveFile(tmpname)
+    raise
+
+
+def SetEtcHostsEntry(file_name, ip, hostname, aliases):
+  """Sets the name of an IP address and hostname in /etc/hosts.
+
+  """
+  # Ensure aliases are unique
+  aliases = UniqueSequence([hostname] + aliases)[1:]
+
+  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
+  try:
+    out = os.fdopen(fd, 'w')
+    try:
+      f = open(file_name, 'r')
+      try:
+        written = False
+        for line in f:
+          fields = line.split()
+          if fields and not fields[0].startswith('#') and ip == fields[0]:
+            continue
           out.write(line)
 
-      out.flush()
-      os.rename(tmpname, file_name)
+        out.write("%s\t%s" % (ip, hostname))
+        if aliases:
+          out.write(" %s" % ' '.join(aliases))
+        out.write('\n')
+
+        out.flush()
+        os.fsync(out)
+        os.rename(tmpname, file_name)
+      finally:
+        f.close()
     finally:
-      f.close()
-  finally:
-    out.close()
+      out.close()
+  except:
+    RemoveFile(tmpname)
+    raise
+
+
+def RemoveEtcHostsEntry(file_name, hostname):
+  """Removes a hostname from /etc/hosts.
+
+  IP addresses without names are removed from the file.
+  """
+  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
+  try:
+    out = os.fdopen(fd, 'w')
+    try:
+      f = open(file_name, 'r')
+      try:
+        for line in f:
+          fields = line.split()
+          if len(fields) > 1 and not fields[0].startswith('#'):
+            names = fields[1:]
+            if hostname in names:
+              while hostname in names:
+                names.remove(hostname)
+              if names:
+                out.write("%s %s\n" % (fields[0], ' '.join(names)))
+              continue
+
+          out.write(line)
+
+        out.flush()
+        os.fsync(out)
+        os.rename(tmpname, file_name)
+      finally:
+        f.close()
+    finally:
+      out.close()
+  except:
+    RemoveFile(tmpname)
+    raise
 
 
 def CreateBackup(file_name):
@@ -748,10 +874,20 @@ def CreateBackup(file_name):
     raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
                                 file_name)
 
-  # Warning: the following code contains a race condition when we create more
-  # than one backup of the same file in a second.
-  backup_name = file_name + '.backup-%d' % int(time.time())
-  shutil.copyfile(file_name, backup_name)
+  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
+  dir_name = os.path.dirname(file_name)
+
+  fsrc = open(file_name, 'rb')
+  try:
+    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
+    fdst = os.fdopen(fd, 'wb')
+    try:
+      shutil.copyfileobj(fsrc, fdst)
+    finally:
+      fdst.close()
+  finally:
+    fsrc.close()
+
   return backup_name
 
 
@@ -772,42 +908,8 @@ def ShellQuoteArgs(args):
   return ' '.join([ShellQuote(i) for i in args])
 
 
-def _ParseIpOutput(output):
-  """Parsing code for GetLocalIPAddresses().
-
-  This function is split out, so we can unit test it.
-
-  """
-  re_ip = re.compile('^(\d+\.\d+\.\d+\.\d+)(?:/\d+)$')
-
-  ips = []
-  for line in output.splitlines(False):
-    fields = line.split()
-    if len(line) < 4:
-      continue
-    m = re_ip.match(fields[3])
-    if m:
-      ips.append(m.group(1))
-
-  return ips
-
-
-def GetLocalIPAddresses():
-  """Gets a list of all local IP addresses.
-
-  Should this break one day, a small Python module written in C could
-  use the API call getifaddrs().
-
-  """
-  result = RunCmd(["ip", "-family", "inet", "-oneline", "addr", "show"])
-  if result.failed:
-    raise errors.OpExecError("Command '%s' failed, error: %s,"
-      " output: %s" % (result.cmd, result.fail_reason, result.output))
-
-  return _ParseIpOutput(result.output)
 
-
-def TcpPing(source, target, port, timeout=10, live_port_needed=True):
+def TcpPing(source, target, port, timeout=10, live_port_needed=False):
   """Simple ping implementation using TCP connect(2).
 
   Try to do a TCP connect(2) from the specified source IP to the specified
@@ -838,3 +940,138 @@ def TcpPing(source, target, port, timeout=10, live_port_needed=True):
     success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
 
   return success
+
+
+def ListVisibleFiles(path):
+  """Returns a list of all visible files in a directory.
+
+  """
+  files = [i for i in os.listdir(path) if not i.startswith(".")]
+  files.sort()
+  return files
+
+
+def GetHomeDir(user, default=None):
+  """Try to get the homedir of the given user.
+
+  The user can be passed either as a string (denoting the name) or as
+  an integer (denoting the user id). If the user is not found, the
+  'default' argument is returned, which defaults to None.
+
+  """
+  try:
+    if isinstance(user, basestring):
+      result = pwd.getpwnam(user)
+    elif isinstance(user, (int, long)):
+      result = pwd.getpwuid(user)
+    else:
+      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
+                                   type(user))
+  except KeyError:
+    return default
+  return result.pw_dir
+
+
+def NewUUID():
+  """Returns a random UUID.
+
+  """
+  f = open("/proc/sys/kernel/random/uuid", "r")
+  try:
+    return f.read(128).rstrip("\n")
+  finally:
+    f.close()
+
+
+def WriteFile(file_name, fn=None, data=None,
+              mode=None, uid=-1, gid=-1,
+              atime=None, mtime=None):
+  """(Over)write a file atomically.
+
+  The file_name and either fn (a function taking one argument, the
+  file descriptor, and which should write the data to it) or data (the
+  contents of the file) must be passed. The other arguments are
+  optional and allow setting the file mode, owner and group, and the
+  mtime/atime of the file.
+
+  If the function doesn't raise an exception, it has succeeded and the
+  target file has the new contents. If the file has raised an
+  exception, an existing target file should be unmodified and the
+  temporary file should be removed.
+
+  """
+  if not os.path.isabs(file_name):
+    raise errors.ProgrammerError("Path passed to WriteFile is not"
+                                 " absolute: '%s'" % file_name)
+
+  if [fn, data].count(None) != 1:
+    raise errors.ProgrammerError("fn or data required")
+
+  if [atime, mtime].count(None) == 1:
+    raise errors.ProgrammerError("Both atime and mtime must be either"
+                                 " set or None")
+
+
+  dir_name, base_name = os.path.split(file_name)
+  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
+  # here we need to make sure we remove the temp file, if any error
+  # leaves it in place
+  try:
+    if uid != -1 or gid != -1:
+      os.chown(new_name, uid, gid)
+    if mode:
+      os.chmod(new_name, mode)
+    if data is not None:
+      os.write(fd, data)
+    else:
+      fn(fd)
+    os.fsync(fd)
+    if atime is not None and mtime is not None:
+      os.utime(new_name, (atime, mtime))
+    os.rename(new_name, file_name)
+  finally:
+    os.close(fd)
+    RemoveFile(new_name)
+
+
+def all(seq, pred=bool):
+  "Returns True if pred(x) is True for every element in the iterable"
+  for elem in itertools.ifilterfalse(pred, seq):
+    return False
+  return True
+
+
+def any(seq, pred=bool):
+  "Returns True if pred(x) is True for at least one element in the iterable"
+  for elem in itertools.ifilter(pred, seq):
+    return True
+  return False
+
+
+def UniqueSequence(seq):
+  """Returns a list with unique elements.
+
+  Element order is preserved.
+  """
+  seen = set()
+  return [i for i in seq if i not in seen and not seen.add(i)]
+
+
+def IsValidMac(mac):
+  """Predicate to check if a MAC address is valid.
+
+  Checks wether the supplied MAC address is formally correct, only
+  accepts colon separated format.
+  """
+  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
+  return mac_check.match(mac) is not None
+
+
+def TestDelay(duration):
+  """Sleep for a fixed amount of time.
+
+  """
+  if duration < 0:
+    return False
+  time.sleep(duration)
+  return True