Coverage for fluent/sender.py: 99%

162 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-02-29 15:22 +0900

1import errno 

2import socket 

3import struct 

4import threading 

5import time 

6import traceback 

7 

8import msgpack 

9 

10_global_sender = None 

11 

12 

13def _set_global_sender(sender): # pragma: no cover 

14 """[For testing] Function to set global sender directly""" 

15 global _global_sender 

16 _global_sender = sender 

17 

18 

19def setup(tag, **kwargs): # pragma: no cover 

20 global _global_sender 

21 _global_sender = FluentSender(tag, **kwargs) 

22 

23 

24def get_global_sender(): # pragma: no cover 

25 return _global_sender 

26 

27 

28def close(): # pragma: no cover 

29 get_global_sender().close() 

30 

31 

32class EventTime(msgpack.ExtType): 

33 def __new__(cls, timestamp, nanoseconds=None): 

34 if nanoseconds is None: 

35 seconds = int(timestamp) 

36 nanoseconds = int(timestamp % 1 * 10**9) 

37 else: 

38 seconds = int(timestamp) 

39 return super().__new__( 

40 cls, 

41 code=0, 

42 data=struct.pack(">II", seconds, nanoseconds), 

43 ) 

44 

45 @classmethod 

46 def from_unix_nano(cls, unix_nano): 

47 seconds, nanos = divmod(unix_nano, 10**9) 

48 return cls(seconds, nanos) 

49 

50 

51class FluentSender: 

52 def __init__( 

53 self, 

54 tag, 

55 host="localhost", 

56 port=24224, 

57 bufmax=1 * 1024 * 1024, 

58 timeout=3.0, 

59 verbose=False, 

60 buffer_overflow_handler=None, 

61 nanosecond_precision=False, 

62 msgpack_kwargs=None, 

63 *, 

64 forward_packet_error=True, 

65 **kwargs, 

66 ): 

67 """ 

68 :param kwargs: This kwargs argument is not used in __init__. This will be removed in the next major version. 

69 """ 

70 self.tag = tag 

71 self.host = host 

72 self.port = port 

73 self.bufmax = bufmax 

74 self.timeout = timeout 

75 self.verbose = verbose 

76 self.buffer_overflow_handler = buffer_overflow_handler 

77 self.nanosecond_precision = nanosecond_precision 

78 self.forward_packet_error = forward_packet_error 

79 self.msgpack_kwargs = {} if msgpack_kwargs is None else msgpack_kwargs 

80 

81 self.socket = None 

82 self.pendings = None 

83 self.lock = threading.Lock() 

84 self._closed = False 

85 self._last_error_threadlocal = threading.local() 

86 

87 def emit(self, label, data): 

88 if self.nanosecond_precision: 

89 cur_time = EventTime.from_unix_nano(time.time_ns()) 

90 else: 

91 cur_time = int(time.time()) 

92 return self.emit_with_time(label, cur_time, data) 

93 

94 def emit_with_time(self, label, timestamp, data): 

95 try: 

96 bytes_ = self._make_packet(label, timestamp, data) 

97 except Exception as e: 

98 if not self.forward_packet_error: 98 ↛ 99line 98 didn't jump to line 99, because the condition on line 98 was never true

99 raise 

100 self.last_error = e 

101 bytes_ = self._make_packet( 

102 label, 

103 timestamp, 

104 { 

105 "level": "CRITICAL", 

106 "message": "Can't output to log", 

107 "traceback": traceback.format_exc(), 

108 }, 

109 ) 

110 return self._send(bytes_) 

111 

112 @property 

113 def last_error(self): 

114 return getattr(self._last_error_threadlocal, "exception", None) 

115 

116 @last_error.setter 

117 def last_error(self, err): 

118 self._last_error_threadlocal.exception = err 

119 

120 def clear_last_error(self, _thread_id=None): 

121 if hasattr(self._last_error_threadlocal, "exception"): 

122 delattr(self._last_error_threadlocal, "exception") 

123 

124 def close(self): 

125 with self.lock: 

126 if self._closed: 

127 return 

128 self._closed = True 

