feat: Add sync sharding (#1891)

erezrokah · web-flow · commit e1823f82fd3c · 2024-09-18T10:14:14.000Z
#### Summary Goes with cloudquery/plugin-pb-go#401. ~~Still testing this so in draft~~ Part of cloudquery/cloudquery-issues#2214 (internal issue) ---
diff --git a/examples/simple_plugin/go.mod b/examples/simple_plugin/go.mod
@@ -54,6 +54,7 @@ require (
 	github.com/oapi-codegen/runtime v1.1.1 // indirect
 	github.com/pierrec/lz4/v4 v4.1.21 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/samber/lo v1.47.0 // indirect
 	github.com/santhosh-tekuri/jsonschema/v6 v6.0.1 // indirect
 	github.com/spf13/cobra v1.8.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
diff --git a/examples/simple_plugin/go.sum b/examples/simple_plugin/go.sum
@@ -126,6 +126,8 @@ github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
 github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
 github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/samber/lo v1.47.0 h1:z7RynLwP5nbyRscyvcD043DWYoOcYRv3mV8lBeqOCLc=
+github.com/samber/lo v1.47.0/go.mod h1:RmDH9Ct32Qy3gduHQuKJ3gW1fMHAnE/fAzQuf6He5cU=
 github.com/santhosh-tekuri/jsonschema/v6 v6.0.1 h1:PKK9DyHxif4LZo+uQSgXNqs0jj5+xZwwfKHgph2lxBw=
 github.com/santhosh-tekuri/jsonschema/v6 v6.0.1/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU=
 github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
diff --git a/go.mod b/go.mod
@@ -22,6 +22,7 @@ require (
 	github.com/hashicorp/go-retryablehttp v0.7.7
 	github.com/invopop/jsonschema v0.12.0
 	github.com/rs/zerolog v1.33.0
+	github.com/samber/lo v1.47.0
 	github.com/santhosh-tekuri/jsonschema/v6 v6.0.1
 	github.com/spf13/cobra v1.8.1
 	github.com/stretchr/testify v1.9.0
diff --git a/go.sum b/go.sum
@@ -126,6 +126,8 @@ github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
 github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
 github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/samber/lo v1.47.0 h1:z7RynLwP5nbyRscyvcD043DWYoOcYRv3mV8lBeqOCLc=
+github.com/samber/lo v1.47.0/go.mod h1:RmDH9Ct32Qy3gduHQuKJ3gW1fMHAnE/fAzQuf6He5cU=
 github.com/santhosh-tekuri/jsonschema/v6 v6.0.1 h1:PKK9DyHxif4LZo+uQSgXNqs0jj5+xZwwfKHgph2lxBw=
 github.com/santhosh-tekuri/jsonschema/v6 v6.0.1/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU=
 github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
diff --git a/internal/servers/plugin/v3/plugin.go b/internal/servers/plugin/v3/plugin.go
@@ -178,6 +178,12 @@ func (s *Server) Sync(req *pb.Sync_Request, stream pb.Plugin_SyncServer) error {
 			Connection: req.Backend.Connection,
 		}
 	}
+	if req.Shard != nil {
+		syncOptions.Shard = &plugin.Shard{
+			Num:   req.Shard.Num,
+			Total: req.Shard.Total,
+		}
+	}
 
 	go func() {
 		defer flushMetrics()
diff --git a/plugin/plugin_source.go b/plugin/plugin_source.go
@@ -15,12 +15,18 @@ type BackendOptions struct {
 	Connection string
 }
 
+type Shard struct {
+	Num   int32
+	Total int32
+}
+
 type SyncOptions struct {
 	Tables              []string
 	SkipTables          []string
 	SkipDependentTables bool
 	DeterministicCQID   bool
 	BackendOptions      *BackendOptions
+	Shard               *Shard
 }
 
 type SourceClient interface {
diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go
@@ -13,6 +13,7 @@ import (
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/rs/zerolog"
+	"github.com/samber/lo"
 	"github.com/thoas/go-funk"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
@@ -90,6 +91,12 @@ func WithInvocationID(invocationID string) Option {
 	}
 }
 
+func WithShard(num int32, total int32) SyncOption {
+	return func(s *syncClient) {
+		s.shard = &shard{num: num, total: total}
+	}
+}
+
 type Client interface {
 	ID() string
 }
@@ -119,6 +126,11 @@ type Scheduler struct {
 	invocationID string
 }
 
+type shard struct {
+	num   int32
+	total int32
+}
+
 type syncClient struct {
 	tables            schema.Tables
 	client            schema.ClientMeta
@@ -128,6 +140,8 @@ type syncClient struct {
 	metrics      *Metrics
 	logger       zerolog.Logger
 	invocationID string
+
+	shard *shard
 }
 
 func NewScheduler(opts ...Option) *Scheduler {
@@ -346,3 +360,24 @@ func maxDepth(tables schema.Tables) uint64 {
 	}
 	return depth
 }
+
+func shardTableClients(tableClients []tableClient, shard *shard) []tableClient {
+	// For sharding to work as expected, tableClients must be deterministic between different shards.
+	if shard == nil || len(tableClients) == 0 {
+		return tableClients
+	}
+	num := int(shard.num)
+	total := int(shard.total)
+	chunkSize := len(tableClients) / total
+	if chunkSize == 0 {
+		chunkSize = 1
+	}
+	chunks := lo.Chunk(tableClients, chunkSize)
+	if num > len(chunks) {
+		return nil
+	}
+	if len(chunks) > total && num == total {
+		return append(chunks[num-1], chunks[num]...)
+	}
+	return chunks[num-1]
+}
diff --git a/scheduler/scheduler_debug.go b/scheduler/scheduler_debug.go
@@ -67,6 +67,7 @@ func (s *syncClient) syncTest(ctx context.Context, syncMultiplier int, resolvedR
 		}
 	}
 	shuffle(allClients, seed)
+	allClients = shardTableClients(allClients, s.shard)
 
 	var wg sync.WaitGroup
 	for _, tc := range allClients {
diff --git a/scheduler/scheduler_dfs.go b/scheduler/scheduler_dfs.go
@@ -40,27 +40,34 @@ func (s *syncClient) syncDfs(ctx context.Context, resolvedResources chan<- *sche
 		s.metrics.initWithClients(table, clients)
 	}
 
-	var wg sync.WaitGroup
+	tableClients := make([]tableClient, 0)
 	for i, table := range s.tables {
-		table := table
-		clients := preInitialisedClients[i]
-		for _, client := range clients {
-			client := client
-			if err := s.scheduler.tableSems[0].Acquire(ctx, 1); err != nil {
-				// This means context was cancelled
-				wg.Wait()
-				return
-			}
-			wg.Add(1)
-			go func() {
-				defer wg.Done()
-				defer s.scheduler.tableSems[0].Release(1)
-				// not checking for error here as nothing much todo.
-				// the error is logged and this happens when context is cancelled
-				s.resolveTableDfs(ctx, table, client, nil, resolvedResources, 1)
-			}()
+		for _, client := range preInitialisedClients[i] {
+			tableClients = append(tableClients, tableClient{table: table, client: client})
 		}
 	}
+	tableClients = shardTableClients(tableClients, s.shard)
+
+	var wg sync.WaitGroup
+	for _, tc := range tableClients {
+		table := tc.table
+		cl := tc.client
+		if err := s.scheduler.tableSems[0].Acquire(ctx, 1); err != nil {
+			// This means context was cancelled
+			wg.Wait()
+			return
+		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer s.scheduler.tableSems[0].Release(1)
+			// not checking for error here as nothing much to do.
+			// the error is logged and this happens when context is cancelled
+			// Round Robin currently uses the DFS algorithm to resolve the tables, but this
+			// may change in the future.
+			s.resolveTableDfs(ctx, table, cl, nil, resolvedResources, 1)
+		}()
+	}
 
 	// Wait for all the worker goroutines to finish
 	wg.Wait()
diff --git a/scheduler/scheduler_round_robin.go b/scheduler/scheduler_round_robin.go
@@ -37,6 +37,7 @@ func (s *syncClient) syncRoundRobin(ctx context.Context, resolvedResources chan<
 	}
 
 	tableClients := roundRobinInterleave(s.tables, preInitialisedClients)
+	tableClients = shardTableClients(tableClients, s.shard)
 
 	var wg sync.WaitGroup
 	for _, tc := range tableClients {
diff --git a/scheduler/scheduler_shuffle.go b/scheduler/scheduler_shuffle.go
@@ -44,6 +44,7 @@ func (s *syncClient) syncShuffle(ctx context.Context, resolvedResources chan<- *
 	// however, if the table order changes, the seed will change and the shuffle order will be different,
 	// so users have a little bit of control over the randomization.
 	seed := hashTableNames(tableNames)
+	tableClients = shardTableClients(tableClients, s.shard)
 	shuffle(tableClients, seed)
 
 	var wg sync.WaitGroup
diff --git a/scheduler/scheduler_test.go b/scheduler/scheduler_test.go
@@ -481,3 +481,120 @@ func TestScheduler_Cancellation(t *testing.T) {
 		}
 	}
 }
+
+func Test_shardTableClients(t *testing.T) {
+	type testCase struct {
+		name         string
+		tableClients []tableClient
+		shard        *shard
+		expected     []tableClient
+	}
+
+	tests := []testCase{
+		{
+			name: "nil shard returns all table clients",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+			},
+			expected: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+			},
+		},
+		{
+			name:         "nil table clients",
+			tableClients: nil,
+			shard:        &shard{num: 1, total: 2},
+			expected:     nil,
+		},
+		{
+			name:         "empty table clients",
+			tableClients: []tableClient{},
+			shard:        &shard{num: 1, total: 2},
+			expected:     []tableClient{},
+		},
+		{
+			name: "even shard 1 of 2",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+			},
+			shard: &shard{num: 1, total: 2},
+			expected: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+			},
+		},
+		{
+			name: "even shard 2 of 2",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+			},
+			shard: &shard{num: 2, total: 2},
+			expected: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+			},
+		},
+		{
+			name: "uneven split 1 of 2",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
+			},
+			shard: &shard{num: 1, total: 2},
+			expected: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+			},
+		},
+		{
+			name: "uneven split 2 of 2",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
+			},
+			shard: &shard{num: 2, total: 2},
+			expected: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
+			},
+		},
+		{
+			name: "more shards than table clients",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+			},
+			shard:    &shard{num: 5, total: 100},
+			expected: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			actual := shardTableClients(tc.tableClients, tc.shard)
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}

