improve job logger

Author: ChristopherHX

protocol/logger/job_logger.go

@@ -6,9 +6,11 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"net/url"
 	"regexp"
+	"runtime/debug"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -193,7 +195,7 @@ func (logger *WebsocketLiveloggerWithFallback) SendLog(wrapper *protocol.Timelin
 	if currentLogger == nil {
 		currentLogger = logger.initialize()
 		if currentLogger == nil {
-			return errors.New("SendLog failure")
+			return errors.New("initialize failure")
 		}
 	}
 	err := currentLogger.SendLog(wrapper)
@@ -237,18 +239,47 @@ func (logger *WebsocketLiveloggerWithFallback) SendLog(wrapper *protocol.Timelin
 
 type internalBufferedLiveLoggerData struct {
 	logchan     chan *protocol.TimelineRecordFeedLinesWrapper
+	logdrain    chan struct{}
 	logfinished chan struct{}
 }
 
+// IsZero returns true if the struct is closed
+// * internalBufferedLiveLoggerData is replaced by a zero struct instead of nil to not auto restart
+func (i *internalBufferedLiveLoggerData) IsZero() bool {
+	return i == nil || i.logchan == nil || i.logdrain == nil || i.logfinished == nil
+}
+
 type BufferedLiveLogger struct {
 	LiveLogger
 	data atomic.Pointer[internalBufferedLiveLoggerData]
 }
 
-func (logger *BufferedLiveLogger) sendLogs(logchan chan *protocol.TimelineRecordFeedLinesWrapper, logfinished chan struct{}) {
+func (logger *BufferedLiveLogger) sendLogs(logchan chan *protocol.TimelineRecordFeedLinesWrapper, logdrain chan struct{}, logfinished chan struct{}) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("panic recovered: %v\n%s", r, debug.Stack())
+		}
+	}()
 	defer close(logfinished)
+	var shouldDrain bool
 	for {
-		lines, ok := <-logchan
+		var lines *protocol.TimelineRecordFeedLinesWrapper
+		var ok bool
+		if !shouldDrain {
+			select {
+			case lines, ok = <-logchan:
+			case <-logdrain:
+				shouldDrain = true
+			}
+		}
+		if shouldDrain {
+			// Now try read
+			select {
+			case lines, ok = <-logchan:
+			default:
+				ok = false
+			}
+		}
 		if !ok {
 			return
 		}
@@ -291,26 +322,36 @@ func (logger *BufferedLiveLogger) sendLogs(logchan chan *protocol.TimelineRecord
 }
 
 func (logger *BufferedLiveLogger) Close() error {
-	if data := logger.data.Swap(nil); data != nil {
-		close(data.logchan)
+	if data := logger.data.Swap(&internalBufferedLiveLoggerData{}); !data.IsZero() {
+		// Keep logchan open to avoid races and just let it be freed
+		close(data.logdrain)
 		<-data.logfinished
 	}
 	return nil
 }
 
 func (logger *BufferedLiveLogger) SendLog(wrapper *protocol.TimelineRecordFeedLinesWrapper) error {
 	if data := logger.data.Load(); data != nil {
-		data.logchan <- wrapper
+		if data.IsZero() {
+			return errors.New("buffered live logger is closed")
+		}
+		select {
+		case <-data.logdrain:
+			return errors.New("buffered live logger closing")
+		default:
+			data.logchan <- wrapper
+		}
 	} else {
 		logchan := make(chan *protocol.TimelineRecordFeedLinesWrapper, websocketPingSize)
 		logfinished := make(chan struct{})
 		ndata := internalBufferedLiveLoggerData{
 			logchan:     logchan,
+			logdrain:    make(chan struct{}),
 			logfinished: logfinished,
 		}
 		if logger.data.CompareAndSwap(data, &ndata) {
 			ndata.logchan <- wrapper
-			go logger.sendLogs(logchan, logfinished)
+			go logger.sendLogs(logchan, ndata.logdrain, logfinished)
 		} else {
 			close(ndata.logchan)
 			close(ndata.logfinished)

protocol/logger/job_logger_test.go

@@ -1,9 +1,12 @@
 package logger
 
 import (
+	"fmt"
 	"testing"
+	"time"
 
 	"github.com/ChristopherHX/github-act-runner/protocol"
+	"github.com/stretchr/testify/require"
 )
 
 const (
@@ -35,3 +38,85 @@ func TestJobLogger(t *testing.T) {
 		t.FailNow()
 	}
 }
+
+type TestLiveLogger struct {
+	*testing.T
+}
+
+// Close implements [LiveLogger].
+func (t *TestLiveLogger) Close() error {
+	return nil
+}
+
+// SendLog implements [LiveLogger].
+func (t *TestLiveLogger) SendLog(lines *protocol.TimelineRecordFeedLinesWrapper) error {
+	return nil
+}
+
+func getTestLiveLogger(t *testing.T) LiveLogger {
+	return &TestLiveLogger{
+		T: t,
+	}
+}
+
+func TestBufferedLiveLoggerDrain(t *testing.T) {
+	bufferedLogger := &BufferedLiveLogger{
+		LiveLogger: getTestLiveLogger(t),
+	}
+
+	var logchan chan *protocol.TimelineRecordFeedLinesWrapper = make(chan *protocol.TimelineRecordFeedLinesWrapper)
+	var logdrain chan struct{} = make(chan struct{})
+	var logfinished chan struct{} = make(chan struct{})
+
+	t.Run("forwardLogs", func(t *testing.T) {
+		t.Parallel()
+		bufferedLogger.sendLogs(logchan, logdrain, logfinished)
+	})
+
+	t.Run("sendLogs", func(t *testing.T) {
+		t.Parallel()
+		logchan <- &protocol.TimelineRecordFeedLinesWrapper{
+			Count: 1,
+			Value: []string{"line1"},
+		}
+		logchan <- &protocol.TimelineRecordFeedLinesWrapper{
+			Count: 1,
+			Value: []string{"line2"},
+		}
+		close(logdrain)
+		<-logfinished
+	})
+}
+
+func TestBufferedLiveLogger(t *testing.T) {
+	bufferedLogger := &BufferedLiveLogger{
+		LiveLogger: getTestLiveLogger(t),
+	}
+
+	require.NoError(t, bufferedLogger.SendLog(&protocol.TimelineRecordFeedLinesWrapper{
+		Count: 1,
+		Value: []string{"line1"},
+	}))
+
+	t.Run("close", func(t *testing.T) {
+		t.Parallel()
+		time.Sleep(1122 * time.Microsecond)
+		require.NoError(t, bufferedLogger.Close())
+	})
+
+	t.Run("sendLogs", func(t *testing.T) {
+		t.Parallel()
+		var err error
+		for i := range 100 {
+			time.Sleep(1 * time.Millisecond)
+			err = bufferedLogger.SendLog(&protocol.TimelineRecordFeedLinesWrapper{
+				Count: 1,
+				Value: []string{fmt.Sprintf("line %v", (i + 1))},
+			})
+			if err != nil {
+				break
+			}
+		}
+		require.Error(t, err)
+	})
+}