Coverage for /usr/local/lib/python3.10/site-packages/opal_server-0.0.0-py3.10.egg/opal_server/policy/watcher/task.py: 75%

1import asyncio 

2import os 

3import signal 

4from typing import Any, Coroutine, List, Optional 

5 

6from fastapi_websocket_pubsub import Topic 

7from fastapi_websocket_pubsub.pub_sub_server import PubSubEndpoint 

8from opal_common.http_utils import redact_url_in_text 

9from opal_common.logger import logger 

10from opal_common.sources.base_policy_source import BasePolicySource 

11from opal_server.config import opal_server_config 

12 

13 

14class BasePolicyWatcherTask: 

15 """Manages the asyncio tasks of the policy watcher.""" 

16 

17 def __init__(self, pubsub_endpoint: PubSubEndpoint): 

18 self._tasks: List[asyncio.Task] = [] 

19 self._should_stop: Optional[asyncio.Event] = None 

20 self._pubsub_endpoint = pubsub_endpoint 

21 self._webhook_tasks: List[asyncio.Task] = [] 

22 

23 async def __aenter__(self): 

24 await self.start() 

25 return self 

26 

27 async def __aexit__(self, exc_type, exc, tb): 

28 await self.stop() 

29 

30 async def _on_webhook(self, topic: Topic, data: Any): 

31 logger.info(f"Webhook listener triggered ({len(self._webhook_tasks)})") 

32 # Rebuild rather than remove-while-iterating: list.remove() inside a 

33 # `for t in self._webhook_tasks` loop skips the element after each removal, 

34 # so finished tasks accumulate. Retrieve exceptions before dropping the 

35 # references — otherwise a failed trigger() is only reported by asyncio's 

36 # generic "Task exception was never retrieved" at GC time. 

37 for t in self._webhook_tasks: 

38 if t.done() and not t.cancelled() and t.exception() is not None: 

39 # git exceptions can embed a credentialed remote URL verbatim; 

40 # scrub here as well as in the global log patcher, which only 

41 # runs once configure_logs() has installed it. 

42 exc_text = redact_url_in_text(repr(t.exception())) 

43 logger.error(f"Webhook trigger task failed: {exc_text}") 

44 self._webhook_tasks = [t for t in self._webhook_tasks if not t.done()] 

45 self._webhook_tasks.append(asyncio.create_task(self.trigger(topic, data))) 

46 

47 async def _listen_to_webhook_notifications(self): 

48 # Webhook api route can be hit randomly in all workers, so it publishes a message to the webhook topic. 

49 # This listener, running in the leader's context, would actually trigger the repo pull 

50 

51 async def _subscribe_internal(): 

52 logger.info( 

53 "listening on webhook topic: '{topic}'", 

54 topic=opal_server_config.POLICY_REPO_WEBHOOK_TOPIC, 

55 ) 

56 await self._pubsub_endpoint.subscribe( 

57 [opal_server_config.POLICY_REPO_WEBHOOK_TOPIC], 

58 self._on_webhook, 

59 ) 

60 

61 if self._pubsub_endpoint.broadcaster is not None: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true

62 async with self._pubsub_endpoint.broadcaster.get_listening_context(): 

63 await _subscribe_internal() 

64 await self._pubsub_endpoint.broadcaster.get_reader_task() 

65 

66 # Stop the watcher if broadcaster disconnects 

67 self.signal_stop() 

68 else: 

69 # If no broadcaster is configured, just subscribe, no need to wait on anything 

70 await _subscribe_internal() 

71 

72 async def start(self): 

73 """Starts the policy watcher and registers a failure callback to 

74 terminate gracefully.""" 

75 logger.info("Launching policy watcher") 

76 self._tasks.append(asyncio.create_task(self._listen_to_webhook_notifications())) 

77 self._init_should_stop() 

78 

79 async def stop(self): 

80 """Stops all policy watcher tasks.""" 

81 logger.info("Stopping policy watcher") 

82 for task in self._tasks + self._webhook_tasks: 

83 if not task.done(): 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true

84 task.cancel() 

85 await asyncio.gather(*self._tasks, *self._webhook_tasks, return_exceptions=True) 

86 

87 async def trigger(self, topic: Topic, data: Any): 

88 """Triggers the policy watcher from outside to check for changes (git 

89 pull)""" 

90 raise NotImplementedError() 

91 

92 def wait_until_should_stop(self) -> Coroutine: 

93 """Waits until self.signal_stop() is called on the watcher. 

94 

95 allows us to keep the repo watcher context alive until signalled 

96 to stop from outside. 

97 """ 

98 self._init_should_stop() 

99 return self._should_stop.wait() 

100 

101 def signal_stop(self): 

102 """Signal the repo watcher it should stop.""" 

103 self._init_should_stop() 

104 self._should_stop.set() 

105 

106 def _init_should_stop(self): 

107 if self._should_stop is None: 

108 self._should_stop = asyncio.Event() 

109 

110 async def _fail(self, exc: Exception): 

111 """Called when the watcher fails, and stops all tasks gracefully.""" 

112 logger.error("policy watcher failed with exception: {err}", err=repr(exc)) 

113 self.signal_stop() 

114 # trigger uvicorn graceful shutdown 

115 os.kill(os.getpid(), signal.SIGTERM) 

116 

117 

118class PolicyWatcherTask(BasePolicyWatcherTask): 

119 def __init__(self, policy_source: BasePolicySource, *args, **kwargs): 

120 self._watcher = policy_source 

121 super().__init__(*args, **kwargs) 

122 

123 async def start(self): 

124 await super().start() 

125 self._watcher.add_on_failure_callback(self._fail) 

126 self._tasks.append(asyncio.create_task(self._watcher.run())) 

127 

128 async def stop(self): 

129 await self._watcher.stop() 

130 return await super().stop() 

131 

132 async def trigger(self, topic: Topic, data: Any): 

133 """Triggers the policy watcher from outside to check for changes (git 

134 pull)""" 

135 await self._watcher.check_for_changes()