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 }