fix: Reset timers on flush (#1076)

This updates the ticker logic in the batch writers to reset the ticker when a flush happens. This is better, as it still guarantees that a message won't be delayed by more than batch_timeout, but we don't risk flushing a very small batch because we must flush at regular intervals either.

The choice of resetting _after_ the flush is deliberate: it means that the maximum amount of time between flushes is:

```
max_time = flush_time + batch_timeout
```

otherwise we could do:

```
max_time = batch_timeout
```

but if a flush were to then longer than the batch timeout, we can end up in a cycle of flushing again immediately after the previous flush finishes.

Author: hermanschaaf

writers/batchwriter/batchwriter.go

@@ -122,8 +122,8 @@ func (w *BatchWriter) Close(context.Context) error {
 func (w *BatchWriter) worker(ctx context.Context, tableName string, ch <-chan *message.WriteInsert, flush <-chan chan bool) {
 	sizeBytes := int64(0)
 	resources := make([]*message.WriteInsert, 0, w.batchSize)
-	tick, done := writers.NewTicker(w.batchTimeout)
-	defer done()
+	ticker := writers.NewTicker(w.batchTimeout)
+	defer ticker.Stop()
 	for {
 		select {
 		case r, ok := <-ch:
@@ -136,19 +136,22 @@ func (w *BatchWriter) worker(ctx context.Context, tableName string, ch <-chan *m
 
 			if (w.batchSize > 0 && len(resources) >= w.batchSize) || (w.batchSizeBytes > 0 && sizeBytes+util.TotalRecordSize(r.Record) >= int64(w.batchSizeBytes)) {
 				w.flushTable(ctx, tableName, resources)
+				ticker.Reset(w.batchTimeout)
 				resources, sizeBytes = resources[:0], 0
 			}
 
 			resources = append(resources, r)
 			sizeBytes += util.TotalRecordSize(r.Record)
-		case <-tick:
+		case <-ticker.Chan():
 			if len(resources) > 0 {
 				w.flushTable(ctx, tableName, resources)
+				ticker.Reset(w.batchTimeout)
 				resources, sizeBytes = resources[:0], 0
 			}
 		case done := <-flush:
 			if len(resources) > 0 {
 				w.flushTable(ctx, tableName, resources)
+				ticker.Reset(w.batchTimeout)
 				resources, sizeBytes = resources[:0], 0
 			}
 			done <- true
@@ -170,33 +173,6 @@ func (w *BatchWriter) flushTable(ctx context.Context, tableName string, resource
 	}
 }
 
-// func (*BatchWriter) removeDuplicatesByPK(table *schema.Table, resources []*message.Insert) []*message.Insert {
-// 	pkIndices := table.PrimaryKeysIndexes()
-// 	// special case where there's no PK at all
-// 	if len(pkIndices) == 0 {
-// 		return resources
-// 	}
-
-// 	pks := make(map[string]struct{}, len(resources))
-// 	res := make([]*message.Insert, 0, len(resources))
-// 	for _, r := range resources {
-// 		if r.Record.NumRows() > 1 {
-// 			panic(fmt.Sprintf("record with more than 1 row: %d", r.Record.NumRows()))
-// 		}
-// 		key := pk.String(r.Record)
-// 		_, ok := pks[key]
-// 		if !ok {
-// 			pks[key] = struct{}{}
-// 			res = append(res, r)
-// 			continue
-// 		}
-// 		// duplicate, release
-// 		r.Release()
-// 	}
-
-// 	return res
-// }
-
 func (w *BatchWriter) flushMigrateTables(ctx context.Context) error {
 	w.migrateTableLock.Lock()
 	defer w.migrateTableLock.Unlock()

writers/mixedbatchwriter/mixedbatchwriter.go

@@ -114,8 +114,8 @@ func (w *MixedBatchWriter) Write(ctx context.Context, msgChan <-chan message.Wri
 	}
 	prevMsgType := writers.MsgTypeUnset
 	var err error
-	tick, done := w.tickerFn(w.batchTimeout)
-	defer done()
+	ticker := w.tickerFn(w.batchTimeout)
+	defer ticker.Stop()
 loop:
 	for {
 		select {
@@ -128,6 +128,7 @@ loop:
 				if err := flush(prevMsgType); err != nil {
 					return err
 				}
+				ticker.Reset(w.batchTimeout)
 			}
 			prevMsgType = msgType
 			switch v := msg.(type) {
@@ -143,10 +144,11 @@ loop:
 			if err != nil {
 				return err
 			}
-		case <-tick:
+		case <-ticker.Chan():
 			if err := flush(prevMsgType); err != nil {
 				return err
 			}
+			ticker.Reset(w.batchTimeout)
 			prevMsgType = writers.MsgTypeUnset
 		}
 	}

writers/mixedbatchwriter/mixedbatchwriter_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/apache/arrow/go/v13/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/cloudquery/plugin-sdk/v4/writers"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -214,6 +215,38 @@ func TestMixedBatchWriter(t *testing.T) {
 	}
 }
 
+type mockTicker struct {
+	C       chan time.Time
+	trigger chan struct{}
+}
+
+func (m *mockTicker) Chan() <-chan time.Time {
+	return m.C
+}
+
+func (m *mockTicker) Trigger() chan<- struct{} {
+	return m.trigger
+}
+
+func (m *mockTicker) Stop() {
+	close(m.C)
+}
+
+func (*mockTicker) Reset(_ time.Duration) {}
+
+func newMockTicker(trigger chan struct{}) *mockTicker {
+	c := make(chan time.Time)
+	go func() {
+		for range trigger {
+			c <- time.Now()
+		}
+	}()
+	return &mockTicker{
+		C:       c,
+		trigger: trigger,
+	}
+}
+
 func TestMixedBatchWriterTimeout(t *testing.T) {
 	tm := getTestMessages()
 	cases := []struct {
@@ -240,17 +273,12 @@ func TestMixedBatchWriterTimeout(t *testing.T) {
 				receivedBatches: make([][]message.WriteMessage, 0),
 			}
 			triggerTimeout := make(chan struct{})
+			defer close(triggerTimeout)
 			wr, err := New(client,
 				WithBatchSize(1000),
 				WithBatchSizeBytes(1000000),
-				withTickerFn(func(_ time.Duration) (<-chan time.Time, func()) {
-					c := make(chan time.Time)
-					go func() {
-						for range triggerTimeout {
-							c <- time.Now()
-						}
-					}()
-					return c, func() { close(c) }
+				withTickerFn(func(_ time.Duration) writers.Ticker {
+					return newMockTicker(triggerTimeout)
 				}),
 			)
 			if err != nil {

writers/streamingbatchwriter/mocktimer_test.go

@@ -6,22 +6,26 @@ import (
 	"github.com/cloudquery/plugin-sdk/v4/writers"
 )
 
-type mockTimer struct {
+type mockTicker struct {
 	expire chan time.Time
 }
 
-func (t *mockTimer) timer(time.Duration) (<-chan time.Time, func()) {
-	return t.expire, t.close
+func (t *mockTicker) Stop() {
+	close(t.expire)
 }
 
-func (t *mockTimer) close() {
-	close(t.expire)
+func (*mockTicker) Reset(time.Duration) {}
+
+func (t *mockTicker) Chan() <-chan time.Time {
+	return t.expire
 }
 
-func newMockTimer() (writers.TickerFunc, chan time.Time) {
+func newMockTicker() (writers.TickerFunc, chan<- time.Time) {
 	expire := make(chan time.Time)
-	t := &mockTimer{
+	t := &mockTicker{
 		expire: expire,
 	}
-	return t.timer, expire
+	return func(time.Duration) writers.Ticker {
+		return t
+	}, expire
 }

writers/streamingbatchwriter/streamingbatchwriter.go

@@ -343,8 +343,8 @@ func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup,
 	}
 	defer closeFlush()
 
-	tick, done := s.tickerFn(s.batchTimeout)
-	defer done()
+	ticker := s.tickerFn(s.batchTimeout)
+	defer ticker.Stop()
 	for {
 		select {
 		case r, ok := <-s.ch:
@@ -359,19 +359,21 @@ func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup,
 
 			if (s.batchSizeRows > 0 && sizeRows >= s.batchSizeRows) || (s.batchSizeBytes > 0 && sizeBytes+recSize >= s.batchSizeBytes) {
 				closeFlush()
+				ticker.Reset(s.batchTimeout)
 			}
 
 			ensureOpened()
 			clientCh <- r
 			sizeRows++
 			sizeBytes += recSize
-		case <-tick:
+		case <-ticker.Chan():
 			if sizeRows > 0 {
 				closeFlush()
 			}
 		case done := <-s.flush:
 			if sizeRows > 0 {
 				closeFlush()
+				ticker.Reset(s.batchTimeout)
 			}
 			done <- true
 		}

writers/streamingbatchwriter/streamingbatchwriter_test.go

@@ -227,9 +227,9 @@ func TestStreamingBatchTimeout(t *testing.T) {
 	ch := make(chan message.WriteMessage)
 
 	testClient := newClient()
-	timerFn, timerExpire := newMockTimer()
+	tickerFn, expire := newMockTicker()
 
-	wr, err := New(testClient, withTickerFn(timerFn))
+	wr, err := New(testClient, withTickerFn(tickerFn))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -258,7 +258,7 @@ func TestStreamingBatchTimeout(t *testing.T) {
 	}
 
 	// flush
-	close(timerExpire)
+	close(expire)
 	waitForLength(t, testClient.MessageLen, messageTypeInsert, 1)
 
 	close(ch)
@@ -332,8 +332,8 @@ func TestStreamingBatchUpserts(t *testing.T) {
 	ch := make(chan message.WriteMessage)
 
 	testClient := newClient()
-	timerFn, timerExpire := newMockTimer()
-	wr, err := New(testClient, WithBatchSizeRows(2), withTickerFn(timerFn))
+	tickerFn, expire := newMockTicker()
+	wr, err := New(testClient, WithBatchSizeRows(2), withTickerFn(tickerFn))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -363,7 +363,7 @@ func TestStreamingBatchUpserts(t *testing.T) {
 	time.Sleep(50 * time.Millisecond)
 
 	// flush the batch
-	close(timerExpire)
+	close(expire)
 	waitForLength(t, testClient.MessageLen, messageTypeInsert, 2)
 
 	close(ch)
@@ -379,9 +379,10 @@ func TestStreamingBatchUpserts(t *testing.T) {
 func waitForLength(t *testing.T, checkLen func(messageType) int, msgType messageType, want int) {
 	t.Helper()
 	lastValue := -1
+	timeout := time.After(5 * time.Second)
 	for {
 		select {
-		case <-time.After(time.Second):
+		case <-timeout:
 			t.Fatalf("timed out waiting for %v message length %d (last value: %d)", msgType, want, lastValue)
 		default:
 			if lastValue = checkLen(msgType); lastValue == want {

writers/ticker.go

@@ -4,14 +4,35 @@ import (
 	"time"
 )
 
-type TickerFunc func(interval time.Duration) (ch <-chan time.Time, done func())
+type TickerFunc func(time.Duration) Ticker
 
-func NewTicker(interval time.Duration) (<-chan time.Time, func()) {
+type Ticker interface {
+	Stop()
+	Reset(d time.Duration)
+	Chan() <-chan time.Time
+}
+
+func NewTicker(interval time.Duration) Ticker {
 	if interval <= 0 {
-		return nil, nop
+		return nopTicker{}
 	}
-	t := time.NewTicker(interval)
-	return t.C, t.Stop
+	return &ticker{time.NewTicker(interval)}
+}
+
+type ticker struct {
+	*time.Ticker
+}
+
+func (t *ticker) Chan() <-chan time.Time {
+	return t.C
 }
 
-func nop() {}
+type nopTicker struct{}
+
+func (nopTicker) Stop() {}
+
+func (nopTicker) Reset(_ time.Duration) {}
+
+func (nopTicker) Chan() <-chan time.Time {
+	return nil
+}