Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/grpc/_observability.py: 56%

1# Copyright 2023 The gRPC authors. 

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 

15from __future__ import annotations 

16 

17import abc 

18import contextlib 

19import logging 

20import threading 

21from typing import ( 

22 Any, 

23 Generator, 

24 Generic, 

25 List, 

26 Optional, 

27 Tuple, 

28 TypeVar, 

29 Union, 

30) 

31 

32from grpc._cython import cygrpc as _cygrpc 

33from grpc._typing import ChannelArgumentType 

34 

35_LOGGER = logging.getLogger(__name__) 

36 

37_channel = Any # _channel.py imports this module. 

38ClientCallTracerCapsule = TypeVar("ClientCallTracerCapsule") 

39ServerCallTracerFactoryCapsule = TypeVar("ServerCallTracerFactoryCapsule") 

40 

41_plugin_lock: threading.RLock = threading.RLock() 

42_OBSERVABILITY_PLUGIN: Optional["ObservabilityPlugin"] = None 

43_SERVICES_TO_EXCLUDE: List[bytes] = [ 

44 b"google.monitoring.v3.MetricService", 

45 b"google.devtools.cloudtrace.v2.TraceService", 

46] 

47 

48 

49class ServerCallTracerFactory: 

50 """An encapsulation of a ServerCallTracerFactory. 

51 

52 Instances of this class can be passed to a Channel as values for the 

53 grpc.experimental.server_call_tracer_factory option 

54 """ 

55 

56 def __init__(self, address): 

57 self._address = address 

58 

59 def __int__(self): 

60 return self._address 

61 

62 

63class ObservabilityPlugin( 

64 Generic[ClientCallTracerCapsule, ServerCallTracerFactoryCapsule], 

65 metaclass=abc.ABCMeta, 

66): 

67 """Abstract base class for observability plugin. 

68 

69 *This is a semi-private class that was intended for the exclusive use of 

70 the gRPC team.* 

71 

72 The ClientCallTracerCapsule and ClientCallTracerCapsule created by this 

73 plugin should be injected to gRPC core using observability_init at the 

74 start of a program, before any channels/servers are built. 

75 

76 Any future methods added to this interface cannot have the 

77 @abc.abstractmethod annotation. 

78 

79 Attributes: 

80 _stats_enabled: A bool indicates whether tracing is enabled. 

81 _tracing_enabled: A bool indicates whether stats(metrics) is enabled. 

82 _registered_methods: A set which stores the registered method names in 

83 bytes. 

84 """ 

85 

86 _tracing_enabled: bool = False 

87 _stats_enabled: bool = False 

88 

89 @abc.abstractmethod 

90 def create_client_call_tracer( 

91 self, method_name: bytes, target: bytes 

92 ) -> ClientCallTracerCapsule: 

93 """Creates a ClientCallTracerCapsule. 

94 

95 After register the plugin, if tracing or stats is enabled, this method 

96 will be called after a call was created, the ClientCallTracer created 

97 by this method will be saved to call context. 

98 

99 The ClientCallTracer is an object which implements `grpc_core::ClientCallTracer` 

100 interface and wrapped in a PyCapsule using `client_call_tracer` as name. 

101 

102 Args: 

103 method_name: The method name of the call in byte format. 

104 target: The channel target of the call in byte format. 

105 registered_method: Whether this method is pre-registered. 

106 

107 Returns: 

108 A PyCapsule which stores a ClientCallTracer object. 

109 """ 

110 raise NotImplementedError() 

111 

112 @abc.abstractmethod 

113 def save_trace_context( 

114 self, trace_id: str, span_id: str, is_sampled: bool 

115 ) -> None: 

116 """Saves the trace_id and span_id related to the current span. 

117 

118 After register the plugin, if tracing is enabled, this method will be 

119 called after the server finished sending response. 

120 

121 This method can be used to propagate census context. 

122 

123 Args: 

124 trace_id: The identifier for the trace associated with the span as a 

125 32-character hexadecimal encoded string, 

126 e.g. 26ed0036f2eff2b7317bccce3e28d01f 

127 span_id: The identifier for the span as a 16-character hexadecimal encoded 

128 string. e.g. 113ec879e62583bc 

129 is_sampled: A bool indicates whether the span is sampled. 

130 """ 

131 raise NotImplementedError() 

132 

133 @abc.abstractmethod 

134 def create_server_call_tracer_factory( 

135 self, 

136 *, 

137 xds: bool = False, 

138 ) -> Optional[ServerCallTracerFactoryCapsule]: 

139 """Creates a ServerCallTracerFactoryCapsule. 

140 

141 This method will be called at server initialization time to create a 

142 ServerCallTracerFactory, which will be registered to gRPC core. 

143 

144 The ServerCallTracerFactory is an object which implements 

145 `grpc_core::ServerCallTracerFactory` interface and wrapped in a PyCapsule 

146 using `server_call_tracer_factory` as name. 

147 

148 Args: 

149 xds: Whether the server is xds server. 

150 Returns: 

151 A PyCapsule which stores a ServerCallTracerFactory object. Or None if 

152 plugin decides not to create ServerCallTracerFactory. 

153 """ 

154 raise NotImplementedError() 

155 

156 @abc.abstractmethod 

157 def record_rpc_latency( 

158 self, method: str, target: str, rpc_latency: float, status_code: Any 

159 ) -> None: 

