Allow flush to be used multiple times (#1051)

Author: giortzisg

batch_logger.go

@@ -12,17 +12,20 @@ const (
 )
 
 type BatchLogger struct {
-	client    *Client
-	logCh     chan Log
-	cancel    context.CancelFunc
-	wg        sync.WaitGroup
-	startOnce sync.Once
+	client       *Client
+	logCh        chan Log
+	flushCh      chan chan struct{}
+	cancel       context.CancelFunc
+	wg           sync.WaitGroup
+	startOnce    sync.Once
+	shutdownOnce sync.Once
 }
 
 func NewBatchLogger(client *Client) *BatchLogger {
 	return &BatchLogger{
-		client: client,
-		logCh:  make(chan Log, batchSize),
+		client:  client,
+		logCh:   make(chan Log, batchSize),
+		flushCh: make(chan chan struct{}),
 	}
 }
 
@@ -35,17 +38,32 @@ func (l *BatchLogger) Start() {
 	})
 }
 
-func (l *BatchLogger) Flush() {
-	if l.cancel != nil {
-		l.cancel()
-		l.wg.Wait()
+func (l *BatchLogger) Flush(timeout <-chan struct{}) {
+	done := make(chan struct{})
+	select {
+	case l.flushCh <- done:
+		select {
+		case <-done:
+		case <-timeout:
+		}
+	case <-timeout:
 	}
 }
 
+func (l *BatchLogger) Shutdown() {
+	l.shutdownOnce.Do(func() {
+		if l.cancel != nil {
+			l.cancel()
+			l.wg.Wait()
+		}
+	})
+}
+
 func (l *BatchLogger) run(ctx context.Context) {
 	defer l.wg.Done()
 	var logs []Log
 	timer := time.NewTimer(batchTimeout)
+	defer timer.Stop()
 
 	for {
 		select {
@@ -65,8 +83,27 @@ func (l *BatchLogger) run(ctx context.Context) {
 				logs = nil
 			}
 			timer.Reset(batchTimeout)
+		case done := <-l.flushCh:
+		flushDrain:
+			for {
+				select {
+				case log := <-l.logCh:
+					logs = append(logs, log)
+				default:
+					break flushDrain
+				}
+			}
+
+			if len(logs) > 0 {
+				l.processEvent(logs)
+				logs = nil
+			}
+			if !timer.Stop() {
+				<-timer.C
+			}
+			timer.Reset(batchTimeout)
+			close(done)
 		case <-ctx.Done():
-			// Drain remaining logs from channel
 		drain:
 			for {
 				select {

client.go

@@ -511,7 +511,17 @@ func (client *Client) RecoverWithContext(
 // call to Init.
 func (client *Client) Flush(timeout time.Duration) bool {
 	if client.batchLogger != nil {
-		client.batchLogger.Flush()
+		start := time.Now()
+		timeoutCh := make(chan struct{})
+		time.AfterFunc(timeout, func() {
+			close(timeoutCh)
+		})
+		client.batchLogger.Flush(timeoutCh)
+		// update the timeout with the time passed
+		timeout -= time.Since(start)
+		if timeout <= 0 {
+			return false
+		}
 	}
 	return client.Transport.Flush(timeout)
 }
@@ -530,7 +540,7 @@ func (client *Client) Flush(timeout time.Duration) bool {
 
 func (client *Client) FlushWithContext(ctx context.Context) bool {
 	if client.batchLogger != nil {
-		client.batchLogger.Flush()
+		client.batchLogger.Flush(ctx.Done())
 	}
 	return client.Transport.FlushWithContext(ctx)
 }

log_test.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/getsentry/sentry-go/attribute"
+	"github.com/getsentry/sentry-go/internal/testutils"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 )
@@ -177,7 +178,7 @@ func Test_sentryLogger_MethodsWithFormat(t *testing.T) {
 			// invalid attribute should be dropped
 			l.SetAttributes(attribute.Builder{Key: "key.invalid", Value: attribute.Value{}})
 			tt.logFunc(ctx, l)
-			Flush(20 * time.Millisecond)
+			Flush(testutils.FlushTimeout())
 
 			opts := cmp.Options{
 				cmpopts.IgnoreFields(Log{}, "Timestamp"),
@@ -321,7 +322,7 @@ func Test_sentryLogger_MethodsWithoutFormat(t *testing.T) {
 			ctx, mockTransport := setupMockTransport()
 			l := NewLogger(ctx)
 			tt.logFunc(ctx, l, tt.args)
-			Flush(20 * time.Millisecond)
+			Flush(testutils.FlushTimeout())
 
 			opts := cmp.Options{
 				cmpopts.IgnoreFields(Log{}, "Timestamp"),
@@ -400,7 +401,7 @@ func Test_sentryLogger_Write(t *testing.T) {
 	if n != len(msg) {
 		t.Errorf("Write returned wrong byte count: got %d, want %d", n, len(msg))
 	}
-	Flush(20 * time.Millisecond)
+	Flush(testutils.FlushTimeout())
 
 	gotEvents := mockTransport.Events()
 	if len(gotEvents) != 1 {
@@ -426,7 +427,7 @@ func Test_sentryLogger_FlushAttributesAfterSend(t *testing.T) {
 
 	l.SetAttributes(attribute.String("string", "some str"))
 	l.Warn(ctx, msg)
-	Flush(20 * time.Millisecond)
+	Flush(testutils.FlushTimeout())
 
 	gotEvents := mockTransport.Events()
 	if len(gotEvents) != 1 {
@@ -444,7 +445,7 @@ func Test_batchLogger_Flush(t *testing.T) {
 	ctx, mockTransport := setupMockTransport()
 	l := NewLogger(context.Background())
 	l.Info(ctx, "context done log")
-	Flush(20 * time.Millisecond)
+	Flush(testutils.FlushTimeout())
 
 	events := mockTransport.Events()
 	if len(events) != 1 {
@@ -457,7 +458,7 @@ func Test_batchLogger_FlushWithContext(t *testing.T) {
 	l := NewLogger(context.Background())
 	l.Info(ctx, "context done log")
 
-	cancelCtx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	cancelCtx, cancel := context.WithTimeout(context.Background(), testutils.FlushTimeout())
 	FlushWithContext(cancelCtx)
 	defer cancel()
 
@@ -467,6 +468,88 @@ func Test_batchLogger_FlushWithContext(t *testing.T) {
 	}
 }
 
+func Test_batchLogger_FlushMultipleTimes(t *testing.T) {
+	ctx, mockTransport := setupMockTransport()
+	l := NewLogger(ctx)
+
+	for i := 0; i < 5; i++ {
+		l.Info(ctx, "test")
+	}
+
+	Flush(testutils.FlushTimeout())
+
+	events := mockTransport.Events()
+	if len(events) != 1 {
+		t.Logf("Got %d events instead of 1", len(events))
+		for i, event := range events {
+			t.Logf("Event %d: %d logs", i, len(event.Logs))
+		}
+		t.Fatalf("expected 1 event after first flush, got %d", len(events))
+	}
+	if len(events[0].Logs) != 5 {
+		t.Fatalf("expected 5 logs in first batch, got %d", len(events[0].Logs))
+	}
+
+	mockTransport.events = nil
+
+	for i := 0; i < 3; i++ {
+		l.Info(ctx, "test")
+	}
+
+	Flush(testutils.FlushTimeout())
+	events = mockTransport.Events()
+	if len(events) != 1 {
+		t.Fatalf("expected 1 event after second flush, got %d", len(events))
+	}
+	if len(events[0].Logs) != 3 {
+		t.Fatalf("expected 3 logs in second batch, got %d", len(events[0].Logs))
+	}
+
+	mockTransport.events = nil
+
+	Flush(testutils.FlushTimeout())
+	events = mockTransport.Events()
+	if len(events) != 0 {
+		t.Fatalf("expected 0 events after third flush with no logs, got %d", len(events))
+	}
+}
+
+func Test_batchLogger_Shutdown(t *testing.T) {
+	ctx, mockTransport := setupMockTransport()
+	l := NewLogger(ctx)
+	for i := 0; i < 3; i++ {
+		l.Info(ctx, "test")
+	}
+
+	hub := GetHubFromContext(ctx)
+	hub.Client().batchLogger.Shutdown()
+
+	events := mockTransport.Events()
+	if len(events) != 1 {
+		t.Fatalf("expected 1 event after shutdown, got %d", len(events))
+	}
+	if len(events[0].Logs) != 3 {
+		t.Fatalf("expected 3 logs in shutdown batch, got %d", len(events[0].Logs))
+	}
+
+	mockTransport.events = nil
+
+	// Test that shutdown can be called multiple times safely
+	hub.Client().batchLogger.Shutdown()
+	hub.Client().batchLogger.Shutdown()
+
+	events = mockTransport.Events()
+	if len(events) != 0 {
+		t.Fatalf("expected 0 events after multiple shutdowns, got %d", len(events))
+	}
+
+	Flush(testutils.FlushTimeout())
+	events = mockTransport.Events()
+	if len(events) != 0 {
+		t.Fatalf("expected 0 events after flush on shutdown logger, got %d", len(events))
+	}
+}
+
 func Test_sentryLogger_BeforeSendLog(t *testing.T) {
 	ctx := context.Background()
 	mockTransport := &MockTransport{}
@@ -492,7 +575,7 @@ func Test_sentryLogger_BeforeSendLog(t *testing.T) {
 
 	l := NewLogger(ctx)
 	l.Info(ctx, "context done log")
-	Flush(20 * time.Millisecond)
+	Flush(testutils.FlushTimeout())
 
 	events := mockTransport.Events()
 	if len(events) != 0 {
@@ -528,7 +611,7 @@ func Test_sentryLogger_TracePropagationWithTransaction(t *testing.T) {
 	logger := NewLogger(txn.Context())
 	logger.Info(txn.Context(), "message with tracing")
 
-	Flush(20 * time.Millisecond)
+	Flush(testutils.FlushTimeout())
 
 	events := mockTransport.Events()
 	if len(events) != 1 {