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     (out,err),proc = server.popen_ssh_command(
  85         cmd,
  86         host = host,
  87         port = port,
  88         user = user,
  89         agent = agent,
  90         ident_key = ident_key,
  91         server_key = server_key,
  92         tty = tty ,
  93         hostip = hostip
  94         )
  95
  96     if proc.wait():
  97         raise RuntimeError, "Failed to set up application: %s %s" % (out,err,)
  98
  99     return (out,err),proc
 100
 101 @server.eintr_retry
 102 def remote_check_pid(pidfile,
 103         host = None, port = None, user = None, agent = None,
 104         ident_key = None, server_key = None, hostip = None):
 105     """
 106     Check the pidfile of a process spawned with remote_spawn.
 107
 108     Parameters:
 109         pidfile: the pidfile passed to remote_span
 110
 111         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 112
 113     Returns:
 114
 115         A (pid, ppid) tuple useful for calling remote_status and remote_kill,
 116         or None if the pidfile isn't valid yet (maybe the process is still starting).
 117     """
 118
 119     (out,err),proc = server.popen_ssh_command(
 120         "cat %(pidfile)s" % {
 121             'pidfile' : pidfile,
 122         },
 123         host = host,
 124         port = port,
 125         user = user,
 126         agent = agent,
 127         ident_key = ident_key,
 128         server_key = server_key,
 129         hostip = hostip
 130         )
 131
 132     if proc.wait():
 133         return None
 134
 135     if out:
 136         try:
 137             return map(int,out.strip().split(' ',1))
 138         except:
 139             # Ignore, many ways to fail that don't matter that much
 140             return None
 141
 142
 143 @server.eintr_retry
 144 def remote_status(pid, ppid,
 145         host = None, port = None, user = None, agent = None,
 146         ident_key = None, server_key = None, hostip = None):
 147     """
 148     Check the status of a process spawned with remote_spawn.
 149
 150     Parameters:
 151         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 152
 153         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 154
 155     Returns:
 156
 157         One of NOT_STARTED, RUNNING, FINISHED
 158     """
 159
 160     (out,err),proc = server.popen_ssh_command(
 161         "ps --pid %(pid)d -o pid | grep -c %(pid)d ; true" % {
 162             'ppid' : ppid,
 163             'pid' : pid,
 164         },
 165         host = host,
 166         port = port,
 167         user = user,
 168         agent = agent,
 169         ident_key = ident_key,
 170         server_key = server_key,
 171         hostip = hostip
 172         )
 173
 174     if proc.wait():
 175         return NOT_STARTED
 176
 177     status = False
 178     if out:
 179         try:
 180             status = bool(int(out.strip()))
 181         except:
 182             if out or err:
 183                 logging.warn("Error checking remote status:\n%s%s\n", out, err)
 184             # Ignore, many ways to fail that don't matter that much
 185             return NOT_STARTED
 186     return RUNNING if status else FINISHED
 187
 188
 189 @server.eintr_retry
 190 def remote_kill(pid, ppid, sudo = False,
 191         host = None, port = None, user = None, agent = None,
 192         ident_key = None, server_key = None, hostip = None,
 193         nowait = False):
 194     """
 195     Kill a process spawned with remote_spawn.
 196
 197     First tries a SIGTERM, and if the process does not end in 10 seconds,
 198     it sends a SIGKILL.
 199
 200     Parameters:
 201         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 202
 203         sudo: whether the command was run with sudo - careful killing like this.
 204
 205         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 206
 207     Returns:
 208
 209         Nothing, should have killed the process
 210     """
 211
 212     if sudo:
 213         subkill = "$(ps --ppid %(pid)d -o pid h)" % { 'pid' : pid }
 214     else:
 215         subkill = ""
 216     cmd = """
 217 SUBKILL="%(subkill)s" ;
 218 %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 219 %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 220 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
 221     sleep 0.2
 222     if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` == '0' ]; then
 223         break
 224     else
 225         %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 226         %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 227     fi
 228     sleep 1.8
 229 done
 230 if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` != '0' ]; then
 231     %(sudo)s kill -9 -- -%(pid)d $SUBKILL || /bin/true
 232     %(sudo)s kill -9 %(pid)d $SUBKILL || /bin/true
 233 fi
 234 """
 235     if nowait:
 236         cmd = "( %s ) >/dev/null 2>/dev/null </dev/null &" % (cmd,)
 237
 238     (out,err),proc = server.popen_ssh_command(
 239         cmd % {
 240             'ppid' : ppid,
 241             'pid' : pid,
 242             'sudo' : 'sudo -S' if sudo else '',
 243             'subkill' : subkill,
 244         },
 245         host = host,
 246         port = port,
 247         user = user,
 248         agent = agent,
 249         ident_key = ident_key,
 250         server_key = server_key,
 251         hostip = hostip
 252         )
 253
 254     # wait, don't leave zombies around
 255     proc.wait()
 256
 257