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