1#
2# Module providing various facilities to other parts of the package
3#
4# multiprocessing/util.py
5#
6# Copyright (c) 2006-2008, R Oudkerk
7# All rights reserved.
8#
9# Redistribution and use in source and binary forms, with or without
10# modification, are permitted provided that the following conditions
11# are met:
12#
13# 1. Redistributions of source code must retain the above copyright
14#    notice, this list of conditions and the following disclaimer.
15# 2. Redistributions in binary form must reproduce the above copyright
16#    notice, this list of conditions and the following disclaimer in the
17#    documentation and/or other materials provided with the distribution.
18# 3. Neither the name of author nor the names of any contributors may be
19#    used to endorse or promote products derived from this software
20#    without specific prior written permission.
21#
22# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND
23# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
24# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
25# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
26# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
27# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
28# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
29# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
30# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
31# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
32# SUCH DAMAGE.
33#
34
35import itertools
36import weakref
37import atexit
38import threading        # we want threading to install it's
39                        # cleanup function before multiprocessing does
40
41from multiprocess.process import current_process, active_children
42
43__all__ = [
44    'sub_debug', 'debug', 'info', 'sub_warning', 'get_logger',
45    'log_to_stderr', 'get_temp_dir', 'register_after_fork',
46    'is_exiting', 'Finalize', 'ForkAwareThreadLock', 'ForkAwareLocal',
47    'SUBDEBUG', 'SUBWARNING',
48    ]
49
50#
51# Logging
52#
53
54NOTSET = 0
55SUBDEBUG = 5
56DEBUG = 10
57INFO = 20
58SUBWARNING = 25
59
60LOGGER_NAME = 'multiprocess'
61DEFAULT_LOGGING_FORMAT = '[%(levelname)s/%(processName)s] %(message)s'
62
63_logger = None
64_log_to_stderr = False
65
66def sub_debug(msg, *args):
67    if _logger:
68        _logger.log(SUBDEBUG, msg, *args)
69
70def debug(msg, *args):
71    if _logger:
72        _logger.log(DEBUG, msg, *args)
73
74def info(msg, *args):
75    if _logger:
76        _logger.log(INFO, msg, *args)
77
78def sub_warning(msg, *args):
79    if _logger:
80        _logger.log(SUBWARNING, msg, *args)
81
82def get_logger():
83    '''
84    Returns logger used by multiprocessing
85    '''
86    global _logger
87    import logging, atexit
88
89    logging._acquireLock()
90    try:
91        if not _logger:
92
93            _logger = logging.getLogger(LOGGER_NAME)
94            _logger.propagate = 0
95            logging.addLevelName(SUBDEBUG, 'SUBDEBUG')
96            logging.addLevelName(SUBWARNING, 'SUBWARNING')
97
98            # XXX multiprocessing should cleanup before logging
99            if hasattr(atexit, 'unregister'):
100                atexit.unregister(_exit_function)
101                atexit.register(_exit_function)
102            else:
103                atexit._exithandlers.remove((_exit_function, (), {}))
104                atexit._exithandlers.append((_exit_function, (), {}))
105
106    finally:
107        logging._releaseLock()
108
109    return _logger
110
111def log_to_stderr(level=None):
112    '''
113    Turn on logging and add a handler which prints to stderr
114    '''
115    global _log_to_stderr
116    import logging
117
118    logger = get_logger()
119    formatter = logging.Formatter(DEFAULT_LOGGING_FORMAT)
120    handler = logging.StreamHandler()
121    handler.setFormatter(formatter)
122    logger.addHandler(handler)
123
124    if level:
125        logger.setLevel(level)
126    _log_to_stderr = True
127    return _logger
128
129#
130# Function returning a temp directory which will be removed on exit
131#
132
133def get_temp_dir():
134    # get name of a temp directory which will be automatically cleaned up
135    if current_process()._tempdir is None:
136        import shutil, tempfile
137        tempdir = tempfile.mkdtemp(prefix='pymp-')
138        info('created temp directory %s', tempdir)
139        Finalize(None, shutil.rmtree, args=[tempdir], exitpriority=-100)
140        current_process()._tempdir = tempdir
141    return current_process()._tempdir
142
143#
144# Support for reinitialization of objects when bootstrapping a child process
145#
146
147_afterfork_registry = weakref.WeakValueDictionary()
148_afterfork_counter = itertools.count()
149
150def _run_after_forkers():
151    items = list(_afterfork_registry.items())
152    items.sort()
153    for (index, ident, func), obj in items:
154        try:
155            func(obj)
156        except Exception as e:
157            info('after forker raised exception %s', e)
158
159def register_after_fork(obj, func):
160    _afterfork_registry[(next(_afterfork_counter), id(obj), func)] = obj
161
162#
163# Finalization using weakrefs
164#
165
166_finalizer_registry = {}
167_finalizer_counter = itertools.count()
168
169
170class Finalize(object):
171    '''
172    Class which supports object finalization using weakrefs
173    '''
174    def __init__(self, obj, callback, args=(), kwargs=None, exitpriority=None):
175        assert exitpriority is None or type(exitpriority) is int
176
177        if obj is not None:
178            self._weakref = weakref.ref(obj, self)
179        else:
180            assert exitpriority is not None
181
182        self._callback = callback
183        self._args = args
184        self._kwargs = kwargs or {}
185        self._key = (exitpriority, next(_finalizer_counter))
186
187        _finalizer_registry[self._key] = self
188
189    def __call__(self, wr=None):
190        '''
191        Run the callback unless it has already been called or cancelled
192        '''
193        try:
194            del _finalizer_registry[self._key]
195        except KeyError:
196            sub_debug('finalizer no longer registered')
197        else:
198            sub_debug('finalizer calling %s with args %s and kwargs %s',
199                     self._callback, self._args, self._kwargs)
200            res = self._callback(*self._args, **self._kwargs)
201            self._weakref = self._callback = self._args = \
202                            self._kwargs = self._key = None
203            return res
204
205    def cancel(self):
206        '''
207        Cancel finalization of the object
208        '''
209        try:
210            del _finalizer_registry[self._key]
211        except KeyError:
212            pass
213        else:
214            self._weakref = self._callback = self._args = \
215                            self._kwargs = self._key = None
216
217    def still_active(self):
218        '''
219        Return whether this finalizer is still waiting to invoke callback
220        '''
221        return self._key in _finalizer_registry
222
223    def __repr__(self):
224        try:
225            obj = self._weakref()
226        except (AttributeError, TypeError):
227            obj = None
228
229        if obj is None:
230            return '<Finalize object, dead>'
231
232        x = '<Finalize object, callback=%s' % \
233            getattr(self._callback, '__name__', self._callback)
234        if self._args:
235            x += ', args=' + str(self._args)
236        if self._kwargs:
237            x += ', kwargs=' + str(self._kwargs)
238        if self._key[0] is not None:
239            x += ', exitprority=' + str(self._key[0])
240        return x + '>'
241
242
243def _run_finalizers(minpriority=None):
244    '''
245    Run all finalizers whose exit priority is not None and at least minpriority
246
247    Finalizers with highest priority are called first; finalizers with
248    the same priority will be called in reverse order of creation.
249    '''
250    if _finalizer_registry is None:
251        # This function may be called after this module's globals are
252        # destroyed.  See the _exit_function function in this module for more
253        # notes.
254        return
255
256    if minpriority is None:
257        f = lambda p : p[0][0] is not None
258    else:
259        f = lambda p : p[0][0] is not None and p[0][0] >= minpriority
260
261    items = [x for x in list(_finalizer_registry.items()) if f(x)]
262    items.sort(reverse=True)
263
264    for key, finalizer in items:
265        sub_debug('calling %s', finalizer)
266        try:
267            finalizer()
268        except Exception:
269            import traceback
270            traceback.print_exc()
271
272    if minpriority is None:
273        _finalizer_registry.clear()
274
275#
276# Clean up on exit
277#
278
279def is_exiting():
280    '''
281    Returns true if the process is shutting down
282    '''
283    return _exiting or _exiting is None
284
285_exiting = False
286
287def _exit_function(info=info, debug=debug, _run_finalizers=_run_finalizers,
288                   active_children=active_children,
289                   current_process=current_process):
290    # We hold on to references to functions in the arglist due to the
291    # situation described below, where this function is called after this
292    # module's globals are destroyed.
293
294    global _exiting
295
296    if not _exiting:
297        _exiting = True
298
299        info('process shutting down')
300        debug('running all "atexit" finalizers with priority >= 0')
301        _run_finalizers(0)
302        if current_process() is not None:
303            # We check if the current process is None here because if
304            # it's None, any call to ``active_children()`` will raise an
305            # AttributeError (active_children winds up trying to get
306            # attributes from util._current_process).  This happens in a
307            # variety of shutdown circumstances that are not well-understood
308            # because module-scope variables are not apparently supposed to
309            # be destroyed until after this function is called.  However,
310            # they are indeed destroyed before this function is called.  See
311            # issues #9775 and #15881.  Also related: #4106, #9205, and #9207.
312            for p in active_children():
313                if p._daemonic:
314                    info('calling terminate() for daemon %s', p.name)
315                    p._popen.terminate()
316
317            for p in active_children():
318                info('calling join() for process %s', p.name)
319                p.join()
320
321        debug('running the remaining "atexit" finalizers')
322        _run_finalizers()
323
324atexit.register(_exit_function)
325
326#
327# Some fork aware types
328#
329
330class ForkAwareThreadLock(object):
331    def __init__(self):
332        self._lock = threading.Lock()
333        self.acquire = self._lock.acquire
334        self.release = self._lock.release
335        register_after_fork(self, ForkAwareThreadLock.__init__)
336
337class ForkAwareLocal(threading.local):
338    def __init__(self):
339        register_after_fork(self, lambda obj : obj.__dict__.clear())
340    def __reduce__(self):
341        return type(self), ()
342