1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
|
# Copyright 2009-2011 Canonical Ltd. This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).
"""Utilities for doing the sort of thing the os module does."""
__metaclass__ = type
__all__ = [
'ensure_directory_exists',
'get_pid_from_file',
'kill_by_pidfile',
'get_pid_from_file',
'open_for_writing',
'override_environ',
'remove_if_exists',
'remove_tree',
'two_stage_kill',
'until_no_eintr',
]
from contextlib import contextmanager
import errno
import os.path
import shutil
from signal import (
SIGKILL,
SIGTERM,
)
import socket
import time
def remove_tree(path):
"""Remove the tree at 'path' from disk."""
if os.path.exists(path):
shutil.rmtree(path)
def set_environ(new_values):
"""Set the environment variables as specified by new_values.
:return: a dict of the old values
"""
old_values = {}
for name, value in new_values.iteritems():
old_values[name] = os.environ.get(name)
if value is None:
if old_values[name] is not None:
del os.environ[name]
else:
os.environ[name] = value
return old_values
@contextmanager
def override_environ(**kwargs):
"""Override environment variables with the kwarg values.
If a value is None, the environment variable is deleted. Variables are
restored to their previous state when exiting the context.
"""
old_values = set_environ(kwargs)
try:
yield
finally:
set_environ(old_values)
def until_no_eintr(retries, function, *args, **kwargs):
"""Run 'function' until it doesn't raise EINTR errors.
:param retries: The maximum number of times to try running 'function'.
:param function: The function to run.
:param *args: Arguments passed to the function.
:param **kwargs: Keyword arguments passed to the function.
:return: The return value of 'function'.
"""
if not retries:
return
for i in range(retries):
try:
return function(*args, **kwargs)
except (IOError, OSError), e:
if e.errno == errno.EINTR:
continue
raise
except socket.error, e:
# In Python 2.6 we can use IOError instead. It also has
# reason.errno but we might be using 2.5 here so use the
# index hack.
if e[0] == errno.EINTR:
continue
raise
else:
raise
def ensure_directory_exists(directory, mode=0777):
"""Create 'directory' if it doesn't exist.
:return: True if the directory had to be created, False otherwise.
"""
try:
os.makedirs(directory, mode=mode)
except OSError, e:
if e.errno == errno.EEXIST:
return False
raise
return True
def open_for_writing(filename, mode, dirmode=0777):
"""Open 'filename' for writing, creating directories if necessary.
:param filename: The path of the file to open.
:param mode: The mode to open the filename with. Should be 'w', 'a' or
something similar. See ``open`` for more details. If you pass in
a read-only mode (e.g. 'r'), then we'll just accept that and return
a read-only file-like object.
:param dirmode: The mode to use to create directories, if necessary.
:return: A file-like object that can be used to write to 'filename'.
"""
try:
return open(filename, mode)
except IOError, e:
if e.errno == errno.ENOENT:
os.makedirs(os.path.dirname(filename), mode=dirmode)
return open(filename, mode)
def _kill_may_race(pid, signal_number):
"""Kill a pid accepting that it may not exist."""
try:
os.kill(pid, signal_number)
except OSError, e:
if e.errno in (errno.ESRCH, errno.ECHILD):
# Process has already been killed.
return
# Some other issue (e.g. different user owns it)
raise
def two_stage_kill(pid, poll_interval=0.1, num_polls=50):
"""Kill process 'pid' with SIGTERM. If it doesn't die, SIGKILL it.
:param pid: The pid of the process to kill.
:param poll_interval: The polling interval used to check if the
process is still around.
:param num_polls: The number of polls to do before doing a SIGKILL.
"""
# Kill the process.
_kill_may_race(pid, SIGTERM)
# Poll until the process has ended.
for i in range(num_polls):
try:
# Reap the child process and get its return value. If it's not
# gone yet, continue.
new_pid, result = os.waitpid(pid, os.WNOHANG)
if new_pid:
return result
time.sleep(poll_interval)
except OSError, e:
if e.errno in (errno.ESRCH, errno.ECHILD):
# Raised if the process is gone by the time we try to get the
# return value.
return
# The process is still around, so terminate it violently.
_kill_may_race(pid, SIGKILL)
def get_pid_from_file(pidfile_path):
"""Retrieve the PID from the given file, if it exists, None otherwise."""
if not os.path.exists(pidfile_path):
return None
# Get the pid.
pid = open(pidfile_path, 'r').read().split()[0]
try:
pid = int(pid)
except ValueError:
# pidfile contains rubbish
return None
return pid
def kill_by_pidfile(pidfile_path, poll_interval=0.1, num_polls=50):
"""Kill a process identified by the pid stored in a file.
The pid file is removed from disk.
"""
try:
pid = get_pid_from_file(pidfile_path)
if pid is None:
return
two_stage_kill(pid, poll_interval, num_polls)
finally:
remove_if_exists(pidfile_path)
def remove_if_exists(path):
"""Remove the given file if it exists."""
try:
os.remove(path)
except OSError, e:
if e.errno != errno.ENOENT:
raise
|