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