2022-08-24 02:24:25 +00:00
|
|
|
package internal
|
|
|
|
|
|
|
|
import (
|
|
|
|
"bytes"
|
|
|
|
"fmt"
|
|
|
|
"io"
|
|
|
|
"sync"
|
2022-10-24 20:14:32 +00:00
|
|
|
|
|
|
|
"github.com/go-logr/logr"
|
|
|
|
"github.com/go-logr/logr/funcr"
|
2022-08-24 02:24:25 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
type WriterMode uint
|
|
|
|
|
|
|
|
const (
|
|
|
|
WriterModeStreamAndBuffer WriterMode = iota
|
|
|
|
WriterModeBufferOnly
|
|
|
|
)
|
|
|
|
|
|
|
|
type WriterInterface interface {
|
|
|
|
io.Writer
|
|
|
|
|
|
|
|
Truncate()
|
|
|
|
Bytes() []byte
|
|
|
|
}
|
|
|
|
|
2023-02-01 17:06:36 +00:00
|
|
|
//Writer implements WriterInterface and GinkgoWriterInterface
|
2022-08-24 02:24:25 +00:00
|
|
|
type Writer struct {
|
|
|
|
buffer *bytes.Buffer
|
|
|
|
outWriter io.Writer
|
|
|
|
lock *sync.Mutex
|
|
|
|
mode WriterMode
|
|
|
|
|
|
|
|
teeWriters []io.Writer
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewWriter(outWriter io.Writer) *Writer {
|
|
|
|
return &Writer{
|
2023-02-01 17:06:36 +00:00
|
|
|
buffer: &bytes.Buffer{},
|
|
|
|
lock: &sync.Mutex{},
|
|
|
|
outWriter: outWriter,
|
|
|
|
mode: WriterModeStreamAndBuffer,
|
2022-08-24 02:24:25 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) SetMode(mode WriterMode) {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
w.mode = mode
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Write(b []byte) (n int, err error) {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
|
|
|
|
for _, teeWriter := range w.teeWriters {
|
|
|
|
teeWriter.Write(b)
|
|
|
|
}
|
|
|
|
|
|
|
|
if w.mode == WriterModeStreamAndBuffer {
|
2023-02-01 17:06:36 +00:00
|
|
|
w.outWriter.Write(b)
|
2022-08-24 02:24:25 +00:00
|
|
|
}
|
|
|
|
return w.buffer.Write(b)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Truncate() {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
w.buffer.Reset()
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Bytes() []byte {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
b := w.buffer.Bytes()
|
|
|
|
copied := make([]byte, len(b))
|
|
|
|
copy(copied, b)
|
|
|
|
return copied
|
|
|
|
}
|
|
|
|
|
2023-02-01 17:06:36 +00:00
|
|
|
//GinkgoWriterInterface
|
2022-08-24 02:24:25 +00:00
|
|
|
func (w *Writer) TeeTo(writer io.Writer) {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
|
|
|
|
w.teeWriters = append(w.teeWriters, writer)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) ClearTeeWriters() {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
|
|
|
|
w.teeWriters = []io.Writer{}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Print(a ...interface{}) {
|
|
|
|
fmt.Fprint(w, a...)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Printf(format string, a ...interface{}) {
|
|
|
|
fmt.Fprintf(w, format, a...)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Println(a ...interface{}) {
|
|
|
|
fmt.Fprintln(w, a...)
|
|
|
|
}
|
2022-10-24 20:14:32 +00:00
|
|
|
|
|
|
|
func GinkgoLogrFunc(writer *Writer) logr.Logger {
|
|
|
|
return funcr.New(func(prefix, args string) {
|
|
|
|
writer.Printf("%s", args)
|
|
|
|
}, funcr.Options{})
|
|
|
|
}
|