1package gexec
2
3import (
4	"io"
5	"sync"
6)
7
8/*
9PrefixedWriter wraps an io.Writer, emiting the passed in prefix at the beginning of each new line.
10This can be useful when running multiple gexec.Sessions concurrently - you can prefix the log output of each
11session by passing in a PrefixedWriter:
12
13gexec.Start(cmd, NewPrefixedWriter("[my-cmd] ", GinkgoWriter), NewPrefixedWriter("[my-cmd] ", GinkgoWriter))
14*/
15type PrefixedWriter struct {
16	prefix        []byte
17	writer        io.Writer
18	lock          *sync.Mutex
19	atStartOfLine bool
20}
21
22func NewPrefixedWriter(prefix string, writer io.Writer) *PrefixedWriter {
23	return &PrefixedWriter{
24		prefix:        []byte(prefix),
25		writer:        writer,
26		lock:          &sync.Mutex{},
27		atStartOfLine: true,
28	}
29}
30
31func (w *PrefixedWriter) Write(b []byte) (int, error) {
32	w.lock.Lock()
33	defer w.lock.Unlock()
34
35	toWrite := []byte{}
36
37	for _, c := range b {
38		if w.atStartOfLine {
39			toWrite = append(toWrite, w.prefix...)
40		}
41
42		toWrite = append(toWrite, c)
43
44		w.atStartOfLine = c == '\n'
45	}
46
47	_, err := w.writer.Write(toWrite)
48	if err != nil {
49		return 0, err
50	}
51
52	return len(b), nil
53}
54