Repository: cloudspannerecosystem/harbourbridge
Pull request #486 (2023-03-08 05:10:12)

Parallel schema creation
Create a parallel task processor. Used it for processing tables in parallel.

- [x] Tests pass
- [x] Appropriate changes to README are included in PR


Patch:
diff --git a/common/utils/utils.go b/common/utils/utils.go
index 1e9c2ed..3c5f8d2 100644
--- a/common/utils/utils.go
+++ b/common/utils/utils.go
@@ -561,7 +561,7 @@ func GetLegacyModeSupportedDrivers() []string {
 // ReadSpannerSchema fills conv by querying Spanner infoschema treating Spanner as both the source and dest.
 func ReadSpannerSchema(ctx context.Context, conv *internal.Conv, client *sp.Client) error {
 	infoSchema := spanner.InfoSchemaImpl{Client: client, Ctx: ctx, TargetDb: conv.TargetDb}
-	err := common.ProcessSchema(conv, infoSchema)
+	err := common.ProcessSchema(conv, infoSchema, common.DefaultWorkers)
 	if err != nil {
 		return fmt.Errorf("error trying to read and convert spanner schema: %v", err)
 	}
diff --git a/conversion/conversion.go b/conversion/conversion.go
index 09b91f3..d256309 100644
--- a/conversion/conversion.go
+++ b/conversion/conversion.go
@@ -172,7 +172,7 @@ func schemaFromDatabase(sourceProfile profiles.SourceProfile, targetProfile prof
 	if err != nil {
 		return conv, err
 	}
-	return conv, common.ProcessSchema(conv, infoSchema)
+	return conv, common.ProcessSchema(conv, infoSchema, common.DefaultWorkers)
 }
 
 func performSnapshotMigration(config writer.BatchWriterConfig, conv *internal.Conv, client *sp.Client, infoSchema common.InfoSchema) *writer.BatchWriter {
diff --git a/sources/common/infoschema.go b/sources/common/infoschema.go
index 70b477a..40d4904 100644
--- a/sources/common/infoschema.go
+++ b/sources/common/infoschema.go
@@ -17,6 +17,7 @@ package common
 import (
 	"context"
 	"fmt"
+	"sync"
 
 	sp "cloud.google.com/go/spanner"
 
@@ -25,6 +26,8 @@ import (
 	"github.com/cloudspannerecosystem/harbourbridge/spanner/ddl"
 )
 
+const DefaultWorkers = 20 // Default to 20 - observed diminishing returns above this value
+
 // InfoSchema contains database information.
 type InfoSchema interface {
 	GetToDdl() ToDdl
@@ -58,18 +61,36 @@ type FkConstraint struct {
 // ProcessSchema performs schema conversion for source database
 // 'db'. Information schema tables are a broadly supported ANSI standard,
 // and we use them to obtain source database's schema information.
-func ProcessSchema(conv *internal.Conv, infoSchema InfoSchema) error {
+// numWorkers decides the parallelism while converting tables
+func ProcessSchema(conv *internal.Conv, infoSchema InfoSchema, numWorkers int) error {
 	tables, err := infoSchema.GetTables()
+	fmt.Println("fetched tables", tables)
 	if err != nil {
 		return err
 	}
-	for _, t := range tables {
-		if err := processTable(conv, t, infoSchema); err != nil {
-			return err
-		}
+
+	if numWorkers < 1 {
+		numWorkers = DefaultWorkers
+	}
+
+	asyncProcessTable := func(t SchemaAndName, mutex *sync.Mutex) TaskResult[SchemaAndName] {
+		table, e := processTable(conv, t, infoSchema)
+		mutex.Lock()
+		conv.SrcSchema[table.Name] = table
+		mutex.Unlock()
+		res := TaskResult[SchemaAndName]{t, e}
+		return res
 	}
+
+	res, e := RunParallelTasks(tables, numWorkers, asyncProcessTable, true)
+	if e != nil {
+		fmt.Println("exiting due to error while processing schema for table", res)
+		return err
+	}
+
 	SchemaToSpannerDDL(conv, infoSchema.GetToDdl())
 	conv.AddPrimaryKeys()
+	fmt.Println("loaded schema")
 	return nil
 }
 
@@ -123,29 +144,31 @@ func SetRowStats(conv *internal.Conv, infoSchema InfoSchema) {
 	}
 }
 
-func processTable(conv *internal.Conv, table SchemaAndName, infoSchema InfoSchema) error {
+func processTable(conv *internal.Conv, table SchemaAndName, infoSchema InfoSchema) (schema.Table, error) {
+	var t schema.Table
+	fmt.Println("processing schema for table", table)
 	primaryKeys, constraints, err := infoSchema.GetConstraints(conv, table)
 	if err != nil {
-		return fmt.Errorf("couldn't get constraints for table %s.%s: %s", table.Schema, table.Name, err)
+		return t, fmt.Errorf("couldn't get constraints for table %s.%s: %s", table.Schema, table.Name, err)
 	}
 	foreignKeys, err := infoSchema.GetForeignKeys(conv, table)
 	if err != nil {
-		return fmt.Errorf("couldn't get foreign key constraints for table %s.%s: %s", table.Schema, table.Name, err)
+		return t, fmt.Errorf("couldn't get foreign key constraints for table %s.%s: %s", table.Schema, table.Name, err)
 	}
 	indexes, err := infoSchema.GetIndexes(conv, table)
 	if err != nil {
-		return fmt.Errorf("couldn't get indexes for table %s.%s: %s", table.Schema, table.Name, err)
+		return t, fmt.Errorf("couldn't get indexes for table %s.%s: %s", table.Schema, table.Name, err)
 	}
 	colDefs, colNames, err := infoSchema.GetColumns(conv, table, constraints, primaryKeys)
 	if err != nil {
-		return fmt.Errorf("couldn't get schema for table %s.%s: %s", table.Schema, table.Name, err)
+		return t, fmt.Errorf("couldn't get schema for table %s.%s: %s", table.Schema, table.Name, err)
 	}
 	name := infoSchema.GetTableName(table.Schema, table.Name)
 	var schemaPKeys []schema.Key
 	for _, k := range primaryKeys {
 		schemaPKeys = append(schemaPKeys, schema.Key{Column: k})
 	}
-	conv.SrcSchema[name] = schema.Table{
+	t = schema.Table{
 		Name:        name,
 		Schema:      table.Schema,
 		ColNames:    colNames,
@@ -153,5 +176,5 @@ func processTable(conv *internal.Conv, table SchemaAndName, infoSchema InfoSchem
 		PrimaryKeys: schemaPKeys,
 		Indexes:     indexes,
 		ForeignKeys: foreignKeys}
-	return nil
+	return t, nil
 }
diff --git a/sources/common/utils.go b/sources/common/utils.go
index fdc08ed..ec8851b 100644
--- a/sources/common/utils.go
+++ b/sources/common/utils.go
@@ -16,8 +16,10 @@ package common
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/cloudspannerecosystem/harbourbridge/internal"
+	"github.com/cloudspannerecosystem/harbourbridge/logger"
 	"github.com/cloudspannerecosystem/harbourbridge/schema"
 	"github.com/cloudspannerecosystem/harbourbridge/spanner/ddl"
 )
@@ -46,3 +48,66 @@ func GetColsAndSchemas(conv *internal.Conv, srcTable string) (schema.Table, stri
 	}
 	return srcSchema, spTable, spCols, spSchema, err
 }
+
+type TaskResult[O any] struct {
+	Result O
+	Err    error
+}
+
+// Run multiple tasks in parallel. The tasks are expected to be thread safe
+// input: 		List of inputs
+// numWorkers: 	Size of worker pool
+// f:			Function to execute the task
+// fastExit: 	If an error is encountered, gracefully exit all running or queued tasks
+// Returns an array of TaskResults and last error
+func RunParallelTasks[I any, O any](input []I, numWorkers int, f func(i I, mutex *sync.Mutex) TaskResult[O],
+	fastExit bool) ([]TaskResult[O], error) {
+	inputChannel := make(chan I, len(input))
+	outputChannel := make(chan TaskResult[O], len(input))
+
+	wg := &sync.WaitGroup{}
+	defer func() {
+		for range inputChannel {
+			logger.Log.Debug(fmt.Sprint("clearing out pending tasks"))
+		}
+		wg.Wait()
+	}()
+
+	mutex := &sync.Mutex{}
+	logger.Log.Debug(fmt.Sprint("initiating workers", numWorkers))
+	for w := 0; w < numWorkers; w++ {
+		go processAsync(f, inputChannel, outputChannel, mutex, wg)
+	}
+
+	for _, in := range input {
+		inputChannel <- in
+	}
+	close(inputChannel)
+
+	out := []TaskResult[O]{}
+	for i := 0; i < len(input); i++ {
+		res := <-outputChannel
+		if fastExit && res.Err != nil {
+			logger.Log.Debug(fmt.Sprint("stopping worker pool due to encountered error", res.Err))
+			for range inputChannel {
+				logger.Log.Debug(fmt.Sprint("ignoring task to fast exit"))
+			}
+			wg.Wait()
+			return out, res.Err
+		}
+		out = append(out, res)
+	}
+	logger.Log.Debug(fmt.Sprint("completed processing of %n tasks", len(out)))
+	return out, nil
+}
+
+func processAsync[I any, O any](f func(i I, mutex *sync.Mutex) TaskResult[O], in chan I,
+	out chan TaskResult[O], mutex *sync.Mutex, wg *sync.WaitGroup) {
+	wg.Add(1)
+	for i := range in {
+		logger.Log.Debug(fmt.Sprint("processing task for input", i))
+		out <- f(i, mutex)
+	}
+	logger.Log.Debug(fmt.Sprint("worker complete"))
+	wg.Done()
+}
diff --git a/web/web.go b/web/web.go
index a0a2964..8970f00 100644
--- a/web/web.go
+++ b/web/web.go
@@ -144,13 +144,13 @@ func convertSchemaSQL(w http.ResponseWriter, r *http.Request) {
 	var err error
 	switch sessionState.driver {
 	case constants.MYSQL:
-		err = common.ProcessSchema(conv, mysql.InfoSchemaImpl{DbName: sessionState.dbName, Db: sessionState.sourceDB})
+		err = common.ProcessSchema(conv, mysql.InfoSchemaImpl{DbName: sessionState.dbName, Db: sessionState.sourceDB}, common.DefaultWorkers)
 	case constants.POSTGRES:
-		err = common.ProcessSchema(conv, postgres.InfoSchemaImpl{Db: sessionState.sourceDB})
+		err = common.ProcessSchema(conv, postgres.InfoSchemaImpl{Db: sessionState.sourceDB}, common.DefaultWorkers)
 	case constants.SQLSERVER:
-		err = common.ProcessSchema(conv, sqlserver.InfoSchemaImpl{DbName: sessionState.dbName, Db: sessionState.sourceDB})
+		err = common.ProcessSchema(conv, sqlserver.InfoSchemaImpl{DbName: sessionState.dbName, Db: sessionState.sourceDB}, common.DefaultWorkers)
 	case constants.ORACLE:
-		err = common.ProcessSchema(conv, oracle.InfoSchemaImpl{DbName: strings.ToUpper(sessionState.dbName), Db: sessionState.sourceDB})
+		err = common.ProcessSchema(conv, oracle.InfoSchemaImpl{DbName: strings.ToUpper(sessionState.dbName), Db: sessionState.sourceDB}, common.DefaultWorkers)
 	default:
 		http.Error(w, fmt.Sprintf("Driver : '%s' is not supported", sessionState.driver), http.StatusBadRequest)
 		return
diff --git a/webv2/web.go b/webv2/web.go
index 96f05de..dabdaff 100644
--- a/webv2/web.go
+++ b/webv2/web.go
@@ -123,15 +123,15 @@ type progressDetails struct {
 }
 
 type migrationDetails struct {
-	TargetDetails targetDetails `json:"TargetDetails"`
+	TargetDetails  targetDetails  `json:"TargetDetails"`
 	DataflowConfig dataflowConfig `json:"DataflowConfig"`
-	MigrationMode string        `json:MigrationMode`
-	MigrationType string        `json:MigrationType`
+	MigrationMode  string         `json:MigrationMode`
+	MigrationType  string         `json:MigrationType`
 }
 
 type dataflowConfig struct {
-	Network string `json:Network`
-	Subnetwork string  `json:Subnetwork`
+	Network    string `json:Network`
+	Subnetwork string `json:Subnetwork`
 }
 
 type targetDetails struct {
@@ -147,9 +147,9 @@ type StreamingCfg struct {
 	TmpDir        string        `json:"tmpDir"`
 }
 type DataflowCfg struct {
-	JobName  string `json:"JobName"`
-	Location string `json:"Location"`
-	Network string `json:"Network"`
+	JobName    string `json:"JobName"`
+	Location   string `json:"Location"`
+	Network    string `json:"Network"`
 	Subnetwork string `json:"Subnetwork"`
 }
 type ConnectionConfig struct {
@@ -240,13 +240,13 @@ func convertSchemaSQL(w http.ResponseWriter, r *http.Request) {
 	var err error
 	switch sessionState.Driver {
 	case constants.MYSQL:
-		err = common.ProcessSchema(conv, mysql.InfoSchemaImpl{DbName: sessionState.DbName, Db: sessionState.SourceDB})
+		err = common.ProcessSchema(conv, mysql.InfoSchemaImpl{DbName: sessionState.DbName, Db: sessionState.SourceDB}, common.DefaultWorkers)
 	case constants.POSTGRES:
-		err = common.ProcessSchema(conv, postgres.InfoSchemaImpl{Db: sessionState.SourceDB})
+		err = common.ProcessSchema(conv, postgres.InfoSchemaImpl{Db: sessionState.SourceDB}, common.DefaultWorkers)
 	case constants.SQLSERVER:
-		err = common.ProcessSchema(conv, sqlserver.InfoSchemaImpl{DbName: sessionState.DbName, Db: sessionState.SourceDB})
+		err = common.ProcessSchema(conv, sqlserver.InfoSchemaImpl{DbName: sessionState.DbName, Db: sessionState.SourceDB}, common.DefaultWorkers)
 	case constants.ORACLE:
-		err = common.ProcessSchema(conv, oracle.InfoSchemaImpl{DbName: strings.ToUpper(sessionState.DbName), Db: sessionState.SourceDB})
+		err = common.ProcessSchema(conv, oracle.InfoSchemaImpl{DbName: strings.ToUpper(sessionState.DbName), Db: sessionState.SourceDB}, common.DefaultWorkers)
 	default:
 		http.Error(w, fmt.Sprintf("Driver : '%s' is not supported", sessionState.Driver), http.StatusBadRequest)
 		return
@@ -1745,9 +1745,9 @@ func createStreamingCfgFile(sessionState *session.SessionState, targetDetails ta
 			},
 		},
 		DataflowCfg: DataflowCfg{
-			JobName:  "",
-			Location: sessionState.Region,
-			Network: dataflowConfig.Network,
+			JobName:    "",
+			Location:   sessionState.Region,
+			Network:    dataflowConfig.Network,
 			Subnetwork: dataflowConfig.Subnetwork,
 		},
 		TmpDir: "gs://" + sessionState.Bucket + sessionState.RootPath,


Test patch:
diff --git a/sources/common/utils_test.go b/sources/common/utils_test.go
index b0a8e85..cca0b74 100644
--- a/sources/common/utils_test.go
+++ b/sources/common/utils_test.go
@@ -15,8 +15,12 @@
 package common
 
 import (
+	"fmt"
+	"math/rand"
 	"reflect"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/zap"
@@ -121,3 +125,17 @@ func TestGetColsAndSchemas(t *testing.T) {
 		})
 	}
 }
+
+func TestWorkerPool(t *testing.T) {
+	input := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
+
+	f := func(i int, mutex *sync.Mutex) TaskResult[int] {
+		sleepTime := time.Duration(rand.Intn(1000 * 1000))
+		time.Sleep(sleepTime)
+		res := TaskResult[int]{Result: i, Err: nil}
+		return res
+	}
+
+	out, _ := RunParallelTasks(input, 5, f, false)
+	assert.Equal(t, len(input), len(out), fmt.Sprintln("jobs not processed"))
+}
diff --git a/sources/dynamodb/schema_test.go b/sources/dynamodb/schema_test.go
index a5ee813..53708b5 100644
--- a/sources/dynamodb/schema_test.go
+++ b/sources/dynamodb/schema_test.go
@@ -188,7 +188,7 @@ func TestProcessSchema(t *testing.T) {
 	sampleSize := int64(10000)
 
 	conv := internal.MakeConv()
-	err := common.ProcessSchema(conv, InfoSchemaImpl{client, nil, sampleSize})
+	err := common.ProcessSchema(conv, InfoSchemaImpl{client, nil, sampleSize}, 1)
 
 	assert.Nil(t, err)
 	expectedSchema := map[string]ddl.CreateTable{
@@ -291,7 +291,7 @@ func TestProcessSchema_FullDataTypes(t *testing.T) {
 	sampleSize := int64(10000)
 
 	conv := internal.MakeConv()
-	err := common.ProcessSchema(conv, InfoSchemaImpl{client, nil, sampleSize})
+	err := common.ProcessSchema(conv, InfoSchemaImpl{client, nil, sampleSize}, 1)
 
 	assert.Nil(t, err)
 	expectedSchema := map[string]ddl.CreateTable{
diff --git a/sources/mysql/infoschema_test.go b/sources/mysql/infoschema_test.go
index 9e21ce2..13418bc 100644
--- a/sources/mysql/infoschema_test.go
+++ b/sources/mysql/infoschema_test.go
@@ -197,7 +197,7 @@ func TestProcessSchemaMYSQL(t *testing.T) {
 	db := mkMockDB(t, ms)
 	conv := internal.MakeConv()
 	isi := InfoSchemaImpl{"test", db, profiles.SourceProfile{}, profiles.TargetProfile{}}
-	err := common.ProcessSchema(conv, isi)
+	err := common.ProcessSchema(conv, isi, 1)
 	assert.Nil(t, err)
 	expectedSchema := map[string]ddl.CreateTable{
 		"user": ddl.CreateTable{
@@ -378,7 +378,7 @@ func TestProcessData_MultiCol(t *testing.T) {
 	db := mkMockDB(t, ms)
 	conv := internal.MakeConv()
 	isi := InfoSchemaImpl{"test", db, profiles.SourceProfile{}, profiles.TargetProfile{}}
-	err := common.ProcessSchema(conv, isi)
+	err := common.ProcessSchema(conv, isi, 1)
 	assert.Nil(t, err)
 	expectedSchema := map[string]ddl.CreateTable{
 		"test": ddl.CreateTable{
diff --git a/sources/oracle/infoschema_test.go b/sources/oracle/infoschema_test.go
index 6a254d9..e8386d5 100644
--- a/sources/oracle/infoschema_test.go
+++ b/sources/oracle/infoschema_test.go
@@ -153,7 +153,7 @@ func TestProcessSchemaOracle(t *testing.T) {
 	}
 	db := mkMockDB(t, ms)
 	conv := internal.MakeConv()
-	err := common.ProcessSchema(conv, InfoSchemaImpl{"test", db, profiles.SourceProfile{}, profiles.TargetProfile{}})
+	err := common.ProcessSchema(conv, InfoSchemaImpl{"test", db, profiles.SourceProfile{}, profiles.TargetProfile{}}, 1)
 	assert.Nil(t, err)
 	expectedSchema := map[string]ddl.CreateTable{
 		"USER": {
diff --git a/sources/postgres/infoschema_test.go b/sources/postgres/infoschema_test.go
index 1031edf..bca438a 100644
--- a/sources/postgres/infoschema_test.go
+++ b/sources/postgres/infoschema_test.go
@@ -202,7 +202,7 @@ func TestProcessSchema(t *testing.T) {
 	}
 	db := mkMockDB(t, ms)
 	conv := internal.MakeConv()
-	err := common.ProcessSchema(conv, InfoSchemaImpl{db, profiles.SourceProfile{}, profiles.TargetProfile{}})
+	err := common.ProcessSchema(conv, InfoSchemaImpl{db, profiles.SourceProfile{}, profiles.TargetProfile{}}, 1)
 	assert.Nil(t, err)
 	expectedSchema := map[string]ddl.CreateTable{
 		"user": ddl.CreateTable{
@@ -458,7 +458,7 @@ func TestConvertSqlRow_MultiCol(t *testing.T) {
 	}
 	db := mkMockDB(t, ms)
 	conv := internal.MakeConv()
-	err := common.ProcessSchema(conv, InfoSchemaImpl{db, profiles.SourceProfile{}, profiles.TargetProfile{}})
+	err := common.ProcessSchema(conv, InfoSchemaImpl{db, profiles.SourceProfile{}, profiles.TargetProfile{}}, 1)
 	assert.Nil(t, err)
 	conv.SetDataMode()
 	var rows []spannerData
diff --git a/sources/sqlserver/infoschema_test.go b/sources/sqlserver/infoschema_test.go
index b85dac5..0e638a2 100644
--- a/sources/sqlserver/infoschema_test.go
+++ b/sources/sqlserver/infoschema_test.go
@@ -227,7 +227,7 @@ func TestProcessSchema(t *testing.T) {
 	}
 	db := mkMockDB(t, ms)
 	conv := internal.MakeConv()
-	err := common.ProcessSchema(conv, InfoSchemaImpl{"test", db})
+	err := common.ProcessSchema(conv, InfoSchemaImpl{"test", db}, 1)
 	assert.Nil(t, err)
 	expectedSchema := map[string]ddl.CreateTable{
 		"user": {