160 """Record the latency of the RPC. 

161 

162 After register the plugin, if stats is enabled, this method will be 

163 called at the end of each RPC. 

164 

165 Args: 

166 method: The fully-qualified name of the RPC method being invoked. 

167 target: The target name of the RPC method being invoked. 

168 rpc_latency: The latency for the RPC in seconds, equals to the time between 

169 when the client invokes the RPC and when the client receives the status. 

170 status_code: An element of grpc.StatusCode in string format representing the 

171 final status for the RPC. 

172 """ 

173 raise NotImplementedError() 

174 

175 def set_tracing(self, enable: bool) -> None: 

176 """Enable or disable tracing. 

177 

178 Args: 

179 enable: A bool indicates whether tracing should be enabled. 

180 """ 

181 self._tracing_enabled = enable 

182 

183 def set_stats(self, enable: bool) -> None: 

184 """Enable or disable stats(metrics). 

185 

186 Args: 

187 enable: A bool indicates whether stats should be enabled. 

188 """ 

189 self._stats_enabled = enable 

190 

191 def save_registered_method(self, method_name: bytes) -> None: 

192 """Saves the method name to registered_method list. 

193 

194 When exporting metrics, method name for unregistered methods will be replaced 

195 with 'other' by default. 

196 

197 Args: 

198 method_name: The method name in bytes. 

199 """ 

200 raise NotImplementedError() 

201 

202 @property 

203 def tracing_enabled(self) -> bool: 

204 return self._tracing_enabled 

205 

206 @property 

207 def stats_enabled(self) -> bool: 

208 return self._stats_enabled 

209 

210 @property 

211 def observability_enabled(self) -> bool: 

212 return self.tracing_enabled or self.stats_enabled 

213 

214 

215@contextlib.contextmanager 

216def get_plugin() -> Generator[Optional[ObservabilityPlugin], None, None]: 

217 """Get the ObservabilityPlugin in _observability module. 

218 

219 Returns: 

220 The ObservabilityPlugin currently registered with the _observability 

221 module. Or None if no plugin exists at the time of calling this method. 

222 """ 

223 with _plugin_lock: 

224 yield _OBSERVABILITY_PLUGIN 

225 

226 

227def set_plugin(observability_plugin: Optional[ObservabilityPlugin]) -> None: 

228 """Save ObservabilityPlugin to _observability module. 

229 

230 Args: 

231 observability_plugin: The ObservabilityPlugin to save. 

232 

233 Raises: 

234 ValueError: If an ObservabilityPlugin was already registered at the 

235 time of calling this method. 

236 """ 

237 global _OBSERVABILITY_PLUGIN # pylint: disable=global-statement 

238 with _plugin_lock: 

239 if observability_plugin and _OBSERVABILITY_PLUGIN: 

240 raise ValueError("observability_plugin was already set!") 

241 _OBSERVABILITY_PLUGIN = observability_plugin 

242 

243 

244def observability_init(observability_plugin: ObservabilityPlugin) -> None: 

245 """Initialize observability with provided ObservabilityPlugin. 

246 

247 This method have to be called at the start of a program, before any 

248 channels/servers are built. 

249 

250 Args: 

251 observability_plugin: The ObservabilityPlugin to use. 

252 

253 Raises: 

254 ValueError: If an ObservabilityPlugin was already registered at the 

255 time of calling this method. 

256 """ 

257 set_plugin(observability_plugin) 

258 

259 

260def observability_deinit() -> None: 

261 """Clear the observability context, including ObservabilityPlugin and 

262 ServerCallTracerFactory 

263 

264 This method have to be called after exit observability context so that 

265 it's possible to re-initialize again. 

266 """ 

267 set_plugin(None) 

268 _cygrpc.clear_server_call_tracer_factory() 

269 

270 

271def maybe_record_rpc_latency(state: "_channel._RPCState") -> None: 

272 """Record the latency of the RPC, if the plugin is registered and stats is enabled. 

273 

274 This method will be called at the end of each RPC. 

275 

276 Args: 

277 state: a grpc._channel._RPCState object which contains the stats related to the 

278 RPC. 

279 """ 

280 # TODO(xuanwn): use channel args to exclude those metrics. 

281 for exclude_prefix in _SERVICES_TO_EXCLUDE: 

282 if exclude_prefix in state.method.encode("utf8"): 

283 return 

284 with get_plugin() as plugin: 

285 if plugin and plugin.stats_enabled: 

286 rpc_latency_s = state.rpc_end_time - state.rpc_start_time 

287 rpc_latency_ms = rpc_latency_s * 1000 

288 plugin.record_rpc_latency( 

289 state.method, state.target, rpc_latency_ms, state.code 

290 ) 

291 

292 

293def create_server_call_tracer_factory_option( 

294 xds: bool, 

295) -> Union[Tuple[ChannelArgumentType], Tuple[()]]: 

296 with get_plugin() as plugin: 

297 if plugin and plugin.stats_enabled: 

298 server_call_tracer_factory_address = ( 

299 _cygrpc.get_server_call_tracer_factory_address(plugin, xds) 

300 ) 

301 if server_call_tracer_factory_address: 

302 return ( 

303 ( 

304 "grpc.experimental.server_call_tracer_factory", 

305 ServerCallTracerFactory( 

306 server_call_tracer_factory_address 

307 ), 

308 ), 

309 ) 

310 return ()