1#
2# Copyright (C) 2018 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import errno
18import logging
19import os
20import signal
21
22from functools import wraps
23from vts.runners.host import errors
24
25
26class TimeoutError(errors.VtsError):
27    """Timeout exception class to throw when a function times out."""
28
29
30def timeout(seconds, message=os.strerror(errno.ETIME), no_exception=False):
31    """Timeout decorator for functions.
32
33    Args:
34        seconds: int, number of seconds before timing out the decorated function.
35        message: string, error message when the decorated function times out.
36        no_exception: bool, whether to raise exception when decorated function times out.
37                      If set to False, the function will stop execution and return None.
38
39    Returns:
40        Decorated function returns if no timeout, or None if timeout but no_exception is True.
41
42    Raises:
43        TimeoutError if decorated function times out.
44        TypeError if seconds is not integer
45    """
46    def _handler_timeout(signum, frame):
47       raise TimeoutError(message)
48
49    def decorator(func):
50        def wrapper(*args, **kwargs):
51            if seconds > 0:
52                signal.signal(signal.SIGALRM, _handler_timeout)
53                signal.alarm(seconds)
54
55            try:
56                result = func(*args, **kwargs)
57            except TimeoutError as e:
58                if no_exception:
59                    logging.error(message)
60                    return None
61                else:
62                    raise e
63            finally:
64                signal.alarm(0)
65
66            return result
67
68        return wraps(func)(wrapper)
69
70    return decorator
71