github.com/mongodb/grip@v0.0.0-20240213223901-f906268d82b9/send/buffered.go (about) 1 package send 2 3 import ( 4 "context" 5 "errors" 6 "sync" 7 "time" 8 9 "github.com/mongodb/grip/message" 10 ) 11 12 const ( 13 minFlushInterval = 5 * time.Second 14 defaultFlushInterval = time.Minute 15 defaultBufferSize = 100 16 ) 17 18 // BufferedSenderOptions configure the buffered sender. 19 type BufferedSenderOptions struct { 20 // FlushInterval is the maximum duration between flushes. The buffer will automatically 21 // flush if it hasn't flushed within the past FlushInterval. 22 // FlushInterval is 1 minute by default. The minimum interval is 5 seconds. 23 // If an interval of less than 5 seconds is specified it is set to 5 seconds. 24 FlushInterval time.Duration 25 // BufferSize is the threshold at which the buffer is flushed. 26 // The size is 100 by default. 27 BufferSize int 28 } 29 30 func (opts *BufferedSenderOptions) validate() error { 31 if opts.FlushInterval < 0 { 32 return errors.New("FlushInterval cannot be negative") 33 } 34 if opts.BufferSize < 0 { 35 return errors.New("BufferSize cannot be negative") 36 } 37 38 if opts.FlushInterval == 0 { 39 opts.FlushInterval = defaultFlushInterval 40 } else if opts.FlushInterval < minFlushInterval { 41 opts.FlushInterval = minFlushInterval 42 } 43 44 if opts.BufferSize == 0 { 45 opts.BufferSize = defaultBufferSize 46 } 47 48 return nil 49 } 50 51 type bufferedSender struct { 52 mu sync.Mutex 53 cancel context.CancelFunc 54 buffer []message.Composer 55 size int 56 lastFlush time.Time 57 closed bool 58 59 Sender 60 } 61 62 // NewBufferedSender provides a Sender implementation that wraps an existing 63 // Sender sending messages in batches, on a specified buffer size or after an 64 // interval has passed. 65 // 66 // This Sender does not own the underlying Sender, so users are responsible for 67 // closing the underlying Sender if/when it is appropriate to release its 68 // resources. 69 func NewBufferedSender(ctx context.Context, sender Sender, opts BufferedSenderOptions) (Sender, error) { 70 if err := opts.validate(); err != nil { 71 return nil, err 72 } 73 74 ctx, cancel := context.WithCancel(ctx) 75 s := &bufferedSender{ 76 Sender: sender, 77 cancel: cancel, 78 buffer: []message.Composer{}, 79 size: opts.BufferSize, 80 } 81 82 go s.intervalFlush(ctx, opts.FlushInterval) 83 84 return s, nil 85 } 86 87 func (s *bufferedSender) Send(msg message.Composer) { 88 if !s.Level().ShouldLog(msg) { 89 return 90 } 91 92 s.mu.Lock() 93 defer s.mu.Unlock() 94 95 if s.closed { 96 return 97 } 98 99 s.buffer = append(s.buffer, msg) 100 if len(s.buffer) >= s.size { 101 s.flush() 102 } 103 } 104 105 func (s *bufferedSender) Flush(_ context.Context) error { 106 s.mu.Lock() 107 defer s.mu.Unlock() 108 109 if !s.closed { 110 s.flush() 111 } 112 113 return nil 114 } 115 116 // Close writes any buffered messages to the underlying Sender. This does not 117 // close the underlying sender. 118 func (s *bufferedSender) Close() error { 119 s.mu.Lock() 120 defer s.mu.Unlock() 121 122 if s.closed { 123 return nil 124 } 125 126 s.cancel() 127 if len(s.buffer) > 0 { 128 s.flush() 129 } 130 s.closed = true 131 132 return nil 133 } 134 135 func (s *bufferedSender) intervalFlush(ctx context.Context, interval time.Duration) { 136 timer := time.NewTimer(interval) 137 defer timer.Stop() 138 139 for { 140 select { 141 case <-ctx.Done(): 142 return 143 case <-timer.C: 144 s.mu.Lock() 145 if len(s.buffer) > 0 && time.Since(s.lastFlush) >= interval { 146 s.flush() 147 } 148 s.mu.Unlock() 149 _ = timer.Reset(interval) 150 } 151 } 152 } 153 154 func (s *bufferedSender) flush() { 155 if len(s.buffer) == 1 { 156 s.Sender.Send(s.buffer[0]) 157 } else { 158 s.Sender.Send(message.NewGroupComposer(s.buffer)) 159 } 160 161 s.buffer = []message.Composer{} 162 s.lastFlush = time.Now() 163 }