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 }