fix(writers): Allow zero timeout, remove unused timeout options from mixedbatchwriter (#1020)

Author: disq

writers/batchwriter/batchwriter.go

@@ -133,7 +133,7 @@ func (w *BatchWriter) worker(ctx context.Context, tableName string, ch <-chan *m
 
 			resources = append(resources, r)
 			sizeBytes += util.TotalRecordSize(r.Record)
-		case <-time.After(w.batchTimeout):
+		case <-timer(w.batchTimeout):
 			if len(resources) > 0 {
 				w.flushTable(ctx, tableName, resources)
 				resources, sizeBytes = resources[:0], 0
@@ -316,3 +316,10 @@ func (w *BatchWriter) startWorker(ctx context.Context, msg *message.WriteInsert)
 	ch <- msg
 	return nil
 }
+
+func timer(timeout time.Duration) <-chan time.Time {
+	if timeout == 0 {
+		return nil
+	}
+	return time.After(timeout)
+}

writers/mixedbatchwriter/mixedbatchwriter.go

@@ -2,7 +2,6 @@ package mixedbatchwriter
 
 import (
 	"context"
-	"time"
 
 	"github.com/apache/arrow/go/v13/arrow/util"
 	"github.com/cloudquery/plugin-sdk/v4/message"
@@ -20,7 +19,6 @@ type Client interface {
 type MixedBatchWriter struct {
 	client         Client
 	logger         zerolog.Logger
-	batchTimeout   time.Duration
 	batchSize      int
 	batchSizeBytes int
 }
@@ -36,12 +34,6 @@ func WithLogger(logger zerolog.Logger) Option {
 	}
 }
 
-func WithBatchTimeout(timeout time.Duration) Option {
-	return func(p *MixedBatchWriter) {
-		p.batchTimeout = timeout
-	}
-}
-
 func WithBatchSize(size int) Option {
 	return func(p *MixedBatchWriter) {
 		p.batchSize = size
@@ -58,7 +50,6 @@ func New(client Client, opts ...Option) (*MixedBatchWriter, error) {
 	c := &MixedBatchWriter{
 		client:         client,
 		logger:         zerolog.Nop(),
-		batchTimeout:   writers.DefaultBatchTimeoutSeconds * time.Second,
 		batchSize:      writers.DefaultBatchSize,
 		batchSizeBytes: writers.DefaultBatchSizeBytes,
 	}

writers/streamingbatchwriter/mocktimer_test.go

@@ -0,0 +1,19 @@
+package streamingbatchwriter
+
+import "time"
+
+type mockTimer struct {
+	expire chan time.Time
+}
+
+func (t *mockTimer) timer(time.Duration) <-chan time.Time {
+	return t.expire
+}
+
+func newMockTimer() (timerFn, chan time.Time) {
+	expire := make(chan time.Time)
+	t := &mockTimer{
+		expire: expire,
+	}
+	return t.timer, expire
+}

writers/streamingbatchwriter/streamingbatchwriter.go

@@ -60,8 +60,12 @@ type StreamingBatchWriter struct {
 	batchTimeout   time.Duration
 	batchSizeRows  int64
 	batchSizeBytes int64
+
+	timerFn timerFn
 }
 
+type timerFn func(timeout time.Duration) <-chan time.Time
+
 // Assert at compile-time that StreamingBatchWriter implements the Writer interface
 var _ writers.Writer = (*StreamingBatchWriter)(nil)
 
@@ -91,6 +95,12 @@ func WithBatchSizeBytes(size int64) Option {
 	}
 }
 
+func withTimerFn(timer timerFn) Option {
+	return func(p *StreamingBatchWriter) {
+		p.timerFn = timer
+	}
+}
+
 func New(client Client, opts ...Option) (*StreamingBatchWriter, error) {
 	c := &StreamingBatchWriter{
 		client:         client,
@@ -99,6 +109,7 @@ func New(client Client, opts ...Option) (*StreamingBatchWriter, error) {
 		batchTimeout:   writers.DefaultBatchTimeoutSeconds * time.Second,
 		batchSizeRows:  writers.DefaultBatchSize,
 		batchSizeBytes: writers.DefaultBatchSizeBytes,
+		timerFn:        timer,
 	}
 	for _, opt := range opts {
 		opt(c)
@@ -214,6 +225,7 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 
 			batchSizeRows: w.batchSizeRows,
 			batchTimeout:  w.batchTimeout,
+			timerFn:       w.timerFn,
 		}
 
 		w.workersWaitGroup.Add(1)
@@ -265,6 +277,7 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			batchSizeRows:  w.batchSizeRows,
 			batchSizeBytes: w.batchSizeBytes,
 			batchTimeout:   w.batchTimeout,
+			timerFn:        w.timerFn,
 		}
 		w.workersLock.Lock()
 		w.insertWorkers[tableName] = wr
@@ -290,6 +303,7 @@ type streamingWorkerManager[T message.WriteMessage] struct {
 	batchSizeRows  int64
 	batchSizeBytes int64
 	batchTimeout   time.Duration
+	timerFn        timerFn
 }
 
 func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup, tableName string) {
@@ -351,7 +365,7 @@ func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup,
 			clientCh <- r
 			sizeRows++
 			sizeBytes += recSize
-		case <-time.After(s.batchTimeout):
+		case <-s.timerFn(s.batchTimeout):
 			if sizeRows > 0 {
 				closeFlush()
 			}
@@ -370,3 +384,10 @@ func DummyHandler[T message.WriteMessage](ch <-chan T) {
 	for range ch {
 	}
 }
+
+func timer(timeout time.Duration) <-chan time.Time {
+	if timeout == 0 {
+		return nil
+	}
+	return time.After(timeout)
+}