feat!: Support table_concurrency and resource_concurrency (#268)

#### Summary

Closes cloudquery/cloudquery#2313 cloudquery/cloudquery#159 #264

This renames `concurrency` to `table_concurrency` (this has the same behaviour as previously but the name was too generic).

Also this introduce a new option `resource_concurrency` - this limits the number of go routines that resolve a specific resource (useful when a lot of column resolvers that have api calls and `PreResourceResolver`).

Importe Note: both options are only for top-level concurrency control i.e we only spawn go-routines for parent tables and for resources of parent table. The reason for that is that there is no concurrency model that can work with one variable for recursive calls otherwise it will get deadlock.

So right now the SDK will support only for one level - in the future if we want to support additional level we can do `table_concurrency_1` and `resource_concurrency_1` but I think for now we can skip this.

Author: yevgenypats

plugins/source.go

@@ -31,10 +31,6 @@ type SourcePlugin struct {
 	tables schema.Tables
 }
 
-const (
-	defaultConcurrency = 500000
-)
-
 // Add internal columns
 func addInternalColumns(tables []*schema.Table) {
 	for _, table := range tables {
@@ -106,25 +102,25 @@ func (p *SourcePlugin) Version() string {
 
 // Sync is syncing data from the requested tables in spec to the given channel
 func (p *SourcePlugin) Sync(ctx context.Context, logger zerolog.Logger, spec specs.Source, res chan<- *schema.Resource) (*schema.SyncSummary, error) {
-	c, err := p.newExecutionClient(ctx, logger, spec)
+	spec.SetDefaults()
+	if err := spec.Validate(); err != nil {
+		return nil, fmt.Errorf("invalid spec: %w", err)
+	}
+	tableNames, err := p.listAndValidateTables(spec.Tables, spec.SkipTables)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create execution client for source plugin %s: %w", p.name, err)
+		return nil, err
 	}
 
-	// limiter used to limit the amount of resources fetched concurrently
-	concurrency := spec.Concurrency
-	if concurrency == 0 {
-		concurrency = defaultConcurrency
+	c, err := p.newExecutionClient(ctx, logger, spec)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create execution client for source plugin %s: %w", p.name, err)
 	}
-	logger.Info().Uint64("concurrency", concurrency).Msg("starting source plugin sync")
-	goroutinesSem := semaphore.NewWeighted(helpers.Uint64ToInt64(concurrency))
+	logger.Info().Interface("spec", spec).Msg("starting sync")
+	tableSem := semaphore.NewWeighted(helpers.Uint64ToInt64(spec.TableConcurrency))
 	wg := sync.WaitGroup{}
+	resourceSem := semaphore.NewWeighted(helpers.Uint64ToInt64(spec.ResourceConcurrency))
 	summary := schema.SyncSummary{}
 	startTime := time.Now()
-	tableNames, err := p.listAndValidateTables(spec.Tables, spec.SkipTables)
-	if err != nil {
-		return nil, err
-	}
 
 	logger.Debug().Interface("tables", tableNames).Msg("got table names")
 
@@ -141,16 +137,16 @@ func (p *SourcePlugin) Sync(ctx context.Context, logger zerolog.Logger, spec spe
 		for _, client := range clients {
 			client := client
 			wg.Add(1)
-			if err := goroutinesSem.Acquire(ctx, 1); err != nil {
+			if err := tableSem.Acquire(ctx, 1); err != nil {
 				// This means context was cancelled
 				return nil, err
 			}
 			go func() {
 				defer wg.Done()
-				defer goroutinesSem.Release(1)
+				defer tableSem.Release(1)
 				// TODO: prob introduce client.Identify() to be used in logs
 
-				tableSummary := table.Resolve(ctx, client, nil, res)
+				tableSummary := table.Resolve(ctx, client, nil, resourceSem, res)
 				atomic.AddUint64(&summary.Resources, tableSummary.Resources)
 				atomic.AddUint64(&summary.Errors, tableSummary.Errors)
 				atomic.AddUint64(&summary.Panics, tableSummary.Panics)

plugins/source_test.go

@@ -53,8 +53,10 @@ func TestSync(t *testing.T) {
 	)
 
 	spec := specs.Source{
-		Name:   "testSource",
-		Tables: []string{"*"},
+		Name:         "testSource",
+		Tables:       []string{"*"},
+		Version:      "v1.0.0",
+		Destinations: []string{"test"},
 	}
 
 	resources := make(chan *schema.Resource)

schema/table.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cloudquery/plugin-sdk/helpers"
 	"github.com/getsentry/sentry-go"
 	"github.com/thoas/go-funk"
+	"golang.org/x/sync/semaphore"
 )
 
 // TableResolver is the main entry point when a table is sync is called.
@@ -64,9 +65,6 @@ type Table struct {
 	// Parent is the parent table in case this table is called via parent table (i.e. relation)
 	Parent *Table `json:"-"`
 
-	// Serial is used to force a signature change, which forces new table creation and cascading removal of old table and relations
-	Serial string `json:"-"`
-
 	columnsMap map[string]int
 }
 
@@ -165,13 +163,16 @@ func (t Table) TableNames() []string {
 }
 
 // Call the table resolver with with all of it's relation for every reolved resource
-func (t Table) Resolve(ctx context.Context, meta ClientMeta, parent *Resource, resolvedResources chan<- *Resource) (summary SyncSummary) {
+func (t Table) Resolve(ctx context.Context, meta ClientMeta, parent *Resource, resourcesSem *semaphore.Weighted, resolvedResources chan<- *Resource) (summary SyncSummary) {
 	tableStartTime := time.Now()
 	meta.Logger().Info().Str("table", t.Name).Msg("table resolver started")
 
 	res := make(chan interface{})
 	startTime := time.Now()
+	wg := sync.WaitGroup{}
+	wg.Add(1)
 	go func(sum *SyncSummary) {
+		defer wg.Done()
 		defer func() {
 			if err := recover(); err != nil {
 				stack := fmt.Sprintf("%s\n%s", err, string(debug.Stack()))
@@ -199,9 +200,26 @@ func (t Table) Resolve(ctx context.Context, meta ClientMeta, parent *Resource, r
 			continue
 		}
 		for i := range objects {
-			summary.Merge(t.resolveObject(ctx, meta, parent, objects[i], resolvedResources))
+			i := i
+
+			// right now we support concurrency only for objects/resources of parent tables
+			if resourcesSem == nil {
+				summary.Merge(t.resolveObject(ctx, meta, parent, objects[i], resolvedResources))
+			} else {
+				if err := resourcesSem.Acquire(ctx, 1); err != nil {
+					meta.Logger().Error().Err(err).Msg("failed to acquire semaphore")
+					return summary
+				}
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					//nolint:all
+					summary.Merge(t.resolveObject(ctx, meta, parent, objects[i], resolvedResources))
+				}()
+			}
 		}
 	}
+	wg.Wait()
 	meta.Logger().Info().Str("table", t.Name).Int("total_resources", tableResources).TimeDiff("duration", time.Now(), tableStartTime).Msg("fetch table finished")
 
 	return summary
@@ -270,7 +288,7 @@ func (t Table) resolveObject(ctx context.Context, meta ClientMeta, parent *Resou
 	resolvedResources <- resource
 
 	for _, rel := range t.Relations {
-		summary.Merge(rel.Resolve(ctx, meta, resource, resolvedResources))
+		summary.Merge(rel.Resolve(ctx, meta, resource, nil, resolvedResources))
 	}
 
 	return summary

schema/table_test.go

@@ -217,7 +217,7 @@ func TestTableExecution(t *testing.T) {
 			var summary SyncSummary
 			go func() {
 				defer close(resources)
-				summary = tc.Table.Resolve(ctx, m, nil, resources)
+				summary = tc.Table.Resolve(ctx, m, nil, nil, resources)
 			}()
 			var i = uint64(0)
 			for resource := range resources {

serve/source_test.go

@@ -142,11 +142,12 @@ func TestServeSource(t *testing.T) {
 	resources := make(chan []byte, 1)
 	if err := c.Sync(ctx,
 		specs.Source{
-			Name:     "testSourcePlugin",
-			Version:  "v1.0.0",
-			Registry: specs.RegistryGithub,
-			Tables:   []string{"*"},
-			Spec:     TestSourcePluginSpec{Accounts: []string{"cloudquery/plugin-sdk"}},
+			Name:         "testSourcePlugin",
+			Version:      "v1.0.0",
+			Registry:     specs.RegistryGithub,
+			Tables:       []string{"*"},
+			Spec:         TestSourcePluginSpec{Accounts: []string{"cloudquery/plugin-sdk"}},
+			Destinations: []string{"test"},
 		},
 		resources); err != nil {
 		t.Fatal(err)

specs/source.go

@@ -7,6 +7,11 @@ import (
 	"strings"
 )
 
+const (
+	defaultTableConcurrency    = 500000
+	defaultResourceConcurrency = 500000
+)
+
 // Source is the spec for a source plugin
 type Source struct {
 	// Name of the source plugin to use
@@ -20,8 +25,9 @@ type Source struct {
 	// For the gRPC registry the path will be the address of the gRPC server: host:port
 	Path string `json:"path,omitempty"`
 	// Registry can be github,local,grpc.
-	Registry    Registry `json:"registry,omitempty"`
-	Concurrency uint64   `json:"concurrency,omitempty"`
+	Registry            Registry `json:"registry,omitempty"`
+	TableConcurrency    uint64   `json:"table_concurrency,omitempty"`
+	ResourceConcurrency uint64   `json:"resource_concurrency,omitempty"`
 	// Tables to sync from the source plugin
 	Tables []string `json:"tables,omitempty"`
 	// SkipTables defines tables to skip when syncing data. Useful if a glob pattern is used in Tables
@@ -46,6 +52,13 @@ func (s *Source) SetDefaults() {
 	if s.Tables == nil {
 		s.Tables = []string{"*"}
 	}
+
+	if s.TableConcurrency == 0 {
+		s.TableConcurrency = defaultTableConcurrency
+	}
+	if s.ResourceConcurrency == 0 {
+		s.ResourceConcurrency = defaultResourceConcurrency
+	}
 }
 
 // UnmarshalSpec unmarshals the internal spec into the given interface

specs/source_test.go

@@ -117,12 +117,14 @@ spec:
 `,
 		"",
 		&Source{
-			Name:         "test",
-			Registry:     RegistryGithub,
-			Path:         "cloudquery/test",
-			Version:      "v1.1.0",
-			Tables:       []string{"*"},
-			Destinations: []string{"test"},
+			Name:                "test",
+			Registry:            RegistryGithub,
+			Path:                "cloudquery/test",
+			TableConcurrency:    defaultTableConcurrency,
+			ResourceConcurrency: defaultResourceConcurrency,
+			Version:             "v1.1.0",
+			Tables:              []string{"*"},
+			Destinations:        []string{"test"},
 		},
 	},
 }

specs/testdata/dir/aws.yml

@@ -2,6 +2,6 @@ kind: source
 spec:
   name: aws
   version: v1.0.0
-  concurrency: 10
+  table_concurrency: 10
   registry: local
   destinations: [postgresql]

specs/testdata/gcp.yml

@@ -2,7 +2,7 @@ kind: source
 spec:
   name: gcp
   version: v1.0.0
-  concurrency: 10
+  table_concurrency: 10
   registry: local
   destinations: [postgresqlv2]
 ---