Coverage for faststream / _internal / endpoint / subscriber / supervisor.py: 97%

1from __future__ import annotations (empty)

2 

3import logging (empty)

4import os (empty)

5import time (empty)

6from asyncio import CancelledError, Task (empty)

7from collections import UserDict (empty)

8from typing import TYPE_CHECKING, Any (empty)

9 

10if TYPE_CHECKING: (empty)

11 from collections.abc import Callable, Coroutine 

12 

13 from faststream._internal.endpoint.subscriber.mixins import TasksMixin 

14 

15 

16class _SupervisorCache(UserDict[int, float]): (empty)

17 @property (empty)

18 def ttl(self) -> float: (empty)

19 return float(os.getenv("FASTSTREAM_SUPERVISOR_CACHE_TTL", "3600")) (empty)

20 

21 def _cleanup(self) -> None: (empty)

22 # no need to clean up empty storage 

23 if not len(self): (empty)

24 return (empty)

25 

26 # dict preserves insertion order hence if the first element is not outdated there is no reason to check others 

27 _, timestamp = next(iter(self.items())) (empty)

28 now = time.time() (empty)

29 if now - timestamp < self.ttl: (empty)

30 return (empty)

31 

32 now, to_delete = time.time(), [] (empty)

33 for k, v in self.items(): (empty)

34 if now - v >= self.ttl: 34 ↛ 37line 34 didn't jump to line 37 because the condition on line 34 was always true(empty)

35 to_delete.append(k) (empty)

36 else: 

37 break 

38 

39 for k in to_delete: (empty)

40 del self[k] (empty)

41 

42 def add(self, item: int) -> None: (empty)

43 self[item] = time.time() (empty)

44 

45 def __contains__(self, key: object) -> bool: (empty)

46 self._cleanup() (empty)

47 return super().__contains__(key) (empty)

48 

49 

50class TaskCallbackSupervisor: (empty)

51 """Supervisor for asyncio.Task spawned in TaskMixin implemented via task callback.""" 

52 

53 __slots__ = ( (empty)

54 "args", 

55 "func", 

56 "ignored_exceptions", 

57 "kwargs", 

58 "max_attempts", 

59 "subscriber", 

60 ) 

61 

62 # stores hash identifier of exceptions which whose info was printed 

63 __cache: _SupervisorCache = _SupervisorCache() (empty)

64 

65 def __init__( (empty)

66 self, 

67 func: Callable[..., Coroutine[Any, Any, Any]], (empty)

68 func_args: tuple[Any] | None, (empty)

69 func_kwargs: dict[str, Any] | None, (empty)

70 subscriber: TasksMixin, (empty)

71 *, 

72 ignored_exceptions: tuple[type[BaseException], ...] = (CancelledError,), (empty)

73 ) -> None: (empty)

74 self.subscriber = subscriber (empty)

75 self.func = func (empty)

76 self.args = func_args or () (empty)

77 self.kwargs = func_kwargs or {} (empty)

78 self.ignored_exceptions = ignored_exceptions (empty)

79 

80 @staticmethod (empty)

81 def _get_exception_identifier( (empty)

82 exception: BaseException, *, message_size: int = 4096 

83 ) -> int: 

84 # NOTE: method accepts only raised exceptions (e.g. with __traceback__) 

85 line, message, klass = ( (empty)

86 getattr(exception.__traceback__, "tb_lineno", 0), 

87 str(exception)[:message_size], 

88 str(exception.__class__), 

89 ) 

90 to_hash = f"{line}-{message}-{klass}" (empty)

91 

92 return hash(to_hash) (empty)

93 

94 @property (empty)

95 def is_disabled(self) -> bool: (empty)

96 # supervisor can affect some test cases, so it might be useful to have global killswitch. 

97 return bool(int(os.getenv("FASTSTREAM_SUPERVISOR_DISABLED", "0"))) (empty)

98 

99 def __call__(self, task: Task[Any]) -> None: (empty)

100 logger = self.subscriber._outer_config.logger (empty)

101 

102 logger.log( (empty)

103 f"callback for {task.get_name()} is being executed...", 

104 log_level=logging.INFO, 

105 ) 

106 

107 if task.cancelled() or self.is_disabled: (empty)

108 return (empty)

109 

110 if (exc := task.exception()) and not isinstance(exc, self.ignored_exceptions): (empty)

111 # trace is printed only once, but task is still retried 

112 identifier = self._get_exception_identifier(exc) (empty)

113 if identifier not in self.__cache: (empty)

114 self.__cache.add(identifier) (empty)

115 logger.log( (empty)

116 f"{task.get_name()} raised an exception, retrying...\n" 

117 "If this behavior causes issues, you can disable it via setting the FASTSTREAM_SUPERVISOR_DISABLED env to 1. " 

118 "Also, please consider opening issue on the repository: https://github.com/ag2ai/faststream.", 

119 exc_info=exc, 

120 log_level=logging.ERROR, 

121 ) 

122 

123 self.subscriber.add_task(self.func, self.args, self.kwargs) (empty)