Ticket #14: WIP, only intra-PL TUN connections, required groundwork for cross-backend...
[nepi.git] / 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 def remote_check_pid(pidfile,
100         host = None, port = None, user = None, agent = None, 
101         ident_key = None, server_key = None):
102     """
103     Check the pidfile of a process spawned with remote_spawn.
104     
105     Parameters:
106         pidfile: the pidfile passed to remote_span
107         
108         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
109     
110     Returns:
111         
112         A (pid, ppid) tuple useful for calling remote_status and remote_kill,
113         or None if the pidfile isn't valid yet (maybe the process is still starting).
114     """
115
116     (out,err),proc = server.popen_ssh_command(
117         "cat %(pidfile)s" % {
118             'pidfile' : pidfile,
119         },
120         host = host,
121         port = port,
122         user = user,
123         agent = agent,
124         ident_key = ident_key,
125         server_key = server_key
126         )
127         
128     if proc.wait():
129         return None
130     
131     if out:
132         try:
133             return map(int,out.strip().split(' ',1))
134         except:
135             # Ignore, many ways to fail that don't matter that much
136             return None
137
138
139 def remote_status(pid, ppid, 
140         host = None, port = None, user = None, agent = None, 
141         ident_key = None, server_key = None):
142     """
143     Check the status of a process spawned with remote_spawn.
144     
145     Parameters:
146         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
147         
148         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
149     
150     Returns:
151         
152         One of NOT_STARTED, RUNNING, FINISHED
153     """
154
155     (out,err),proc = server.popen_ssh_command(
156         "ps --ppid %(ppid)d -o pid | grep -c %(pid)d ; true" % {
157             'ppid' : ppid,
158             'pid' : pid,
159         },
160         host = host,
161         port = port,
162         user = user,
163         agent = agent,
164         ident_key = ident_key,
165         server_key = server_key
166         )
167     
168     if proc.wait():
169         return NOT_STARTED
170     
171     status = False
172     if out:
173         try:
174             status = bool(int(out.strip()))
175         except:
176             # Ignore, many ways to fail that don't matter that much
177             return NOT_STARTED
178     return RUNNING if status else FINISHED
179     
180
181 def remote_kill(pid, ppid, sudo = False,
182         host = None, port = None, user = None, agent = None, 
183         ident_key = None, server_key = None):
184     """
185     Kill a process spawned with remote_spawn.
186     
187     First tries a SIGTERM, and if the process does not end in 10 seconds,
188     it sends a SIGKILL.
189     
190     Parameters:
191         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
192         
193         sudo: whether the command was run with sudo - careful killing like this.
194         
195         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
196     
197     Returns:
198         
199         Nothing, should have killed the process
200     """
201
202     (out,err),proc = server.popen_ssh_command(
203         """
204 %(sudo)s kill %(pid)d
205 for x in 1 2 3 4 5 6 7 8 9 0 ; do 
206     sleep 0.1 
207     if [ `ps --pid %(ppid)d -o pid | grep -c %(pid)d` == `0` ]; then
208         break
209     fi
210     sleep 0.9
211 done
212 if [ `ps --pid %(ppid)d -o pid | grep -c %(pid)d` != `0` ]; then
213     %(sudo)s kill -9 %(pid)d %(ppid)d
214 fi
215 """ % {
216             'ppid' : ppid,
217             'pid' : pid,
218             'sudo' : 'sudo -S' if sudo else ''
219         },
220         host = host,
221         port = port,
222         user = user,
223         agent = agent,
224         ident_key = ident_key,
225         server_key = server_key
226         )
227     
228     # wait, don't leave zombies around
229     proc.wait()
230     
231
232