Original file line number	Diff line number	Diff line change
`@@ -178,6 +178,12 @@ func (s Server) Sync(req pb.Sync_Request, stream pb.Plugin_SyncServer) error {`
`178`	`178`	`Connection: req.Backend.Connection,`
`179`	`179`	`}`
`180`	`180`	`}`
	`181`	`+ if req.Shard != nil {`
	`182`	`+ syncOptions.Shard = &plugin.Shard{`
	`183`	`+ Num: req.Shard.Num,`
	`184`	`+ Total: req.Shard.Total,`
	`185`	`+ }`
	`186`	`+ }`
`181`	`187`
`182`	`188`	`go func() {`
`183`	`189`	`defer flushMetrics()`
Original file line number	Diff line number	Diff line change
`@@ -67,6 +67,7 @@ func (s *syncClient) syncTest(ctx context.Context, syncMultiplier int, resolvedR`
`67`	`67`	`}`
`68`	`68`	`}`
`69`	`69`	`shuffle(allClients, seed)`
	`70`	`+ allClients = shardTableClients(allClients, s.shard)`
`70`	`71`
`71`	`72`	`var wg sync.WaitGroup`
`72`	`73`	`for _, tc := range allClients {`
Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ func (s *syncClient) syncRoundRobin(ctx context.Context, resolvedResources chan<`
`37`	`37`	`}`
`38`	`38`
`39`	`39`	`tableClients := roundRobinInterleave(s.tables, preInitialisedClients)`
	`40`	`+ tableClients = shardTableClients(tableClients, s.shard)`
`40`	`41`
`41`	`42`	`var wg sync.WaitGroup`
`42`	`43`	`for _, tc := range tableClients {`