mirror of
https://github.com/ceph/ceph-csi.git
synced 2024-12-23 05:20:19 +00:00
104 lines
1.7 KiB
Go
104 lines
1.7 KiB
Go
|
package internal
|
||
|
|
||
|
import (
|
||
|
"bytes"
|
||
|
"fmt"
|
||
|
"io"
|
||
|
"sync"
|
||
|
)
|
||
|
|
||
|
type WriterMode uint
|
||
|
|
||
|
const (
|
||
|
WriterModeStreamAndBuffer WriterMode = iota
|
||
|
WriterModeBufferOnly
|
||
|
)
|
||
|
|
||
|
type WriterInterface interface {
|
||
|
io.Writer
|
||
|
|
||
|
Truncate()
|
||
|
Bytes() []byte
|
||
|
}
|
||
|
|
||
|
//Writer implements WriterInterface and GinkgoWriterInterface
|
||
|
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{
|
||
|
buffer: &bytes.Buffer{},
|
||
|
lock: &sync.Mutex{},
|
||
|
outWriter: outWriter,
|
||
|
mode: WriterModeStreamAndBuffer,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
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 {
|
||
|
w.outWriter.Write(b)
|
||
|
}
|
||
|
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
|
||
|
}
|
||
|
|
||
|
//GinkgoWriterInterface
|
||
|
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...)
|
||
|
}
|