github.com/mongodb/grip@v0.0.0-20240213223901-f906268d82b9/send/writer.go (about)

     1  package send
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"io"
     7  	"sync"
     8  	"unicode"
     9  
    10  	"github.com/mongodb/grip/level"
    11  	"github.com/mongodb/grip/message"
    12  )
    13  
    14  // WriterSender wraps another sender and also provides an io.Writer.
    15  // (and because Sender is an io.Closer) the type also implements
    16  // io.WriteCloser.
    17  type WriterSender struct {
    18  	Sender
    19  	writer   *bufio.Writer
    20  	buffer   *bytes.Buffer
    21  	priority level.Priority
    22  	mu       sync.Mutex
    23  }
    24  
    25  // NewWriterSender wraps another sender and also provides an io.Writer.
    26  // (and because Sender is an io.Closer) the type also implements
    27  // io.WriteCloser.
    28  //
    29  // While WriteSender itself implements Sender, it also provides a
    30  // Writer method, which allows you to use this Sender to capture
    31  // file-like write operations.
    32  //
    33  // Data sent via the Write method is buffered internally until its
    34  // passed a byte slice that ends with the new line character. If the
    35  // string form of the bytes passed to the write method (including all
    36  // buffered messages) is only whitespace, then it is not sent.
    37  //
    38  // If there are any bytes in the buffer when the Close method is
    39  // called, this sender flushes the buffer. WriterSender does not own the
    40  // underlying Sender, so users are responsible for closing the underlying Sender
    41  // if/when it is appropriate to release its resources.
    42  func NewWriterSender(s Sender) *WriterSender { return MakeWriterSender(s, s.Level().Default) }
    43  
    44  // MakeWriterSender returns an sender interface that also implements
    45  // io.Writer. Specify a priority used as the level for messages sent.
    46  func MakeWriterSender(s Sender, p level.Priority) *WriterSender {
    47  	buffer := new(bytes.Buffer)
    48  
    49  	return &WriterSender{
    50  		Sender:   s,
    51  		priority: p,
    52  		writer:   bufio.NewWriter(buffer),
    53  		buffer:   buffer,
    54  	}
    55  }
    56  
    57  // Write captures a sequence of bytes to the send interface. It never errors.
    58  func (s *WriterSender) Write(p []byte) (int, error) {
    59  	s.mu.Lock()
    60  	defer s.mu.Unlock()
    61  	n, err := s.writer.Write(p)
    62  	if err != nil {
    63  		return n, err
    64  	}
    65  	_ = s.writer.Flush()
    66  
    67  	if s.buffer.Len() > 80 {
    68  		err = s.doSend()
    69  	}
    70  
    71  	return n, err
    72  }
    73  
    74  func (s *WriterSender) doSend() error {
    75  	for {
    76  		line, err := s.buffer.ReadBytes('\n')
    77  		if err == io.EOF {
    78  			s.buffer.Write(line)
    79  			return nil
    80  		}
    81  
    82  		lncp := make([]byte, len(line))
    83  		copy(lncp, line)
    84  
    85  		if err == nil {
    86  			s.Send(message.NewBytesMessage(s.priority, bytes.TrimRightFunc(lncp, unicode.IsSpace)))
    87  			continue
    88  		}
    89  
    90  		s.Send(message.NewBytesMessage(s.priority, bytes.TrimRightFunc(lncp, unicode.IsSpace)))
    91  		return err
    92  	}
    93  }
    94  
    95  // Close writes any buffered messages to the underlying Sender. This does
    96  // not close the underlying sender.
    97  func (s *WriterSender) Close() error {
    98  	s.mu.Lock()
    99  	defer s.mu.Unlock()
   100  
   101  	if err := s.writer.Flush(); err != nil {
   102  		return err
   103  	}
   104  
   105  	s.Send(message.NewBytesMessage(s.priority, bytes.TrimRightFunc(s.buffer.Bytes(), unicode.IsSpace)))
   106  	s.buffer.Reset()
   107  	s.writer.Reset(s.buffer)
   108  	return nil
   109  }