fix(writers): StreamingBatchWriter hangs with non-append mode (#1131)

Fixes the hang
[issue](cloudquery/cloudquery#12793) with
non-append mode when a dest is using `UnimplementedDeleteStale`.

- This also makes it so that it closes things when Write is done:
Otherwise our testing framework confuses things (by invoking Write
multiple times for multiple tests in the same plugin, I think).
We're noticing this now because of
cloudquery/cloudquery#12765, as other plugins
using this writer (file based plugins) don't support migration
operations.

Also makes it so that you don't need to call Close() on plugin deinit
(but it could still be a good idea if Write was interrupted)

---------

Co-authored-by: Kemal Hadimli <[email protected]>

Author: disq

writers/streamingbatchwriter/mocktimer_test.go

@@ -1,17 +1,25 @@
 package streamingbatchwriter
 
 import (
+	"sync"
 	"time"
 
 	"github.com/cloudquery/plugin-sdk/v4/writers"
 )
 
 type mockTicker struct {
-	expire chan time.Time
+	expire  chan time.Time
+	stopped sync.Once
 }
 
 func (t *mockTicker) Stop() {
-	close(t.expire)
+	t.stopped.Do(func() {
+		close(t.expire)
+	})
+}
+
+func (t *mockTicker) Tick() {
+	t.expire <- time.Now()
 }
 
 func (*mockTicker) Reset(time.Duration) {}
@@ -20,12 +28,12 @@ func (t *mockTicker) Chan() <-chan time.Time {
 	return t.expire
 }
 
-func newMockTicker() (writers.TickerFunc, chan<- time.Time) {
+func newMockTicker() (writers.TickerFunc, func()) {
 	expire := make(chan time.Time)
 	t := &mockTicker{
 		expire: expire,
 	}
 	return func(time.Duration) writers.Ticker {
 		return t
-	}, expire
+	}, t.Tick
 }

writers/streamingbatchwriter/streamingbatchwriter.go

@@ -156,13 +156,17 @@ func (w *StreamingBatchWriter) Close(context.Context) error {
 	}
 	w.workersWaitGroup.Wait()
 
-	w.insertWorkers = nil
+	w.insertWorkers = make(map[string]*streamingWorkerManager[*message.WriteInsert])
+	w.migrateWorker = nil
+	w.deleteWorker = nil
+	w.lastMsgType = writers.MsgTypeUnset
 
 	return nil
 }
 
 func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessage) error {
 	errCh := make(chan error)
+	defer close(errCh)
 
 	go func() {
 		for err := range errCh {
@@ -172,7 +176,7 @@ func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.Wr
 
 	for msg := range msgs {
 		msgType := writers.MsgID(msg)
-		if w.lastMsgType != msgType {
+		if w.lastMsgType != writers.MsgTypeUnset && w.lastMsgType != msgType {
 			if err := w.Flush(ctx); err != nil {
 				return err
 			}
@@ -183,12 +187,7 @@ func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.Wr
 		}
 	}
 
-	if err := w.Flush(ctx); err != nil {
-		return err
-	}
-
-	close(errCh)
-	return nil
+	return w.Close(ctx)
 }
 
 func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- error, msg message.WriteMessage) error {

writers/streamingbatchwriter/streamingbatchwriter_test.go

@@ -227,7 +227,7 @@ func TestStreamingBatchTimeout(t *testing.T) {
 	ch := make(chan message.WriteMessage)
 
 	testClient := newClient()
-	tickerFn, expire := newMockTicker()
+	tickerFn, tickFn := newMockTicker()
 
 	wr, err := New(testClient, withTickerFn(tickerFn))
 	if err != nil {
@@ -258,7 +258,7 @@ func TestStreamingBatchTimeout(t *testing.T) {
 	}
 
 	// flush
-	close(expire)
+	tickFn()
 	waitForLength(t, testClient.MessageLen, messageTypeInsert, 1)
 
 	close(ch)
@@ -332,7 +332,7 @@ func TestStreamingBatchUpserts(t *testing.T) {
 	ch := make(chan message.WriteMessage)
 
 	testClient := newClient()
-	tickerFn, expire := newMockTicker()
+	tickerFn, tickFn := 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(expire)
+	tickFn()
 	waitForLength(t, testClient.MessageLen, messageTypeInsert, 2)
 
 	close(ch)

writers/streamingbatchwriter/unimplemented.go

@@ -18,9 +18,12 @@ func (IgnoreMigrateTable) MigrateTable(_ context.Context, ch <-chan *message.Wri
 	return nil
 }
 
-// UnimplementedDeleteStale is a dummy handler to error on DeleteStale messages
+// UnimplementedDeleteStale is a dummy handler to consume and error on DeleteStale messages
 type UnimplementedDeleteStale struct{}
 
-func (UnimplementedDeleteStale) DeleteStale(_ context.Context, _ <-chan *message.WriteDeleteStale) error {
+func (UnimplementedDeleteStale) DeleteStale(_ context.Context, ch <-chan *message.WriteDeleteStale) error {
+	// nolint:revive
+	for range ch {
+	}
 	return fmt.Errorf("DeleteStale: %w", plugin.ErrNotImplemented)
 }