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