129 if self.pendings: 

130 try: 

131 self._send_data(self.pendings) 

132 except Exception: 

133 self._call_buffer_overflow_handler(self.pendings) 

134 

135 self._close() 

136 self.pendings = None 

137 

138 def _make_packet(self, label, timestamp, data): 

139 if label: 

140 tag = f"{self.tag}.{label}" if self.tag else label 

141 else: 

142 tag = self.tag 

143 if self.nanosecond_precision and isinstance(timestamp, float): 

144 timestamp = EventTime(timestamp) 

145 packet = (tag, timestamp, data) 

146 if self.verbose: 

147 print(packet) 

148 return msgpack.packb(packet, **self.msgpack_kwargs) 

149 

150 def _send(self, bytes_): 

151 with self.lock: 

152 if self._closed: 

153 return False 

154 return self._send_internal(bytes_) 

155 

156 def _send_internal(self, bytes_): 

157 # buffering 

158 if self.pendings: 

159 self.pendings += bytes_ 

160 bytes_ = self.pendings 

161 

162 try: 

163 self._send_data(bytes_) 

164 

165 # send finished 

166 self.pendings = None 

167 

168 return True 

169 except OSError as e: 

170 self.last_error = e 

171 

172 # close socket 

173 self._close() 

174 

175 # clear buffer if it exceeds max buffer size 

176 if self.pendings and (len(self.pendings) > self.bufmax): 

177 self._call_buffer_overflow_handler(self.pendings) 

178 self.pendings = None 

179 else: 

180 self.pendings = bytes_ 

181 

182 return False 

183 

184 def _check_recv_side(self): 

185 try: 

186 self.socket.settimeout(0.0) 

187 try: 

188 recvd = self.socket.recv(4096) 

189 except OSError as recv_e: 

190 if recv_e.errno != errno.EWOULDBLOCK: 

191 raise 

192 return 

193 

194 if recvd == b"": 

195 raise OSError(errno.EPIPE, "Broken pipe") 

196 finally: 

197 self.socket.settimeout(self.timeout) 

198 

199 def _send_data(self, bytes_): 

200 # reconnect if possible 

201 self._reconnect() 

202 # send message 

203 bytes_to_send = len(bytes_) 

204 bytes_sent = 0 

205 self._check_recv_side() 

206 while bytes_sent < bytes_to_send: 

207 sent = self.socket.send(bytes_[bytes_sent:]) 

208 if sent == 0: 

209 raise OSError(errno.EPIPE, "Broken pipe") 

210 bytes_sent += sent 

211 self._check_recv_side() 

212 

213 def _reconnect(self): 

214 if not self.socket: 

215 try: 

216 if self.host.startswith("unix://"): 

217 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) 

218 sock.settimeout(self.timeout) 

219 sock.connect(self.host[len("unix://") :]) 

220 else: 

221 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 

222 sock.settimeout(self.timeout) 

223 # This might be controversial and may need to be removed 

224 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) 

225 sock.connect((self.host, self.port)) 

226 except Exception as e: 

227 try: 

228 sock.close() 

229 except Exception: # pragma: no cover 

230 pass 

231 raise e 

232 else: 

233 self.socket = sock 

234 

235 def _call_buffer_overflow_handler(self, pending_events): 

236 try: 

237 if self.buffer_overflow_handler: 

238 self.buffer_overflow_handler(pending_events) 

239 except Exception: 

240 # User should care any exception in handler 

241 pass 

242 

243 def _close(self): 

244 try: 

245 sock = self.socket 

246 if sock: 

247 try: 

248 try: 

249 sock.shutdown(socket.SHUT_RDWR) 

250 except OSError: # pragma: no cover 

251 pass 

252 finally: 

253 try: 

254 sock.close() 

255 except OSError: # pragma: no cover 

256 pass 

257 finally: 

258 self.socket = None 

259 

260 def __enter__(self): 

261 return self 

262 

263 def __exit__(self, typ, value, traceback): 

264 try: 

265 self.close() 

266 except Exception as e: # pragma: no cover 

267 self.last_error = e