src/nepi/testbeds/planetlab/rspawn.py

   1 # Utility library for spawning remote asynchronous tasks
   2 from nepi.util import server
   3 import getpass
   4 import logging
   5
   6 class STDOUT:
   7     """
   8     Special value that when given to remote_spawn in stderr causes stderr to
   9     redirect to whatever stdout was redirected to.
  10     """
  11
  12 class RUNNING:
  13     """
  14     Process is still running
  15     """
  16
  17 class FINISHED:
  18     """
  19     Process is finished
  20     """
  21
  22 class NOT_STARTED:
  23     """
  24     Process hasn't started running yet (this should be very rare)
  25     """
  26
  27 def remote_spawn(command, pidfile, stdout='/dev/null', stderr=STDOUT, stdin='/dev/null', home=None, create_home=False, sudo=False,
  28         host = None, port = None, user = None, agent = None,
  29         ident_key = None, server_key = None,
  30         tty = False, hostip = None):
  31     """
  32     Spawn a remote command such that it will continue working asynchronously.
  33
  34     Parameters:
  35         command: the command to run - it should be a single line.
  36
  37         pidfile: path of a (ideally unique to this task) pidfile for tracking the process.
  38
  39         stdout: path of a file to redirect standard output to - must be a string.
  40             Defaults to /dev/null
  41         stderr: path of a file to redirect standard error to - string or the special STDOUT value
  42             to redirect to the same file stdout was redirected to. Defaults to STDOUT.
  43         stdin: path of a file with input to be piped into the command's standard input
  44
  45         home: path of a folder to use as working directory - should exist, unless you specify create_home
  46
  47         create_home: if True, the home folder will be created first with mkdir -p
  48
  49         sudo: whether the command needs to be executed as root
  50
  51         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
  52
  53     Returns:
  54         (stdout, stderr), process
  55
  56         Of the spawning process, which only captures errors at spawning time.
  57         Usually only useful for diagnostics.
  58     """
  59     # Start process in a "daemonized" way, using nohup and heavy
  60     # stdin/out redirection to avoid connection issues
  61     if stderr is STDOUT:
  62         stderr = '&1'
  63     else:
  64         stderr = ' ' + stderr
  65
  66     daemon_command = '{ { %(command)s  > %(stdout)s 2>%(stderr)s < %(stdin)s & } ; echo $! 1 > %(pidfile)s ; }' % {
  67         'command' : command,
  68         'pidfile' : server.shell_escape(pidfile),
  69
  70         'stdout' : stdout,
  71         'stderr' : stderr,
  72         'stdin' : stdin,
  73     }
  74
  75     cmd = "%(create)s%(gohome)s rm -f %(pidfile)s ; %(sudo)s nohup bash -c %(command)s " % {
  76             'command' : server.shell_escape(daemon_command),
  77
  78             'sudo' : 'sudo -S' if sudo else '',
  79
  80             'pidfile' : server.shell_escape(pidfile),
  81             'gohome' : 'cd %s ; ' % (server.shell_escape(home),) if home else '',
  82             'create' : 'mkdir -p %s ; ' % (server.shell_escape,) if create_home else '',
  83         }
  84
  85     (out,err),proc = server.popen_ssh_command(
  86         cmd,
  87         host = host,
  88         port = port,
  89         user = user,
  90         agent = agent,
  91         ident_key = ident_key,
  92         server_key = server_key,
  93         tty = tty ,
  94         hostip = hostip
  95         )
  96
  97     if proc.wait():
  98         raise RuntimeError, "Failed to set up application on host %s: %s %s" % (host, out,err,)
  99
 100     return (out,err),proc
 101
 102 @server.eintr_retry
 103 def remote_check_pid(pidfile,
 104         host = None, port = None, user = None, agent = None,
 105         ident_key = None, server_key = None, hostip = None):
 106     """
 107     Check the pidfile of a process spawned with remote_spawn.
 108
 109     Parameters:
 110         pidfile: the pidfile passed to remote_span
 111
 112         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 113
 114     Returns:
 115
 116         A (pid, ppid) tuple useful for calling remote_status and remote_kill,
 117         or None if the pidfile isn't valid yet (maybe the process is still starting).
 118     """
 119
 120     (out,err),proc = server.popen_ssh_command(
 121         "cat %(pidfile)s" % {
 122             'pidfile' : pidfile,
 123         },
 124         host = host,
 125         port = port,
 126         user = user,
 127         agent = agent,
 128         ident_key = ident_key,
 129         server_key = server_key,
 130         hostip = hostip
 131         )
 132
 133     if proc.wait():
 134         return None
 135
 136     if out:
 137         try:
 138             return map(int,out.strip().split(' ',1))
 139         except:
 140             # Ignore, many ways to fail that don't matter that much
 141             return None
 142
 143
 144 @server.eintr_retry
 145 def remote_status(pid, ppid,
 146         host = None, port = None, user = None, agent = None,
 147         ident_key = None, server_key = None, hostip = None):
 148     """
 149     Check the status of a process spawned with remote_spawn.
 150
 151     Parameters:
 152         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 153
 154         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 155
 156     Returns:
 157
 158         One of NOT_STARTED, RUNNING, FINISHED
 159     """
 160
 161     (out,err),proc = server.popen_ssh_command(
 162         "ps --pid %(pid)d -o pid | grep -c %(pid)d ; true" % {
 163             'ppid' : ppid,
 164             'pid' : pid,
 165         },
 166         host = host,
 167         port = port,
 168         user = user,
 169         agent = agent,
 170         ident_key = ident_key,
 171         server_key = server_key,
 172         hostip = hostip
 173         )
 174
 175     if proc.wait():
 176         return NOT_STARTED
 177
 178     status = False
 179     if out:
 180         try:
 181             status = bool(int(out.strip()))
 182         except:
 183             if out or err:
 184                 logging.warn("Error checking remote status:\n%s%s\n", out, err)
 185             # Ignore, many ways to fail that don't matter that much
 186             return NOT_STARTED
 187     return RUNNING if status else FINISHED
 188
 189
 190 @server.eintr_retry
 191 def remote_kill(pid, ppid, sudo = False,
 192         host = None, port = None, user = None, agent = None,
 193         ident_key = None, server_key = None, hostip = None,
 194         nowait = False):
 195     """
 196     Kill a process spawned with remote_spawn.
 197
 198     First tries a SIGTERM, and if the process does not end in 10 seconds,
 199     it sends a SIGKILL.
 200
 201     Parameters:
 202         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 203
 204         sudo: whether the command was run with sudo - careful killing like this.
 205
 206         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 207
 208     Returns:
 209
 210         Nothing, should have killed the process
 211     """
 212
 213     if sudo:
 214         subkill = "$(ps --ppid %(pid)d -o pid h)" % { 'pid' : pid }
 215     else:
 216         subkill = ""
 217     cmd = """
 218 SUBKILL="%(subkill)s" ;
 219 %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 220 %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 221 for x in 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 ; do
 222     sleep 0.2
 223     if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` == '0' ]; then
 224         break
 225     else
 226         %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 227         %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 228     fi
 229     sleep 1.8
 230 done
 231 if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` != '0' ]; then
 232     %(sudo)s kill -9 -- -%(pid)d $SUBKILL || /bin/true
 233     %(sudo)s kill -9 %(pid)d $SUBKILL || /bin/true
 234 fi
 235 """
 236     if nowait:
 237         cmd = "( %s ) >/dev/null 2>/dev/null </dev/null &" % (cmd,)
 238
 239     (out,err),proc = server.popen_ssh_command(
 240         cmd % {
 241             'ppid' : ppid,
 242             'pid' : pid,
 243             'sudo' : 'sudo -S' if sudo else '',
 244             'subkill' : subkill,
 245         },
 246         host = host,
 247         port = port,
 248         user = user,
 249         agent = agent,
 250         ident_key = ident_key,
 251         server_key = server_key,
 252         hostip = hostip
 253         )
 254
 255     # wait, don't leave zombies around
 256     proc.wait()
 257
 258