src/nepi/testbeds/planetlab/rspawn.py

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