blake.io/pqx@v0.2.2-0.20231231055241-83f2254c0a07/internal/logplex/logplex.go (about)

     1  package logplex
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"sync"
     8  	"unicode"
     9  )
    10  
    11  type Logplex struct {
    12  	Sink io.Writer
    13  
    14  	// if nil, all lines go to Drain
    15  	Split func(line []byte) (key, message []byte)
    16  
    17  	lineBuf bytes.Buffer
    18  
    19  	mu       sync.Mutex
    20  	sinks    map[string]io.Writer
    21  	lastSeen []byte
    22  }
    23  
    24  func (lp *Logplex) Watch(prefix string, w io.Writer) {
    25  	lp.mu.Lock()
    26  	defer lp.mu.Unlock()
    27  	if lp.sinks == nil {
    28  		lp.sinks = map[string]io.Writer{}
    29  	}
    30  	lp.sinks[prefix] = w
    31  }
    32  
    33  var newline = []byte{'\n'}
    34  
    35  // Write writes p to the underlying buffer and flushes each newline to their
    36  // corresponding sinks.
    37  func (lp *Logplex) Write(p []byte) (int, error) {
    38  	lp.mu.Lock()
    39  	defer lp.mu.Unlock()
    40  
    41  	p0 := p
    42  	for {
    43  		before, after, hasNewline := bytes.Cut(p, newline)
    44  		lp.lineBuf.Write(before)
    45  		if hasNewline {
    46  			lp.lineBuf.Write(newline)
    47  			if err := lp.flushLocked(); err != nil {
    48  				return 0, err
    49  			}
    50  			p = after
    51  		} else {
    52  			return len(p0), nil
    53  		}
    54  	}
    55  }
    56  
    57  // caller must hold mu
    58  func (lp *Logplex) flushLocked() error {
    59  	defer lp.lineBuf.Reset()
    60  
    61  	// There is only one line in the buffer. Check the prefix and send to
    62  	// the appropriate sink.
    63  	if lp.Split == nil {
    64  		_, err := lp.Sink.Write(lp.lineBuf.Bytes())
    65  		return err
    66  	}
    67  
    68  	line := lp.lineBuf.Bytes()
    69  	sent, err := lp.sendLine(line)
    70  	if err != nil {
    71  		return err
    72  	}
    73  	if sent {
    74  		return nil
    75  	}
    76  	_, err = lp.Sink.Write(lp.lineBuf.Bytes())
    77  	return err
    78  }
    79  
    80  func (lp *Logplex) Unwatch(prefix string) {
    81  	lp.mu.Lock()
    82  	defer lp.mu.Unlock()
    83  	delete(lp.sinks, prefix)
    84  }
    85  
    86  // Flush flushes the any underlying buffered contents to any corresponding sink.
    87  //
    88  // The contents flushed may not be a complete line, or have enough data to
    89  // determine the proper sink and instead send to Sink.
    90  //
    91  // Flush should only be called when no more data will be written to Write.
    92  func (lp *Logplex) Flush() error {
    93  	lp.mu.Lock()
    94  	defer lp.mu.Unlock()
    95  	return lp.flushLocked()
    96  }
    97  
    98  // caller must hold mu
    99  func (lp *Logplex) sendLine(line []byte) (sent bool, err error) {
   100  	key, message := lp.Split(line)
   101  	if isContinuation(message) {
   102  		key = lp.lastSeen
   103  	}
   104  	for prefix, w := range lp.sinks {
   105  		if string(key) == prefix {
   106  			lp.lastSeen = append(lp.lastSeen[:0], key...)
   107  			_, err := w.Write(message)
   108  			return true, err
   109  		}
   110  	}
   111  	return false, nil
   112  }
   113  
   114  type logfWriter struct {
   115  	logf func(string, ...any)
   116  }
   117  
   118  func LogfWriter(logf func(string, ...any)) io.Writer {
   119  	return &logfWriter{logf}
   120  }
   121  
   122  func (w *logfWriter) Write(p []byte) (int, error) {
   123  	w.logf(string(p))
   124  	return len(p), nil
   125  }
   126  
   127  func LogfFromWriter(w io.Writer) func(string, ...any) {
   128  	return func(format string, args ...any) {
   129  		fmt.Fprintf(w, format, args...)
   130  	}
   131  }
   132  
   133  func isContinuation(line []byte) bool {
   134  	return len(line) > 0 && unicode.IsSpace(rune(line[0]))
   135  }