Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/utils/context.py: 90.91%

22 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-06-11 15:37 +0000

1import logging 

2from threading import Event 

3from typing import Callable, List 

4 

5from grpc import ServicerContext 

6 

7LOGGER = logging.getLogger(__name__) 

8 

9 

10class CancellationContext: 

11 def __init__(self, context: ServicerContext) -> None: 

12 """ 

13 Creates a wrapper for a grpc.ServicerContext which allows determining if a gRPC request has been 

14 cancelled by the client. Callbacks may be added to this context which will be invoked if the 

15 underlying grpc.ServicerContext is triggered. 

16 """ 

17 

18 self._event = Event() 

19 self._context = context 

20 self._callbacks: List[Callable[[], None]] = [] 

21 context.add_callback(self._on_callback) 

22 

23 def is_cancelled(self) -> bool: 

24 return self._event.is_set() 

25 

26 def _on_callback(self) -> None: 

27 LOGGER.debug("Request cancelled") 

28 self._event.set() 

29 for callback in self._callbacks: 

30 callback() 

31 

32 def on_cancel(self, callback: Callable[[], None]) -> None: 

33 self._callbacks.append(callback) 

34 if self.is_cancelled(): 

35 callback()