Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/airflow/utils/timeout.py: 43%

1# 

2# Licensed to the Apache Software Foundation (ASF) under one 

3# or more contributor license agreements. See the NOTICE file 

4# distributed with this work for additional information 

5# regarding copyright ownership. The ASF licenses this file 

6# to you under the Apache License, Version 2.0 (the 

7# "License"); you may not use this file except in compliance 

8# with the License. You may obtain a copy of the License at 

9# 

10# http://www.apache.org/licenses/LICENSE-2.0 

11# 

12# Unless required by applicable law or agreed to in writing, 

13# software distributed under the License is distributed on an 

14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

15# KIND, either express or implied. See the License for the 

16# specific language governing permissions and limitations 

17# under the License. 

18from __future__ import annotations 

19 

20import os 

21import signal 

22from threading import Timer 

23from typing import ContextManager 

24 

25from airflow.exceptions import AirflowTaskTimeout 

26from airflow.utils.log.logging_mixin import LoggingMixin 

27from airflow.utils.platform import IS_WINDOWS 

28 

29_timeout = ContextManager[None] 

30 

31 

32class TimeoutWindows(_timeout, LoggingMixin): 

33 """Windows timeout version: To be used in a ``with`` block and timeout its content.""" 

34 

35 def __init__(self, seconds=1, error_message="Timeout"): 

36 super().__init__() 

37 self._timer: Timer | None = None 

38 self.seconds = seconds 

39 self.error_message = error_message + ", PID: " + str(os.getpid()) 

40 

41 def handle_timeout(self, *args): 

42 """Logs information and raises AirflowTaskTimeout.""" 

43 self.log.error("Process timed out, PID: %s", str(os.getpid())) 

44 raise AirflowTaskTimeout(self.error_message) 

45 

46 def __enter__(self): 

47 if self._timer: 

48 self._timer.cancel() 

49 self._timer = Timer(self.seconds, self.handle_timeout) 

50 self._timer.start() 

51 

52 def __exit__(self, type_, value, traceback): 

53 if self._timer: 

54 self._timer.cancel() 

55 self._timer = None 

56 

57 

58class TimeoutPosix(_timeout, LoggingMixin): 

59 """POSIX Timeout version: To be used in a ``with`` block and timeout its content.""" 

60 

61 def __init__(self, seconds=1, error_message="Timeout"): 

62 super().__init__() 

63 self.seconds = seconds 

64 self.error_message = error_message + ", PID: " + str(os.getpid()) 

65 

66 def handle_timeout(self, signum, frame): 

67 """Logs information and raises AirflowTaskTimeout.""" 

68 self.log.error("Process timed out, PID: %s", str(os.getpid())) 

69 raise AirflowTaskTimeout(self.error_message) 

70 

71 def __enter__(self): 

72 try: 

73 signal.signal(signal.SIGALRM, self.handle_timeout) 

74 signal.setitimer(signal.ITIMER_REAL, self.seconds) 

75 except ValueError: 

76 self.log.warning("timeout can't be used in the current context", exc_info=True) 

77 

78 def __exit__(self, type_, value, traceback): 

79 try: 

80 signal.setitimer(signal.ITIMER_REAL, 0) 

81 except ValueError: 

82 self.log.warning("timeout can't be used in the current context", exc_info=True) 

83 

84 

85if IS_WINDOWS: 

86 timeout: type[TimeoutWindows | TimeoutPosix] = TimeoutWindows 

87else: 

88 timeout = TimeoutPosix