Import-data: Reuse worker pool across tasks, implement sequential task picker (#2257)

- Re-use worker pool for all tasks pertaining to all tables.
- Implement Sequential Task Picker, that picks one task at a time.

Author: makalaaneesh

migtests/tests/import-file/validate

@@ -132,7 +132,9 @@ TABLE_NAME_DATA_MAP = {
 def file_import_done_checks(tgt):
 	fetched_row_cnts = tgt.row_count_of_all_tables()
 	print(f"Row counts after import data file: {fetched_row_cnts}")
-	assert fetched_row_cnts == EXPECTED
+	for table_name, rc in EXPECTED.items():
+		assert fetched_row_cnts[table_name] == rc, f"Row count mismatch for table {table_name}, expected: {rc}, got: {fetched_row_cnts[table_name]}"
+
 
 	fetched_row_cnts_non_public = tgt.row_count_of_all_tables("non_public")
 	print(f"Row counts after import data file in non_public schema: {fetched_row_cnts_non_public}")

yb-voyager/cmd/importData.go

@@ -155,6 +155,7 @@ func importDataCommandFn(cmd *cobra.Command, args []string) {
 	// TODO: handle case-sensitive in table names with oracle ff-db
 	// quoteTableNameIfRequired()
 	importFileTasks := discoverFilesToImport()
+	log.Debugf("Discovered import file tasks: %v", importFileTasks)
 	if importerRole == TARGET_DB_IMPORTER_ROLE {
 
 		importType = record.ExportType
@@ -316,6 +317,10 @@ type ImportFileTask struct {
 	FileSize     int64
 }
 
+func (task *ImportFileTask) String() string {
+	return fmt.Sprintf("{ID: %d, FilePath: %s, TableName: %s, RowCount: %d, FileSize: %d}", task.ID, task.FilePath, task.TableNameTup.ForOutput(), task.RowCount, task.FileSize)
+}
+
 // func quoteTableNameIfRequired() {
 // 	if tconf.TargetDBType != ORACLE {
 // 		return
@@ -543,6 +548,8 @@ func importData(importFileTasks []*ImportFileTask) {
 			utils.ErrExit("Failed to classify tasks: %s", err)
 		}
 	}
+	log.Infof("pending tasks: %v", pendingTasks)
+	log.Infof("completed tasks: %v", completedTasks)
 
 	//TODO: BUG: we are applying table-list filter on importFileTasks, but here we are considering all tables as per
 	// export-data table-list. Should be fine because we are only disabling and re-enabling, but this is still not ideal.
@@ -584,28 +591,36 @@ func importData(importFileTasks []*ImportFileTask) {
 				controlPlane.UpdateImportedRowCount(importDataAllTableMetrics)
 			}
 
-			for _, task := range pendingTasks {
-				// The code can produce `poolSize` number of batches at a time. But, it can consume only
-				// `parallelism` number of batches at a time.
-				batchImportPool = pool.New().WithMaxGoroutines(poolSize)
-				log.Infof("created batch import pool of size: %d", poolSize)
-
-				taskImporter, err := NewFileTaskImporter(task, state, batchImportPool, progressReporter)
+			useTaskPicker := utils.GetEnvAsBool("USE_TASK_PICKER_FOR_IMPORT", true)
+			if useTaskPicker {
+				err := importTasksViaTaskPicker(pendingTasks, state, progressReporter, poolSize)
 				if err != nil {
-					utils.ErrExit("Failed to create file task importer: %s", err)
+					utils.ErrExit("Failed to import tasks via task picker: %s", err)
 				}
+			} else {
+				for _, task := range pendingTasks {
+					// The code can produce `poolSize` number of batches at a time. But, it can consume only
+					// `parallelism` number of batches at a time.
+					batchImportPool = pool.New().WithMaxGoroutines(poolSize)
+					log.Infof("created batch import pool of size: %d", poolSize)
 
-				for !taskImporter.AllBatchesSubmitted() {
-					err := taskImporter.SubmitNextBatch()
+					taskImporter, err := NewFileTaskImporter(task, state, batchImportPool, progressReporter)
 					if err != nil {
-						utils.ErrExit("Failed to submit next batch: task:%v err: %s", task, err)
+						utils.ErrExit("Failed to create file task importer: %s", err)
 					}
-				}
 
-				batchImportPool.Wait() // Wait for the file import to finish.
-				taskImporter.PostProcess()
+					for !taskImporter.AllBatchesSubmitted() {
+						err := taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
+						if err != nil {
+							utils.ErrExit("Failed to submit next batch: task:%v err: %s", task, err)
+						}
+					}
+
+					batchImportPool.Wait() // wait for file import to finish
+					taskImporter.PostProcess()
+				}
+				time.Sleep(time.Second * 2)
 			}
-			time.Sleep(time.Second * 2)
 		}
 		utils.PrintAndLog("snapshot data import complete\n\n")
 	}
@@ -684,6 +699,75 @@ func importData(importFileTasks []*ImportFileTask) {
 
 }
 
+/*
+1. Initialize a worker pool
+2. Create a task picker which helps the importer choose which task to process in each iteration.
+3. Loop until all tasks are done:
+  - Pick a task from the task picker.
+  - If the task is not already being processed, create a new FileTaskImporter for the task.
+  - For the task that is picked, produce the next batch and submit it to the worker pool. Worker will asynchronously import the batch.
+  - If task is done, mark it as done in the task picker.
+*/
+func importTasksViaTaskPicker(pendingTasks []*ImportFileTask, state *ImportDataState, progressReporter *ImportDataProgressReporter, poolSize int) error {
+	// The code can produce `poolSize` number of batches at a time. But, it can consume only
+	// `parallelism` number of batches at a time.
+	batchImportPool = pool.New().WithMaxGoroutines(poolSize)
+	log.Infof("created batch import pool of size: %d", poolSize)
+
+	taskPicker, err := NewSequentialTaskPicker(pendingTasks, state)
+	if err != nil {
+		return fmt.Errorf("create task picker: %w", err)
+	}
+	taskImporters := map[int]*FileTaskImporter{}
+
+	for taskPicker.HasMoreTasks() {
+		task, err := taskPicker.Pick()
+		if err != nil {
+			return fmt.Errorf("get next task: %w", err)
+		}
+		var taskImporter *FileTaskImporter
+		var ok bool
+		taskImporter, ok = taskImporters[task.ID]
+		if !ok {
+			taskImporter, err = NewFileTaskImporter(task, state, batchImportPool, progressReporter)
+			if err != nil {
+				return fmt.Errorf("create file task importer: %s", err)
+			}
+			log.Infof("created file task importer for table: %s, task: %v", task.TableNameTup.ForOutput(), task)
+			taskImporters[task.ID] = taskImporter
+		}
+
+		if taskImporter.AllBatchesSubmitted() {
+			// All batches for this task have been submitted.
+			// task could have been completed (all batches imported) OR still in progress
+			// in case task is done, we should inform task picker so that we stop picking that task.
+			taskDone, err := taskImporter.AllBatchesImported()
+			if err != nil {
+				return fmt.Errorf("check if all batches are imported: task: %v err :%w", task, err)
+			}
+			if taskDone {
+				taskImporter.PostProcess()
+				err = taskPicker.MarkTaskAsDone(task)
+				if err != nil {
+					return fmt.Errorf("mark task as done: task: %v, err: %w", task, err)
+				}
+				continue
+			} else {
+				// some batches are still in progress, wait for them to complete as decided by the picker.
+				// don't want to busy-wait, so in case of sequentialTaskPicker, we sleep.
+				taskPicker.WaitForTasksBatchesTobeImported()
+				continue
+			}
+
+		}
+		err = taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
+		if err != nil {
+			return fmt.Errorf("submit next batch: task:%v err: %s", task, err)
+		}
+	}
+	return nil
+}
+
 func startAdaptiveParallelism() (bool, error) {
 	yb, ok := tdb.(*tgtdb.TargetYugabyteDB)
 	if !ok {

yb-voyager/cmd/importDataFileBatchProducer_test.go

@@ -39,7 +39,7 @@ func TestBasicFileBatchProducer(t *testing.T) {
 
 	fileContents := `id,val
 1, "hello"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -71,7 +71,7 @@ func TestFileBatchProducerBasedOnRowsThreshold(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -122,7 +122,7 @@ func TestFileBatchProducerBasedOnSizeThreshold(t *testing.T) {
 2, "ghijk"
 3, "mnopq"
 4, "stuvw"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -190,7 +190,7 @@ func TestFileBatchProducerThrowsErrorWhenSingleRowGreaterThanMaxBatchSize(t *tes
 2, "ghijk"
 3, "mnopq1234567899876543"
 4, "stuvw"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -226,7 +226,7 @@ func TestFileBatchProducerResumable(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -282,7 +282,7 @@ func TestFileBatchProducerResumeAfterAllBatchesProduced(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)

yb-voyager/cmd/importDataFileCommand.go

@@ -179,7 +179,8 @@ func getImportFileTasks(currFileTableMapping string) []*ImportFileTask {
 		return result
 	}
 	kvs := strings.Split(currFileTableMapping, ",")
-	for i, kv := range kvs {
+	idCounter := 0
+	for _, kv := range kvs {
 		globPattern, table := strings.Split(kv, ":")[0], strings.Split(kv, ":")[1]
 		filePaths, err := dataStore.Glob(globPattern)
 		if err != nil {
@@ -198,12 +199,13 @@ func getImportFileTasks(currFileTableMapping string) []*ImportFileTask {
 				utils.ErrExit("calculating file size in bytes: %q: %v", filePath, err)
 			}
 			task := &ImportFileTask{
-				ID:           i,
+				ID:           idCounter,
 				FilePath:     filePath,
 				TableNameTup: tableNameTuple,
 				FileSize:     fileSize,
 			}
 			result = append(result, task)
+			idCounter++
 		}
 	}
 	return result

yb-voyager/cmd/importDataFileTaskImporter.go

@@ -30,6 +30,11 @@ import (
 	"github.com/yugabyte/yb-voyager/yb-voyager/src/utils/sqlname"
 )
 
+/*
+FileTaskImporter is responsible for importing an ImportFileTask.
+It uses a FileBatchProducer to produce batches. It submits each batch to a provided
+worker pool for processing. It also maintains and updates the progress of the task.
+*/
 type FileTaskImporter struct {
 	task                 *ImportFileTask
 	state                *ImportDataState
@@ -75,12 +80,15 @@ func (fti *FileTaskImporter) AllBatchesSubmitted() bool {
 	return fti.batchProducer.Done()
 }
 
-func (fti *FileTaskImporter) AllBatchesImported() error {
-	// TODO: check importDataState for status.
-	panic("not implemented")
+func (fti *FileTaskImporter) AllBatchesImported() (bool, error) {
+	taskStatus, err := fti.state.GetFileImportState(fti.task.FilePath, fti.task.TableNameTup)
+	if err != nil {
+		return false, fmt.Errorf("getting file import state: %s", err)
+	}
+	return taskStatus == FILE_IMPORT_COMPLETED, nil
 }
 
-func (fti *FileTaskImporter) SubmitNextBatch() error {
+func (fti *FileTaskImporter) ProduceAndSubmitNextBatchToWorkerPool() error {
 	if fti.AllBatchesSubmitted() {
 		return fmt.Errorf("no more batches to submit")
 	}
@@ -234,6 +242,7 @@ func getImportBatchArgsProto(tableNameTup sqlname.NameTuple, filePath string) *t
 	}
 	// If `columns` is unset at this point, no attribute list is passed in the COPY command.
 	fileFormat := dataFileDescriptor.FileFormat
+
 	// from export data with ora2pg, it comes as an SQL file, with COPY command having data.
 	// Import-data also reads it appropriately with the help of sqlDataFile.
 	// But while running COPY for a batch, we need to set the format as TEXT (SQL does not make sense)

yb-voyager/cmd/importDataFileTaskImporter_test.go

@@ -47,7 +47,7 @@ func TestBasicTaskImport(t *testing.T) {
 	fileContents := `id,val
 1, "hello"
 2, "world"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_basic")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_basic", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)
@@ -56,7 +56,7 @@ func TestBasicTaskImport(t *testing.T) {
 	testutils.FatalIfError(t, err)
 
 	for !taskImporter.AllBatchesSubmitted() {
-		err := taskImporter.SubmitNextBatch()
+		err := taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
 		assert.NoError(t, err)
 	}
 
@@ -88,15 +88,15 @@ func TestImportAllBatchesAndResume(t *testing.T) {
 	fileContents := `id,val
 1, "hello"
 2, "world"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_all")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_all", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)
 	workerPool := pool.New().WithMaxGoroutines(2)
 	taskImporter, err := NewFileTaskImporter(task, state, workerPool, progressReporter)
 
 	for !taskImporter.AllBatchesSubmitted() {
-		err := taskImporter.SubmitNextBatch()
+		err := taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
 		assert.NoError(t, err)
 	}
 
@@ -139,7 +139,7 @@ func TestTaskImportResumable(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)
@@ -148,7 +148,7 @@ func TestTaskImportResumable(t *testing.T) {
 	testutils.FatalIfError(t, err)
 
 	// submit 1 batch
-	err = taskImporter.SubmitNextBatch()
+	err = taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
 	assert.NoError(t, err)
 
 	// check that the first batch was imported
@@ -165,7 +165,7 @@ func TestTaskImportResumable(t *testing.T) {
 	testutils.FatalIfError(t, err)
 
 	// submit second batch, not first batch again as it was already imported
-	err = taskImporter.SubmitNextBatch()
+	err = taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
 	assert.NoError(t, err)
 
 	assert.Equal(t, true, taskImporter.AllBatchesSubmitted())
@@ -198,7 +198,7 @@ func TestTaskImportResumableNoPK(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume_no_pk")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume_no_pk", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)
@@ -207,7 +207,7 @@ func TestTaskImportResumableNoPK(t *testing.T) {
 	testutils.FatalIfError(t, err)
 
 	// submit 1 batch
-	err = taskImporter.SubmitNextBatch()
+	err = taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
 	assert.NoError(t, err)
 
 	// check that the first batch was imported
@@ -224,7 +224,7 @@ func TestTaskImportResumableNoPK(t *testing.T) {
 	testutils.FatalIfError(t, err)
 
 	// submit second batch, not first batch again as it was already imported
-	err = taskImporter.SubmitNextBatch()
+	err = taskImporter.ProduceAndSubmitNextBatchToWorkerPool()
 	assert.NoError(t, err)
 
 	assert.Equal(t, true, taskImporter.AllBatchesSubmitted())