Coverage for fluent/asyncsender.py: 100%

69 statements  

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

1import threading 

2from queue import Queue, Full, Empty 

3 

4from fluent import sender 

5from fluent.sender import EventTime 

6 

7__all__ = ["EventTime", "FluentSender"] 

8 

9DEFAULT_QUEUE_MAXSIZE = 100 

10DEFAULT_QUEUE_CIRCULAR = False 

11 

12_TOMBSTONE = object() 

13 

14_global_sender = None 

15 

16 

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

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

19 global _global_sender 

20 _global_sender = sender 

21 

22 

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

24 global _global_sender 

25 _global_sender = FluentSender(tag, **kwargs) 

26 

27 

28def get_global_sender(): # pragma: no cover 

29 return _global_sender 

30 

31 

32def close(): # pragma: no cover 

33 get_global_sender().close() 

34 

35 

36class FluentSender(sender.FluentSender): 

37 def __init__( 

38 self, 

39 tag, 

40 host="localhost", 

41 port=24224, 

42 bufmax=1 * 1024 * 1024, 

43 timeout=3.0, 

44 verbose=False, 

45 buffer_overflow_handler=None, 

46 nanosecond_precision=False, 

47 msgpack_kwargs=None, 

48 queue_maxsize=DEFAULT_QUEUE_MAXSIZE, 

49 queue_circular=DEFAULT_QUEUE_CIRCULAR, 

50 queue_overflow_handler=None, 

51 **kwargs, 

52 ): 

53 """ 

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

55 """ 

56 super().__init__( 

57 tag=tag, 

58 host=host, 

59 port=port, 

60 bufmax=bufmax, 

61 timeout=timeout, 

62 verbose=verbose, 

63 buffer_overflow_handler=buffer_overflow_handler, 

64 nanosecond_precision=nanosecond_precision, 

65 msgpack_kwargs=msgpack_kwargs, 

66 **kwargs, 

67 ) 

68 self._queue_maxsize = queue_maxsize 

69 self._queue_circular = queue_circular 

70 if queue_circular and queue_overflow_handler: 

71 self._queue_overflow_handler = queue_overflow_handler 

72 else: 

73 self._queue_overflow_handler = self._queue_overflow_handler_default 

74 

75 self._thread_guard = ( 

76 threading.Event() 

77 ) # This ensures visibility across all variables 

78 self._closed = False 

79 

80 self._queue = Queue(maxsize=queue_maxsize) 

81 self._send_thread = threading.Thread( 

82 target=self._send_loop, name="AsyncFluentSender %d" % id(self) 

83 ) 

84 self._send_thread.daemon = True 

85 self._send_thread.start() 

86 

87 def close(self, flush=True): 

88 with self.lock: 

89 if self._closed: 

90 return 

91 self._closed = True 

92 if not flush: 

93 while True: 

94 try: 

95 self._queue.get(block=False) 

96 except Empty: 

97 break 

98 self._queue.put(_TOMBSTONE) 

99 self._send_thread.join() 

100 

101 @property 

102 def queue_maxsize(self): 

103 return self._queue_maxsize 

104 

105 @property 

106 def queue_blocking(self): 

107 return not self._queue_circular 

108 

109 @property 

110 def queue_circular(self): 

111 return self._queue_circular 

112 

113 def _send(self, bytes_): 

114 with self.lock: 

115 if self._closed: 

116 return False 

117 if self._queue_circular and self._queue.full(): 

118 # discard oldest 

119 try: 

120 discarded_bytes = self._queue.get(block=False) 

121 except Empty: # pragma: no cover 

122 pass 

123 else: 

124 self._queue_overflow_handler(discarded_bytes) 

125 try: 

126 self._queue.put(bytes_, block=(not self._queue_circular)) 

127 except Full: # pragma: no cover 

128 return False # this actually can't happen 

129 

130 return True 

131 

132 def _send_loop(self): 

133 send_internal = super()._send_internal 

134 

135 try: 

136 while True: 

137 bytes_ = self._queue.get(block=True) 

138 if bytes_ is _TOMBSTONE: 

139 break 

140 

141 send_internal(bytes_) 

142 finally: 

143 self._close() 

144 

145 def _queue_overflow_handler_default(self, discarded_bytes): 

146 pass 

147 

148 def __exit__(self, exc_type, exc_val, exc_tb): 

149 self.close()