shutdown logic fix

Author: disq

writers/streamingbatchwriter/streamingbatchwriter.go

@@ -180,14 +180,13 @@ func (w *StreamingBatchWriter) Close(context.Context) error {
 }
 
 func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessage) error {
-	errCh := make(chan error)
-	defer w.Close(ctx)
-
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
+	errCh := make(chan error)
 	go func() {
 		defer close(errCh)
+		defer w.Close(ctx)
 		for msg := range msgs {
 			msgType := writers.MsgID(msg)
 			if w.lastMsgType != writers.MsgTypeUnset && w.lastMsgType != msgType {
@@ -370,6 +369,7 @@ type streamingWorkerManager[T message.WriteMessage] struct {
 	batchTimeout time.Duration
 	tickerFn     writers.TickerFunc
 	failed       *atomic.Bool
+	workerWg     sync.WaitGroup
 
 	inputCh chan T
 	mu      sync.Mutex // protects inputCh
@@ -404,9 +404,11 @@ func (s *streamingWorkerManager[T]) send(ctx context.Context, data T) {
 	}
 
 	s.inputCh = make(chan T)
+	s.workerWg.Add(1)
 
 	// start consuming our new channel
 	go func(ch chan T) {
+		defer s.workerWg.Done()
 		defer func() {
 			if msg := recover(); msg != nil {
 				switch v := msg.(type) {
@@ -453,6 +455,7 @@ func (s *streamingWorkerManager[T]) send(ctx context.Context, data T) {
 
 func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup) {
 	defer wg.Done()
+	defer s.workerWg.Wait()
 	defer s.closeFlush()
 
 	ticker := s.tickerFn(s.batchTimeout)

writers/streamingbatchwriter/streamingbatchwriter_test.go

@@ -30,9 +30,10 @@ type testStreamingBatchClient struct {
 	committed map[messageType]int
 	open      map[messageType][]string
 
-	writeErr      error
-	writeErrAfter int64
-	writeCounter  map[string]int64 // table name to write counter
+	writeErr       error
+	writeErrAfter  int64
+	writeCounter   map[string]int64 // table name to write counter
+	writeCommitErr error
 }
 
 func newClient() *testStreamingBatchClient {
@@ -84,6 +85,11 @@ func (c *testStreamingBatchClient) WriteTable(ctx context.Context, msgs <-chan *
 			return c.writeErr // leave msgs open
 		}
 	}
+
+	if c.writeCommitErr != nil {
+		return c.writeCommitErr
+	}
+
 	return c.handleTypeCommit(ctx, messageTypeInsert, key)
 }
 
@@ -528,6 +534,41 @@ func TestErrorCleanUpSecondMessage(t *testing.T) {
 	waitForLength(t, testClient.MessageLen, messageTypeInsert, 1)  // batch size 1
 }
 
+func TestErrorCleanUpAfterClose(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	ch := make(chan message.WriteMessage)
+
+	testClient := newClient()
+	testClient.writeCommitErr = errors.New("test error")
+
+	wr, err := New(testClient, WithBatchTimeout(0), WithBatchSizeRows(100))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	errCh := make(chan error)
+	go func() {
+		errCh <- wr.Write(ctx, ch)
+	}()
+
+	table := schema.Table{Name: "table1", Columns: []schema.Column{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}}
+	record := getRecord(table.ToArrowSchema(), 1)
+
+	for i := 0; i < 10; i++ {
+		ch <- &message.WriteInsert{
+			Record: record,
+		}
+	}
+
+	waitForLength(t, testClient.InflightLen, messageTypeInsert, 10)
+	close(ch)
+
+	requireErrorCount(t, 1, errCh)
+
+	waitForLength(t, testClient.MessageLen, messageTypeInsert, 0) // batch size 1
+}
+
 func waitForLength(t *testing.T, checkLen func(messageType) int, msgType messageType, want int) {
 	t.Helper()
 	lastValue := -1