Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/metrics_utils.py: 98.80%

1# Copyright (C) 2020 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 time 

16from contextlib import contextmanager 

17from datetime import timedelta 

18from typing import Dict, Iterator, Optional 

19 

20from buildgrid._protos.buildgrid.v2.monitoring_pb2 import MetricRecord 

21from buildgrid.server.context import try_current_instance, try_current_method, try_current_service 

22from buildgrid.server.enums import MetricRecordType 

23from buildgrid.server.monitoring import get_monitoring_bus 

24 

25 

26def _format_metadata(metadata: Optional[Dict[str, str]] = None) -> Dict[str, str]: 

27 if metadata is None: 

28 metadata = {} 

29 

30 metadata = {key: value for key, value in metadata.items() if value is not None} 

31 

32 # If an instance name is not specifically set, try to find one from the context. 

33 if "instanceName" not in metadata: 

34 if (instance_name := try_current_instance()) is not None: 

35 metadata["instanceName"] = instance_name 

36 

37 # If the instance name is set and empty, set it to a default for easy querying. 

38 if metadata.get("instanceName") == "": 

39 metadata["instanceName"] = "unnamed" 

40 

41 # If a service name is not specifically set, try to find one from the context. 

42 if "serviceName" not in metadata: 

43 if (service := try_current_service()) is not None: 

44 metadata["serviceName"] = service 

45 

46 # If a method name is not specifically set, try to find one from the context. 

47 if "method" not in metadata: 

48 if (method := try_current_method()) is not None: 

49 metadata["method"] = method 

50 

51 return metadata 

52 

53 

54def create_counter_record(name: str, count: float, **metadata: str) -> MetricRecord: 

55 counter_record = MetricRecord() 

56 

57 counter_record.creation_timestamp.GetCurrentTime() 

58 counter_record.type = MetricRecordType.COUNTER.value 

59 counter_record.name = name 

60 counter_record.count = count 

61 counter_record.metadata.update(_format_metadata(metadata)) 

62 

63 return counter_record 

64 

65 

66def create_gauge_record(name: str, value: float, **metadata: str) -> MetricRecord: 

67 gauge_record = MetricRecord() 

68 

69 gauge_record.creation_timestamp.GetCurrentTime() 

70 gauge_record.type = MetricRecordType.GAUGE.value 

71 gauge_record.name = name 

72 gauge_record.value = value 

73 gauge_record.metadata.update(_format_metadata(metadata)) 

74 

75 return gauge_record 

76 

77 

78def create_timer_record(name: str, duration: timedelta, **metadata: str) -> MetricRecord: 

79 timer_record = MetricRecord() 

80 

81 timer_record.creation_timestamp.GetCurrentTime() 

82 timer_record.type = MetricRecordType.TIMER.value 

83 timer_record.name = name 

84 timer_record.duration.FromTimedelta(duration) 

85 timer_record.metadata.update(_format_metadata(metadata)) 

86 

87 return timer_record 

88 

89 

90def create_distribution_record(name: str, value: float, **metadata: str) -> MetricRecord: 

91 dist_record = MetricRecord() 

92 

93 dist_record.creation_timestamp.GetCurrentTime() 

94 dist_record.type = MetricRecordType.DISTRIBUTION.value 

95 dist_record.name = name 

96 dist_record.count = value 

97 dist_record.metadata.update(_format_metadata(metadata)) 

98 

99 return dist_record 

100 

101 

102def publish_counter_metric(name: str, count: float, **metadata: str) -> None: 

103 record = create_counter_record(name, count, **metadata) 

104 monitoring_bus = get_monitoring_bus() 

105 monitoring_bus.send_record_nowait(record) 

106 

107 

108def publish_gauge_metric(name: str, value: float, **metadata: str) -> None: 

109 record = create_gauge_record(name, value, **metadata) 

110 monitoring_bus = get_monitoring_bus() 

111 monitoring_bus.send_record_nowait(record) 

112 

113 

114def publish_timer_metric(name: str, duration: timedelta, **metadata: str) -> None: 

115 record = create_timer_record(name, duration, **metadata) 

116 monitoring_bus = get_monitoring_bus() 

117 monitoring_bus.send_record_nowait(record) 

118 

119 

120def publish_distribution_metric(name: str, count: float, **metadata: str) -> None: 

121 record = create_distribution_record(name, float(count), **metadata) 

122 monitoring_bus = get_monitoring_bus() 

123 monitoring_bus.send_record_nowait(record) 

124 

125 

126@contextmanager 

127def timer(metric_name: str, **tags: str) -> Iterator[None]: 

128 start_time = time.perf_counter() 

129 error = "None" 

130 try: 

131 yield 

132 except Exception as e: 

133 error = e.__class__.__name__ 

134 raise 

135 finally: 

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

137 publish_timer_metric(metric_name, run_time, exceptionType=error, **tags)