From 98b738efe6b6be776f83fe8f73cb3265a943d53f Mon Sep 17 00:00:00 2001
From: candiduslynx <candiduslynx@gmail.com>
Date: Tue, 28 Nov 2023 11:45:53 +0200
Subject: [PATCH] chore: Cleanup batch writers

---
 writers/batchwriter/batchwriter.go            |  2 --
 .../streamingbatchwriter.go                   | 20 +++++++++++--------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/writers/batchwriter/batchwriter.go b/writers/batchwriter/batchwriter.go
index 652521b0c4..a996552a34 100644
--- a/writers/batchwriter/batchwriter.go
+++ b/writers/batchwriter/batchwriter.go
@@ -69,7 +69,6 @@ func WithBatchSizeBytes(size int) Option {
 }
 
 type worker struct {
-	count int
 	ch    chan *message.WriteInsert
 	flush chan chan bool
 }
@@ -322,7 +321,6 @@ func (w *BatchWriter) startWorker(_ context.Context, msg *message.WriteInsert) e
 	ch := make(chan *message.WriteInsert)
 	flush := make(chan chan bool)
 	wr = &worker{
-		count: 1,
 		ch:    ch,
 		flush: flush,
 	}
diff --git a/writers/streamingbatchwriter/streamingbatchwriter.go b/writers/streamingbatchwriter/streamingbatchwriter.go
index cb454c6b2e..d7196e6077 100644
--- a/writers/streamingbatchwriter/streamingbatchwriter.go
+++ b/writers/streamingbatchwriter/streamingbatchwriter.go
@@ -226,12 +226,11 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			return nil
 		}
 		ch := make(chan *message.WriteMigrateTable)
-		flush := make(chan chan bool)
 		w.migrateWorker = &streamingWorkerManager[*message.WriteMigrateTable]{
 			ch:        ch,
 			writeFunc: w.client.MigrateTable,
 
-			flush: flush,
+			flush: make(chan chan bool),
 			errCh: errCh,
 
 			batchSizeRows: w.batchSizeRows,
@@ -251,12 +250,11 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			return nil
 		}
 		ch := make(chan *message.WriteDeleteStale)
-		flush := make(chan chan bool)
 		w.deleteStaleWorker = &streamingWorkerManager[*message.WriteDeleteStale]{
 			ch:        ch,
 			writeFunc: w.client.DeleteStale,
 
-			flush: flush,
+			flush: make(chan chan bool),
 			errCh: errCh,
 
 			batchSizeRows: w.batchSizeRows,
@@ -278,12 +276,11 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 		}
 
 		ch := make(chan *message.WriteInsert)
-		flush := make(chan chan bool)
 		wr = &streamingWorkerManager[*message.WriteInsert]{
 			ch:        ch,
 			writeFunc: w.client.WriteTable,
 
-			flush: flush,
+			flush: make(chan chan bool),
 			errCh: errCh,
 
 			batchSizeRows:  w.batchSizeRows,
@@ -292,6 +289,14 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			tickerFn:       w.tickerFn,
 		}
 		w.workersLock.Lock()
+		wrOld, ok := w.insertWorkers[tableName]
+		if ok {
+			w.workersLock.Unlock()
+			// some other goroutine could have already added the worker
+			// just send the message to it & discard our allocated worker
+			wrOld.ch <- m
+			return nil
+		}
 		w.insertWorkers[tableName] = wr
 		w.workersLock.Unlock()
 
@@ -307,13 +312,12 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			return nil
 		}
 		ch := make(chan *message.WriteDeleteRecord)
-		flush := make(chan chan bool)
 		// TODO: flush all workers for nested tables as well (See https://github.com/cloudquery/plugin-sdk/issues/1296)
 		w.deleteRecordWorker = &streamingWorkerManager[*message.WriteDeleteRecord]{
 			ch:        ch,
 			writeFunc: w.client.DeleteRecords,
 
-			flush: flush,
+			flush: make(chan chan bool),
 			errCh: errCh,
 
 			batchSizeRows: w.batchSizeRows,