Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/decorators/time.py: 100.00%

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 

15import functools 

16import inspect 

17import time 

18from datetime import timedelta 

19from typing import Any, Callable, Dict, Iterator, Optional, TypeVar, cast 

20 

21import grpc 

22 

23from buildgrid.server.metrics_utils import publish_timer_metric 

24 

25Func = TypeVar("Func", bound=Callable) # type: ignore[type-arg] 

26 

27 

28def _service_metadata(*args: Any, _error: Optional[str] = None) -> Dict[str, str]: 

29 # If the decorator is being used for service methods, then we assume errors are already being handled. 

30 # If an error does get through somehow, assume this leads to an internal error. 

31 if len(args) == 3: 

32 if isinstance(context := args[2], grpc.ServicerContext): 

33 code = context.code() # type: ignore[attr-defined] 

34 if code is None: 

35 code = grpc.StatusCode.OK if _error is None else grpc.StatusCode.INTERNAL 

36 return {"status": code.name} 

37 

38 # In this case, the error handling is being used for a normal method. Report the error. 

39 if _error is not None: 

40 return {"exceptionType": _error} 

41 

42 return {} 

43 

44 

45def timed(metric_name: str, **tags: str) -> Callable[[Func], Func]: 

46 def decorator(func: Func) -> Func: 

47 @functools.wraps(func) 

48 def return_wrapper(*args: Any, **kwargs: Any) -> Any: 

49 start_time = time.perf_counter() 

50 error: Optional[str] = None 

51 try: 

52 return func(*args, **kwargs) 

53 except Exception as e: 

54 error = e.__class__.__name__ 

55 raise 

56 finally: 

57 run_time = timedelta(seconds=time.perf_counter() - start_time) 

58 publish_timer_metric(metric_name, run_time, **_service_metadata(*args, _error=error), **tags) 

59 

60 @functools.wraps(func) 

61 def yield_wrapper(*args: Any, **kwargs: Any) -> Iterator[Any]: 

62 start_time = time.perf_counter() 

63 error: Optional[str] = None 

64 try: 

65 yield from func(*args, **kwargs) 

66 except Exception as e: 

67 error = e.__class__.__name__ 

68 raise 

69 finally: 

70 run_time = timedelta(seconds=time.perf_counter() - start_time) 

71 publish_timer_metric(metric_name, run_time, **_service_metadata(*args, _error=error), **tags) 

72 

73 if inspect.isgeneratorfunction(func): 

74 return cast(Func, yield_wrapper) 

75 return cast(Func, return_wrapper) 

76 

77 return decorator