src/nepi/util/execfuncs.py

   1 #
   2 #    NEPI, a framework to manage network experiments
   3 #    Copyright (C) 2013 INRIA
   4 #
   5 #    This program is free software: you can redistribute it and/or modify
   6 #    it under the terms of the GNU General Public License as published by
   7 #    the Free Software Foundation, either version 3 of the License, or
   8 #    (at your option) any later version.
   9 #
  10 #    This program is distributed in the hope that it will be useful,
  11 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 #    GNU General Public License for more details.
  14 #
  15 #    You should have received a copy of the GNU General Public License
  16 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17 #
  18 # Author: Alina Quereilhac <alina.quereilhac@inria.fr>
  19
  20 from nepi.util.sshfuncs import ProcStatus, STDOUT
  21
  22 import subprocess
  23
  24 def lexec(command,
  25         user = None,
  26         sudo = False,
  27         stdin = None,
  28         env = None):
  29     """
  30     Executes a local command, returns ((stdout,stderr),process)
  31     """
  32     if env:
  33         export = ''
  34         for envkey, envval in env.iteritems():
  35             export += '%s=%s ' % (envkey, envval)
  36         command = "%s %s" % (export, command)
  37
  38     if sudo:
  39         command = "sudo %s" % command
  40     elif user:
  41         command = "su %s ; %s " % (user, command)
  42
  43
  44     p = subprocess.Popen(command,
  45             stdout = subprocess.PIPE,
  46             stderr = subprocess.PIPE,
  47             stdin  = stdin)
  48
  49     out, err = p.communicate()
  50     return (out, err)
  51
  52 def lcopy(source, dest, recursive = False):
  53     """
  54     Copies from/to localy.
  55     """
  56
  57     if TRACE:
  58         print "scp", source, dest
  59
  60     command = ["cp"]
  61     if recursive:
  62         command.append("-R")
  63
  64     if isinstance(dest, str):
  65         dest = dest.split(";")
  66
  67     if isinstance(src, str):
  68         src = src.split(";")
  69
  70     args.extend(src)
  71
  72     args.extend(dest)
  73
  74     p = subprocess.Popen(command,
  75         stdout=subprocess.PIPE,
  76         stderr=subprocess.PIPE)
  77
  78     out, err = p.communicate()
  79     return ((out, err), proc)
  80
  81 def lspawn(command, pidfile,
  82         stdout = '/dev/null',
  83         stderr = STDOUT,
  84         stdin = '/dev/null',
  85         home = None,
  86         create_home = False,
  87         sudo = False,
  88         user = None):
  89     """
  90     Spawn a local command such that it will continue working asynchronously.
  91
  92     Parameters:
  93         command: the command to run - it should be a single line.
  94
  95         pidfile: path of a (ideally unique to this task) pidfile for tracking the process.
  96
  97         stdout: path of a file to redirect standard output to - must be a string.
  98             Defaults to /dev/null
  99         stderr: path of a file to redirect standard error to - string or the special STDOUT value
 100             to redirect to the same file stdout was redirected to. Defaults to STDOUT.
 101         stdin: path of a file with input to be piped into the command's standard input
 102
 103         home: path of a folder to use as working directory - should exist, unless you specify create_home
 104
 105         create_home: if True, the home folder will be created first with mkdir -p
 106
 107         sudo: whether the command needs to be executed as root
 108
 109     Returns:
 110         (stdout, stderr), process
 111
 112         Of the spawning process, which only captures errors at spawning time.
 113         Usually only useful for diagnostics.
 114     """
 115     # Start process in a "daemonized" way, using nohup and heavy
 116     # stdin/out redirection to avoid connection issues
 117     if stderr is STDOUT:
 118         stderr = '&1'
 119     else:
 120         stderr = ' ' + stderr
 121
 122     daemon_command = '{ { %(command)s  > %(stdout)s 2>%(stderr)s < %(stdin)s & } ; echo $! 1 > %(pidfile)s ; }' % {
 123         'command' : command,
 124         'pidfile' : shell_escape(pidfile),
 125         'stdout' : stdout,
 126         'stderr' : stderr,
 127         'stdin' : stdin,
 128     }
 129
 130     cmd = "%(create)s%(gohome)s rm -f %(pidfile)s ; %(sudo)s nohup bash -c %(command)s " % {
 131             'command' : shell_escape(daemon_command),
 132             'sudo' : 'sudo -S' if sudo else '',
 133             'pidfile' : shell_escape(pidfile),
 134             'gohome' : 'cd %s ; ' % (shell_escape(home),) if home else '',
 135             'create' : 'mkdir -p %s ; ' % (shell_escape(home),) if create_home else '',
 136         }
 137
 138     (out,err),proc = lexec(cmd)
 139
 140     if proc.wait():
 141         raise RuntimeError, "Failed to set up application on host %s: %s %s" % (host, out,err,)
 142
 143     return (out,err),proc
 144
 145 def lgetpid(pidfile):
 146     """
 147     Check the pidfile of a process spawned with remote_spawn.
 148
 149     Parameters:
 150         pidfile: the pidfile passed to remote_span
 151
 152     Returns:
 153
 154         A (pid, ppid) tuple useful for calling remote_status and remote_kill,
 155         or None if the pidfile isn't valid yet (maybe the process is still starting).
 156     """
 157
 158     (out,err),proc = lexec("cat %s" % pidfile )
 159
 160     if proc.wait():
 161         return None
 162
 163     if out:
 164         try:
 165             return map(int,out.strip().split(' ',1))
 166         except:
 167             # Ignore, many ways to fail that don't matter that much
 168             return None
 169
 170 def lstatus(pid, ppid):
 171     """
 172     Check the status of a process spawned with remote_spawn.
 173
 174     Parameters:
 175         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 176
 177     Returns:
 178
 179         One of NOT_STARTED, RUNNING, FINISHED
 180     """
 181
 182     (out,err),proc = lexec(
 183         # Check only by pid. pid+ppid does not always work (especially with sudo)
 184         " (( ps --pid %(pid)d -o pid | grep -c %(pid)d && echo 'wait')  || echo 'done' ) | tail -n 1" % {
 185             'ppid' : ppid,
 186             'pid' : pid,
 187         })
 188
 189     if proc.wait():
 190         return ProcStatus.NOT_STARTED
 191
 192     status = False
 193     if out:
 194         status = (out.strip() == 'wait')
 195     else:
 196         return ProcStatus.NOT_STARTED
 197     return ProcStatus.RUNNING if status else ProcStatus.FINISHED
 198
 199
 200 def lkill(pid, ppid, sudo = False):
 201     """
 202     Kill a process spawned with lspawn.
 203
 204     First tries a SIGTERM, and if the process does not end in 10 seconds,
 205     it sends a SIGKILL.
 206
 207     Parameters:
 208         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 209
 210         sudo: whether the command was run with sudo - careful killing like this.
 211
 212     Returns:
 213
 214         Nothing, should have killed the process
 215     """
 216
 217     subkill = "$(ps --ppid %(pid)d -o pid h)" % { 'pid' : pid }
 218     cmd = """
 219 SUBKILL="%(subkill)s" ;
 220 %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 221 %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 222 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
 223     sleep 0.2
 224     if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` == '0' ]; then
 225         break
 226     else
 227         %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 228         %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 229     fi
 230     sleep 1.8
 231 done
 232 if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` != '0' ]; then
 233     %(sudo)s kill -9 -- -%(pid)d $SUBKILL || /bin/true
 234     %(sudo)s kill -9 %(pid)d $SUBKILL || /bin/true
 235 fi
 236 """
 237     if nowait:
 238         cmd = "( %s ) >/dev/null 2>/dev/null </dev/null &" % (cmd,)
 239
 240     (out,err),proc = lexec(
 241         cmd % {
 242             'ppid' : ppid,
 243             'pid' : pid,
 244             'sudo' : 'sudo -S' if sudo else '',
 245             'subkill' : subkill,
 246         })
 247
 248