Dont close stdout (#289)

* Let the (code)user handle the closing of passed writers when using loginterceptor

* Simplify loginterceptor and fix tests

* Unblock run goroutine when channels are full

* Filter in scan to preserve channel resources

* Log interception buffer limit errors with Debugf

Author: Badlazzor

loginterceptor/loginterceptor.go

@@ -11,31 +11,55 @@ import (
 
 // PrefixInterceptor intercept writes: if a line begins with prefix, it will be written to
 // both writers. Partial writes without newline are buffered until a newline.
+//
+// Note: Callers are responsible for closing intercepted and target writers that implement io.Closer
 type PrefixInterceptor struct {
-	prefixRegexp *regexp.Regexp
-	intercepted  *NonBlockingWriter
-	target       *NonBlockingWriter
-	logger       log.Logger
+	prefixRegexp  *regexp.Regexp
+	targetCh      chan string
+	interceptedCh chan string
+	logger        log.Logger
 
 	// internal pipe and goroutine to scan and route
 	internalReader *io.PipeReader
 	internalWriter *io.PipeWriter
 
-	// close once
+	// closing
 	closeOnce sync.Once
 	closeErr  error
+
+	// TargetDelivered receives a single value when the target goroutine
+	// finishes processing all messages. Callers should consume from this
+	// channel to avoid goroutine leaks, or use Close() and ignore it.
+	TargetDelivered <-chan struct{}
+	// InterceptedDelivered receives a single value when the intercepted goroutine
+	// finishes processing all messages. Callers should consume from this
+	// channel to avoid goroutine leaks, or use Close() and ignore it.
+	InterceptedDelivered <-chan struct{}
 }
 
 // NewPrefixInterceptor returns an io.WriteCloser. Writes are based on line prefix.
+//
+// Note: Callers are responsible for closing intercepted and target writers that implement io.Closer
 func NewPrefixInterceptor(prefixRegexp *regexp.Regexp, intercepted, target io.Writer, logger log.Logger) *PrefixInterceptor {
 	pipeReader, pipeWriter := io.Pipe()
+
+	targetCh := make(chan string, 10000)
+	targetDoneCh := make(chan struct{}, 1)
+	interceptedCh := make(chan string, 10000)
+	interceptedDoneCh := make(chan struct{}, 1)
+
+	go sendingTo(targetCh, targetDoneCh, target, logger)
+	go sendingTo(interceptedCh, interceptedDoneCh, intercepted, logger)
+
 	interceptor := &PrefixInterceptor{
-		prefixRegexp:   prefixRegexp,
-		intercepted:    NewNonBlockingWriter(intercepted, logger),
-		target:         NewNonBlockingWriter(target, logger),
-		logger:         logger,
-		internalReader: pipeReader,
-		internalWriter: pipeWriter,
+		prefixRegexp:         prefixRegexp,
+		targetCh:             targetCh,
+		interceptedCh:        interceptedCh,
+		logger:               logger,
+		internalReader:       pipeReader,
+		internalWriter:       pipeWriter,
+		TargetDelivered:      targetDoneCh,
+		InterceptedDelivered: interceptedDoneCh,
 	}
 	go interceptor.run()
 	return interceptor
@@ -55,15 +79,11 @@ func (i *PrefixInterceptor) Close() error {
 }
 
 func (i *PrefixInterceptor) closeAfterRun() {
-	if err := i.intercepted.Close(); err != nil {
-		i.logger.Errorf("intercepted writer: %v", err)
-	}
-	if err := i.target.Close(); err != nil {
-		i.logger.Errorf("target writer: %v", err)
-	}
 	if err := i.internalReader.Close(); err != nil {
 		i.logger.Errorf("internal reader: %v", err)
 	}
+	close(i.targetCh)
+	close(i.interceptedCh)
 }
 
 // run reads lines (and partial final chunk) and writes them.
@@ -79,18 +99,21 @@ func (i *PrefixInterceptor) run() {
 	for scanner.Scan() {
 		line := scanner.Text() // note: newline removed
 		// re-append newline to preserve same output format
-		outLine := line + "\n"
+		logLine := line + "\n"
+
+		select {
+		case i.targetCh <- logLine:
+		default:
+			i.logger.Debugf("target channel full, dropping message")
+		}
 
-		// Write to intercepted channel if matching regexp
 		if i.prefixRegexp.MatchString(line) {
-			if _, err := io.WriteString(i.intercepted, outLine); err != nil {
-				i.logger.Errorf("intercept writer error: %v", err)
+			select {
+			case i.interceptedCh <- logLine:
+			default:
+				i.logger.Debugf("intercepted channel full, dropping message")
 			}
 		}
-		// Always write to target channel
-		if _, err := io.WriteString(i.target, outLine); err != nil {
-			i.logger.Errorf("writer error: %v", err)
-		}
 	}
 
 	// handle any scanner error
@@ -99,52 +122,18 @@ func (i *PrefixInterceptor) run() {
 	}
 }
 
-// NonBlockingWriter is an io.Writer that writes to a wrapped io.Writer in a non-blocking way.
-type NonBlockingWriter struct {
-	channel chan []byte
-	wrapped io.Writer
-	logger  log.Logger
-}
-
-// NewNonBlockingWriter creates a new NonBlockingWriter.
-func NewNonBlockingWriter(w io.Writer, logger log.Logger) *NonBlockingWriter {
-	writer := &NonBlockingWriter{
-		channel: make(chan []byte, 10000), // buffered channel to avoid blocking
-		wrapped: w,
-		logger:  logger,
-	}
-	go writer.Run()
-	return writer
-}
-
-// Write implements io.Writer. It writes into an internal pipe which the interceptor goroutine consumes.
-func (i *NonBlockingWriter) Write(p []byte) (int, error) {
-	select {
-	case i.channel <- p:
-		return len(p), nil
-	default:
-		i.logger.Debugf("buffer full, dropping log")
-		return 0, nil
-	}
-}
-
-// Close stops the interceptor and closes the pipe.
-func (i *NonBlockingWriter) Close() error {
-	close(i.channel)
-	return nil
-}
-
-// Run consumes the channel and writes to the wrapped writer.
-func (i *NonBlockingWriter) Run() {
-	for msg := range i.channel {
-		if _, err := i.wrapped.Write(msg); err != nil {
-			i.logger.Errorf("NonBlockingWriter: wrapped writer error: %v", err)
+func sendingTo(
+	srcCh <-chan string,
+	done chan<- struct{},
+	writer io.Writer,
+	logger log.Logger,
+) {
+	for msg := range srcCh {
+		if _, err := io.WriteString(writer, msg); err != nil {
+			logger.Errorf(" writer error: %v", err)
 		}
 	}
 
-	if closer, ok := i.wrapped.(io.Closer); ok {
-		if err := closer.Close(); err != nil {
-			i.logger.Errorf("NonBlockingWriter: closing wrapped writer: %v", err)
-		}
-	}
+	done <- struct{}{}
+	close(done)
 }

loginterceptor/loginterceptor_test.go

@@ -1,7 +1,6 @@
 package loginterceptor_test
 
 import (
-	"bytes"
 	"io"
 	"regexp"
 	"sync"
@@ -12,45 +11,42 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+const (
+	msg1 = "Log message without prefix\n"
+	msg2 = "[Bitrise Analytics] Log message with prefixs\n"
+	msg3 = "[Bitrise Build Cache] Log message with prefix\n"
+	msg4 = "Stuff [Bitrise Build Cache] Log message without prefix\n"
+)
+
 func TestPrefixInterceptor(t *testing.T) {
-	interceptReader, interceptWriter := io.Pipe()
-	targetReader, targetWriter := io.Pipe()
+	interceptedMsgs := NewChanWriterCloser()
+	targetMsgs := NewChanWriterCloser()
 	re := regexp.MustCompile(`^\[Bitrise.*\].*`)
 
-	sut := loginterceptor.NewPrefixInterceptor(re, interceptWriter, targetWriter, log.NewLogger())
-
-	msg1 := "Log message without prefix\n"
-	msg2 := "[Bitrise Analytics] Log message with prefix\n"
-	msg3 := "[Bitrise Build Cache] Log message with prefix\n"
-	msg4 := "Stuff [Bitrise Build Cache] Log message without prefix\n"
+	sut := loginterceptor.NewPrefixInterceptor(re, &interceptedMsgs, &targetMsgs, log.NewLogger())
 
 	go func() {
-		//nolint:errCheck
-		defer sut.Close()
-
+		defer func() { _ = sut.Close() }()
 		_, _ = sut.Write([]byte(msg1))
 		_, _ = sut.Write([]byte(msg2))
 		_, _ = sut.Write([]byte(msg3))
 		_, _ = sut.Write([]byte(msg4))
 	}()
 
-	intercepted, target, err := readTwo(interceptReader, targetReader)
-	assert.NoError(t, err)
-	assert.Equal(t, msg2+msg3, string(intercepted))
-	assert.Equal(t, msg1+msg2+msg3+msg4, string(target))
+	waitForBoth(sut)
+	_ = interceptedMsgs.Close()
+	_ = targetMsgs.Close()
+
+	assert.Equal(t, msg2+msg3, interceptedMsgs.Messages())
+	assert.Equal(t, msg1+msg2+msg3+msg4, targetMsgs.Messages())
 }
 
 func TestPrefixInterceptorWithPrematureClose(t *testing.T) {
-	interceptReader, interceptWriter := io.Pipe()
-	targetReader, targetWriter := io.Pipe()
+	interceptedMsgs := NewChanWriterCloser()
+	targetMsgs := NewChanWriterCloser()
 	re := regexp.MustCompile(`^\[Bitrise.*\].*`)
 
-	sut := loginterceptor.NewPrefixInterceptor(re, interceptWriter, targetWriter, log.NewLogger())
-
-	msg1 := "Log message without prefix\n"
-	msg2 := "[Bitrise Analytics] Log message with prefix\n"
-	msg3 := "[Bitrise Build Cache] Log message with prefix\n"
-	msg4 := "Last message that won't be sent\n"
+	sut := loginterceptor.NewPrefixInterceptor(re, &interceptedMsgs, &targetMsgs, log.NewLogger())
 
 	go func() {
 		_, _ = sut.Write([]byte(msg1))
@@ -60,71 +56,80 @@ func TestPrefixInterceptorWithPrematureClose(t *testing.T) {
 		_, _ = sut.Write([]byte(msg4))
 	}()
 
-	intercepted, target, err := readTwo(interceptReader, targetReader)
-	assert.NoError(t, err)
-	assert.Equal(t, msg2+msg3, string(intercepted))
-	assert.Equal(t, msg1+msg2+msg3, string(target))
+	waitForBoth(sut)
+	_ = interceptedMsgs.Close()
+	_ = targetMsgs.Close()
+
+	assert.Equal(t, msg2+msg3, interceptedMsgs.Messages())
+	assert.Equal(t, msg1+msg2+msg3, targetMsgs.Messages())
 }
 
 func TestPrefixInterceptorWithBlockedPipe(t *testing.T) {
-	interceptReader, interceptWriter := io.Pipe()
-	targetReader, targetWriter := io.Pipe()
+	_, interceptWriter := io.Pipe()
+	targetMsgs := NewChanWriterCloser()
 	re := regexp.MustCompile(`^\[Bitrise.*\].*`)
 
-	sut := loginterceptor.NewPrefixInterceptor(re, interceptWriter, targetWriter, log.NewLogger())
-
-	msg1 := "Log message without prefix\n"
-	msg2 := "[Bitrise Analytics] Log message with prefix\n"
-	msg3 := "[Bitrise Build Cache] Log message with prefix\n"
-	msg4 := "Stuff [Bitrise Build Cache] Log message without prefix\n"
+	sut := loginterceptor.NewPrefixInterceptor(re, interceptWriter, &targetMsgs, log.NewLogger())
 
 	go func() {
-		//nolint:errCheck
-		defer sut.Close()
-
 		_, _ = sut.Write([]byte(msg1))
 		_, _ = sut.Write([]byte(msg2))
 		_, _ = sut.Write([]byte(msg3))
+		_ = sut.Close()
 		_, _ = sut.Write([]byte(msg4))
 	}()
 
-	target, err := io.ReadAll(targetReader)
-	assert.NoError(t, err)
-	assert.Equal(t, msg1+msg2+msg3+msg4, string(target))
+	<-sut.TargetDelivered
+	_ = targetMsgs.Close()
 
-	// Reading from interceptReader should block until targetWriter is read
-	intercepted, err := io.ReadAll(interceptReader)
-	assert.NoError(t, err)
-	assert.Equal(t, msg2+msg3, string(intercepted))
+	assert.Equal(t, msg1+msg2+msg3, targetMsgs.Messages())
 }
 
-func readTwo(r1, r2 io.Reader) (out1, out2 []byte, err error) {
-	var (
-		wg     sync.WaitGroup
-		e1, e2 error
-	)
+// --------------------------------
+// Helpers
+// --------------------------------
+func waitForBoth(sut *loginterceptor.PrefixInterceptor) {
+	var wg sync.WaitGroup
 	wg.Add(2)
 
-	var b1, b2 bytes.Buffer
-
-	go func() {
-		defer wg.Done()
-		_, e1 = io.Copy(&b1, r1)
-	}()
+	go func(wg *sync.WaitGroup) {
+		<-sut.TargetDelivered
+		wg.Done()
+	}(&wg)
 
-	go func() {
-		defer wg.Done()
-		_, e2 = io.Copy(&b2, r2)
-	}()
+	go func(wg *sync.WaitGroup) {
+		<-sut.InterceptedDelivered
+		wg.Done()
+	}(&wg)
 
 	wg.Wait()
+}
+
+type ChanWriterCloser struct {
+	channel chan string
+}
 
-	// prefer to return the first non-nil error
-	if e1 != nil {
-		return b1.Bytes(), b2.Bytes(), e1
+func NewChanWriterCloser() ChanWriterCloser {
+	return ChanWriterCloser{
+		channel: make(chan string, 1000),
 	}
-	if e2 != nil {
-		return b1.Bytes(), b2.Bytes(), e2
+}
+
+func (ch *ChanWriterCloser) Write(p []byte) (int, error) {
+	ch.channel <- string(p)
+	return len(p), nil
+}
+
+// Close stops the interceptor and closes the pipe.
+func (ch *ChanWriterCloser) Close() error {
+	close(ch.channel)
+	return nil
+}
+
+func (ch *ChanWriterCloser) Messages() string {
+	var result string
+	for msg := range ch.channel {
+		result += msg
 	}
-	return b1.Bytes(), b2.Bytes(), nil
+	return result
 }