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

1# Copyright (C) 2024 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

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

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15 

16from threading import Event 

17from typing import Callable, List 

18 

19from grpc import ServicerContext 

20 

21from buildgrid.server.logging import buildgrid_logger 

22 

23LOGGER = buildgrid_logger(__name__) 

24 

25 

26class CancellationContext: 

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

28 """ 

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

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

31 underlying grpc.ServicerContext is triggered. 

32 """ 

33 

34 self._event = Event() 

35 self._context = context 

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

37 context.add_callback(self._on_callback) 

38 

39 def is_cancelled(self) -> bool: 

40 return self._event.is_set() 

41 

42 def _on_callback(self) -> None: 

43 LOGGER.debug("Request cancelled.") 

44 self._event.set() 

45 for callback in self._callbacks: 

46 callback() 

47 

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

49 self._callbacks.append(callback) 

50 if self.is_cancelled(): 

51 callback()