Coverage for zenmqtt/instrumentation/tracing.py: 96%

1""" 

2Tracing of messages: a producer span for every publish, and a consumer span for 

3the processing of a received message. The trace context travels in the user 

4properties of PUBLISH, so a consumer span continues the trace of the producer. 

5 

6Tracing is the base class with no-op methods; the OpenTelemetry implementation 

7is zenmqtt.instrumentation.opentelemetry.OpenTelemetryTracing. 

8""" 

9 

10from contextlib import AbstractContextManager, nullcontext 

11from types import TracebackType 

12from typing import Optional 

13 

14from zenmqtt.instrumentation.guard import HookErrors 

15from zenmqtt.mqtt.publish import ( 

16 MQTTMessage, 

17 PublishAcknowledgement, 

18 PublishProperties, 

19) 

20 

21 

22class PublishTrace: 

23 """ 

24 The trace of one publish call. `properties` are sent instead of the 

25 properties of the call, e.g. with the trace context in user properties. 

26 """ 

27 

28 def __init__(self, properties: PublishProperties) -> None: 

29 self.properties = properties 

30 

31 def end( 

32 self, 

33 result: Optional[PublishAcknowledgement], 

34 error: Optional[BaseException], 

35 ) -> None: 

36 """ 

37 Called once when the publish call returns: `result` is the final 

38 acknowledgement (None for QoS 0), `error` is the exception of the call. 

39 """ 

40 

41 

42class Tracing: 

43 def start_publish( 

44 self, 

45 client_id: str, 

46 topic: str, 

47 payload: bytes, 

48 qos: int, 

49 retain: bool, 

50 properties: PublishProperties, 

51 ) -> PublishTrace: 

52 """Called when a publish call starts, before it waits for a connection.""" 

53 return PublishTrace(properties) 

54 

55 def trace( 

56 self, client_id: str, message: MQTTMessage 

57 ) -> AbstractContextManager[None]: 

58 """ 

59 Wraps the processing of a received message, see MQTTClient.trace(). 

60 """ 

61 return nullcontext() 

62 

63 

64class _GuardedPublishTrace(PublishTrace): 

65 def __init__(self, trace: PublishTrace, errors: HookErrors) -> None: 

66 super().__init__(trace.properties) 

67 self._trace = trace 

68 self._errors = errors 

69 

70 def end( 

71 self, 

72 result: Optional[PublishAcknowledgement], 

73 error: Optional[BaseException], 

74 ) -> None: 

75 try: 

76 self._trace.end(result, error) 

77 except Exception as exc: 

78 self._errors.report("PublishTrace.end", exc) 

79 

80 

81class _GuardedTrace(AbstractContextManager[None]): 

82 """ 

83 Exceptions of the tracing context are reported, and it never suppresses an 

84 exception of the processing. 

85 """ 

86 

87 def __init__( 

88 self, context: AbstractContextManager[None], errors: HookErrors 

89 ) -> None: 

90 self._context = context 

91 self._errors = errors 

92 self._entered = False 

93 

94 def __enter__(self) -> None: 

95 try: 

96 self._context.__enter__() 

97 self._entered = True 

98 except Exception as exc: 

99 self._errors.report("trace.__enter__", exc) 

100 

101 def __exit__( 

102 self, 

103 exc_type: Optional[type[BaseException]], 

104 exc: Optional[BaseException], 

105 traceback: Optional[TracebackType], 

106 ) -> None: 

107 if not self._entered: 

108 return 

109 

110 try: 

111 self._context.__exit__(exc_type, exc, traceback) 

112 except Exception as error: 

113 self._errors.report("trace.__exit__", error) 

114 

115 

116class _GuardedTracing(Tracing): 

117 """ 

118 Calls a tracing implementation, its exceptions are reported and never 

119 reach the client: e.g. a publish which the server acknowledged doesn't 

120 fail because its span couldn't be ended. 

121 """ 

122 

123 def __init__(self, tracing: Tracing) -> None: 

124 self.wrapped = tracing 

125 self._errors = HookErrors(type(tracing).__name__, "mqtt_tracing.hook_error") 

126 

127 @property 

128 def errors(self) -> dict[str, int]: 

129 """Number of exceptions per hook.""" 

130 return self._errors.errors 

131 

132 def start_publish( 

133 self, 

134 client_id: str, 

135 topic: str, 

136 payload: bytes, 

137 qos: int, 

138 retain: bool, 

139 properties: PublishProperties, 

140 ) -> PublishTrace: 

141 try: 

142 trace = self.wrapped.start_publish( 

143 client_id, topic, payload, qos, retain, properties 

144 ) 

145 except Exception as exc: 

146 self._errors.report("start_publish", exc) 

147 return PublishTrace(properties) 

148 

149 return _GuardedPublishTrace(trace, self._errors) 

150 

151 def trace( 

152 self, client_id: str, message: MQTTMessage 

153 ) -> AbstractContextManager[None]: 

154 try: 

155 context = self.wrapped.trace(client_id, message) 

156 except Exception as exc: 

157 self._errors.report("trace", exc) 

158 return nullcontext() 

159 

160 return _GuardedTrace(context, self._errors) 

161 

162 

163def guard_tracing(tracing: Optional[Tracing]) -> Optional[Tracing]: 

164 """ 

165 Wraps a tracing implementation, so its exceptions never reach the client; 

166 None and a guarded one are returned as they are. 

167 """ 

168 if tracing is None or isinstance(tracing, _GuardedTracing): 

169 return tracing 

170 

171 return _GuardedTracing(tracing)