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  }