1# This file is part of Ansible
2#
3# Ansible is free software: you can redistribute it and/or modify
4# it under the terms of the GNU General Public License as published by
5# the Free Software Foundation, either version 3 of the License, or
6# (at your option) any later version.
7#
8# Ansible is distributed in the hope that it will be useful,
9# but WITHOUT ANY WARRANTY; without even the implied warranty of
10# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11# GNU General Public License for more details.
12#
13# You should have received a copy of the GNU General Public License
14# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
15
16from __future__ import (absolute_import, division, print_function)
17__metaclass__ = type
18
19import multiprocessing
20import multiprocessing.pool as mp
21
22# timeout function to make sure some fact gathering
23# steps do not exceed a time limit
24
25GATHER_TIMEOUT = None
26DEFAULT_GATHER_TIMEOUT = 10
27
28
29class TimeoutError(Exception):
30    pass
31
32
33def timeout(seconds=None, error_message="Timer expired"):
34    """
35    Timeout decorator to expire after a set number of seconds.  This raises an
36    ansible.module_utils.facts.TimeoutError if the timeout is hit before the
37    function completes.
38    """
39    def decorator(func):
40        def wrapper(*args, **kwargs):
41            timeout_value = seconds
42            if timeout_value is None:
43                timeout_value = globals().get('GATHER_TIMEOUT') or DEFAULT_GATHER_TIMEOUT
44
45            pool = mp.ThreadPool(processes=1)
46            res = pool.apply_async(func, args, kwargs)
47            pool.close()
48            try:
49                return res.get(timeout_value)
50            except multiprocessing.TimeoutError:
51                # This is an ansible.module_utils.common.facts.timeout.TimeoutError
52                raise TimeoutError('Timer expired after %s seconds' % timeout_value)
53
54        return wrapper
55
56    # If we were called as @timeout, then the first parameter will be the
57    # function we are to wrap instead of the number of seconds.  Detect this
58    # and correct it by setting seconds to our default value and return the
59    # inner decorator function manually wrapped around the function
60    if callable(seconds):
61        func = seconds
62        seconds = None
63        return decorator(func)
64
65    # If we were called as @timeout([...]) then python itself will take
66    # care of wrapping the inner decorator around the function
67
68    return decorator
69