feat: Allow sync to be cancelled when in progress (#1334)

In some cases, when the resource channel is processing a large collection of objects, the sync process can continue even if the context cancel has been triggered. This is problematic for billing and usage since a sync can significantly overshoot the remaining rows quota. To allow an in-progress sync to abort quicker, we need to modify a number of channel patterns from using `range` to using `select` - which will allow the `context.Done` channel to be monitored.

This PR also adds the `OnSyncFinisher` hook which is called when the sync has finished adding messages. This allows the client code to call the batch updater close method without missing any updates.

fixes: cloudquery/cloudquery-issues#750

Author: mnorbury

internal/servers/plugin/v3/plugin.go

@@ -210,6 +210,10 @@ func (s *Server) Sync(req *pb.Sync_Request, stream pb.Plugin_SyncServer) error {
 		}
 	}
 
+	if err := s.Plugin.OnSyncFinish(ctx); err != nil {
+		return status.Errorf(codes.Internal, "failed to finish sync: %v", err)
+	}
+
 	return syncErr
 }
 

plugin/plugin.go

@@ -141,6 +141,19 @@ func (p *Plugin) OnBeforeSend(ctx context.Context, msg message.SyncMessage) (mes
 	return msg, nil
 }
 
+// OnSyncFinisher is an interface that can be implemented by a plugin client to be notified when a sync finishes.
+type OnSyncFinisher interface {
+	OnSyncFinish(context.Context) error
+}
+
+// OnSyncFinish gets called after a sync finishes.
+func (p *Plugin) OnSyncFinish(ctx context.Context) error {
+	if v, ok := p.client.(OnSyncFinisher); ok {
+		return v.OnSyncFinish(ctx)
+	}
+	return nil
+}
+
 // IsStaticLinkingEnabled whether static linking is to be enabled
 func (p *Plugin) IsStaticLinkingEnabled() bool {
 	return p.staticLinking

scheduler/scheduler.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/apache/arrow/go/v14/arrow"
 	"runtime/debug"
 	"sync/atomic"
 	"time"
@@ -182,15 +183,23 @@ func (s *Scheduler) Sync(ctx context.Context, client schema.ClientMeta, tables s
 		}
 	}()
 	for resource := range resources {
-		vector := resource.GetValues()
-		bldr := array.NewRecordBuilder(memory.DefaultAllocator, resource.Table.ToArrowSchema())
-		scalar.AppendToRecordBuilder(bldr, vector)
-		rec := bldr.NewRecord()
-		res <- &message.SyncInsert{Record: rec}
+		select {
+		case res <- &message.SyncInsert{Record: resourceToRecord(resource)}:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
 	return nil
 }
 
+func resourceToRecord(resource *schema.Resource) arrow.Record {
+	vector := resource.GetValues()
+	bldr := array.NewRecordBuilder(memory.DefaultAllocator, resource.Table.ToArrowSchema())
+	scalar.AppendToRecordBuilder(bldr, vector)
+	rec := bldr.NewRecord()
+	return rec
+}
+
 func (s *syncClient) logTablesMetrics(tables schema.Tables, client Client) {
 	clientName := client.ID()
 	for _, table := range tables {

scheduler/scheduler_dfs.go

@@ -176,7 +176,10 @@ func (s *syncClient) resolveResourcesDfs(ctx context.Context, table *schema.Tabl
 					atomic.AddUint64(&tableMetrics.Errors, 1)
 					return
 				}
-				resourcesChan <- resolvedResource
+				select {
+				case resourcesChan <- resolvedResource:
+				case <-ctx.Done():
+				}
 			}()
 		}
 		wg.Wait()

scheduler/scheduler_test.go

@@ -2,6 +2,9 @@ package scheduler
 
 import (
 	"context"
+	"fmt"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"testing"
 
 	"github.com/apache/arrow/go/v14/arrow"
@@ -40,6 +43,22 @@ func testColumnResolverPanic(context.Context, schema.ClientMeta, *schema.Resourc
 	panic("ColumnResolver")
 }
 
+func testTableSuccessWithData(data []any) *schema.Table {
+	return &schema.Table{
+		Name: "test_table_success",
+		Resolver: func(_ context.Context, _ schema.ClientMeta, _ *schema.Resource, res chan<- any) error {
+			res <- data
+			return nil
+		},
+		Columns: []schema.Column{
+			{
+				Name: "test_column",
+				Type: arrow.PrimitiveTypes.Int64,
+			},
+		},
+	}
+}
+
 func testTableSuccess() *schema.Table {
 	return &schema.Table{
 		Name:     "test_table_success",
@@ -233,6 +252,72 @@ func TestScheduler(t *testing.T) {
 	}
 }
 
+func TestScheduler_Cancellation(t *testing.T) {
+	data := make([]any, 100)
+
+	tests := []struct {
+		name         string
+		data         []any
+		cancel       bool
+		messageCount int
+	}{
+		{
+			name:         "should consume all message",
+			data:         data,
+			cancel:       false,
+			messageCount: len(data) + 1, // 9 data + 1 migration message
+		},
+		{
+			name:         "should not consume all message on cancel",
+			data:         data,
+			cancel:       true,
+			messageCount: len(data) + 1, // 9 data + 1 migration message
+		},
+	}
+
+	for _, strategy := range AllStrategies {
+		for _, tc := range tests {
+			tc := tc
+			t.Run(fmt.Sprintf("%s_%s", tc.name, strategy.String()), func(t *testing.T) {
+				sc := NewScheduler(WithLogger(zerolog.New(zerolog.NewTestWriter(t))), WithStrategy(strategy))
+
+				messages := make(chan message.SyncMessage)
+				ctx, cancel := context.WithCancel(context.Background())
+				defer cancel()
+
+				go func() {
+					err := sc.Sync(
+						ctx,
+						&testExecutionClient{},
+						[]*schema.Table{testTableSuccessWithData(tc.data)},
+						messages,
+					)
+					if tc.cancel {
+						assert.Equal(t, err, context.Canceled)
+					} else {
+						require.NoError(t, err)
+					}
+					close(messages)
+				}()
+
+				messageConsumed := 0
+				for range messages {
+					if tc.cancel {
+						cancel()
+					}
+					messageConsumed++
+				}
+
+				if tc.cancel {
+					assert.NotEqual(t, tc.messageCount, messageConsumed)
+				} else {
+					assert.Equal(t, tc.messageCount, messageConsumed)
+				}
+			})
+		}
+	}
+}
+
 func testSyncTable(t *testing.T, tc syncTestCase, strategy Strategy, deterministicCQID bool) {
 	ctx := context.Background()
 	tables := []*schema.Table{}