restore/parallel.go - cloudberry-backup - Git at Google

 package restore

 /*
  * This file contains functions related to executing multiple SQL statements in parallel.
  */

 import (
 	"fmt"
 	"strings"
 	"sync"
 	"sync/atomic"

 	"github.com/apache/cloudberry-backup/options"
 	"github.com/apache/cloudberry-backup/toc"
 	"github.com/apache/cloudberry-backup/utils"
 	"github.com/apache/cloudberry-go-libs/gplog"
 )

 var (
 	mutex   = &sync.Mutex{}
 	txMutex = &sync.Mutex{}
 )

 func executeStatementsForConn(statements chan toc.StatementWithType, fatalErr *error, numErrors *int32, progressBar utils.ProgressBar, whichConn int, executeInParallel bool) {
 	for statement := range statements {
 		if wasTerminated || *fatalErr != nil {
 			if executeInParallel {
 				gplog.Error("Error detected on connection %d. Terminating transactions.", whichConn)
 				txMutex.Lock()
 				if connectionPool.Tx[whichConn] != nil {
 					connectionPool.Rollback(whichConn)
 				}
 				txMutex.Unlock()
 			}
 			return
 		}

 		gplog.Debug("Executing statement: %s on connection: %d", strings.TrimSpace(statement.Statement), whichConn)
 		_, err := connectionPool.Exec(statement.Statement, whichConn)
 		if err != nil {
 			gplog.Verbose("Error encountered when executing statement: %s Error was: %s", strings.TrimSpace(statement.Statement), err.Error())
 			if MustGetFlagBool(options.ON_ERROR_CONTINUE) {
 				if executeInParallel {
 					atomic.AddInt32(numErrors, 1)
 					if statement.ObjectType == toc.OBJ_TABLE {
 						mutex.Lock()
 						errorTablesMetadata[statement.Schema+"."+statement.Name] = Empty{}
 						mutex.Unlock()
 					}
 				} else {
 					*numErrors = *numErrors + 1
 					if statement.ObjectType == toc.OBJ_TABLE {
 						errorTablesMetadata[statement.Schema+"."+statement.Name] = Empty{}
 					}
 				}
 			} else {
 				*fatalErr = err
 			}
 		}
 		if progressBar != nil {
 			progressBar.Increment()
 		}
 	}

 	if executeInParallel {
 		txMutex.Lock()
 		if connectionPool.Tx[whichConn] != nil {
 			connectionPool.Commit(whichConn)
 		}
 		txMutex.Unlock()
 	}
 }

 func findLeastBusyConn(connStatementCounts map[int]int) int {
 	minStatements := connStatementCounts[1]
 	leastBusyConn := 1
 	for conn, count := range connStatementCounts {
 		if count < minStatements {
 			leastBusyConn = conn
 			minStatements = count
 		}
 	}
 	return leastBusyConn
 }

 func scheduleStatementsOnWorkers(statements []toc.StatementWithType, numConns int) map[int][]toc.StatementWithType {
 	splitStatements := make(map[int][]toc.StatementWithType, numConns)
 	cohortConnAssignments := make(map[uint32]int)

 	connStatementCounts := make(map[int]int, numConns)
 	for i := 0; i < numConns; i++ {
 		connStatementCounts[i] = 0
 	}
 	// During backup, we track objects into "cohorts" that must be backed up in the same
 	// transaction to avoid lock contention and/or deadlocks. Split up statements to
 	// connections by cohort.
 	for i := 0; i < len(statements); i++ {
 		var cohort uint32
 		if len(statements[i].Tier) == 0 {
 			// older backups will not have a Tier value
 			cohort = 0
 		} else {
 			cohort = statements[i].Tier[1]
 		}
 		keyConn, alreadyPinned := cohortConnAssignments[cohort]
 		if alreadyPinned {
 			splitStatements[keyConn] = append(splitStatements[keyConn], statements[i])
 			connStatementCounts[keyConn] = connStatementCounts[keyConn] + 1
 		} else {
 			leastBusyConn := findLeastBusyConn(connStatementCounts)
 			cohortConnAssignments[cohort] = leastBusyConn
 			splitStatements[leastBusyConn] = append(splitStatements[leastBusyConn], statements[i])
 			connStatementCounts[leastBusyConn] = connStatementCounts[leastBusyConn] + 1
 		}
 	}
 	return splitStatements
 }

 /*
  * This function creates a worker pool of N goroutines to be able to execute up
  * to N statements in parallel.
  */
 func ExecutePredataStatements(statements []toc.StatementWithType, progressBar utils.ProgressBar, executeInParallel bool, whichConn ...int) int32 {
 	var workerPool sync.WaitGroup
 	var fatalErr error
 	var numErrors int32

 	if len(statements) == 0 {
 		return 0
 	}
 	if !executeInParallel {
 		tasks := make(chan toc.StatementWithType, len(statements))
 		for _, statement := range statements {
 			tasks <- statement
 		}
 		close(tasks)
 		connNum := connectionPool.ValidateConnNum(whichConn...)
 		executeStatementsForConn(tasks, &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
 	} else {
 		panicChan := make(chan error)
 		splitStatements := scheduleStatementsOnWorkers(statements, connectionPool.NumConns)
 		chanMap := make(map[int]chan toc.StatementWithType, connectionPool.NumConns)
 		// preinitialize channels to prevent concurrent read and write
 		for i := 0; i < connectionPool.NumConns; i++ {
 			chanMap[i] = make(chan toc.StatementWithType, len(splitStatements[i]))
 			for _, statement := range splitStatements[i] {
 				chanMap[i] <- statement
 			}
 			close(chanMap[i])
 		}
 		for i := 0; i < connectionPool.NumConns; i++ {
 			workerPool.Add(1)
 			go func(connNum int) {
 				defer func() {
 					if panicErr := recover(); panicErr != nil {
 						panicChan <- fmt.Errorf("%v", panicErr)
 					}
 				}()
 				defer workerPool.Done()
 				connNum = connectionPool.ValidateConnNum(connNum)
 				executeStatementsForConn(chanMap[connNum], &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
 			}(i)
 		}
 		workerPool.Wait()
 		// Allow panics to crash from the main process, invoking DoCleanup
 		select {
 		case err := <-panicChan:
 			gplog.Fatal(err, "")
 		default:
 			// no panic, nothing to do
 		}
 	}
 	if fatalErr != nil {
 		fmt.Println("")
 		gplog.Fatal(fatalErr, "")
 	} else if numErrors > 0 {
 		fmt.Println("")
 		gplog.Error("Encountered %d errors during metadata restore; see log file %s for a list of failed statements.", numErrors, gplog.GetLogFilePath())
 	}

 	return numErrors
 }

 func ExecuteStatements(statements []toc.StatementWithType, progressBar utils.ProgressBar, executeInParallel bool, whichConn ...int) int32 {
 	var workerPool sync.WaitGroup
 	var fatalErr error
 	var numErrors int32
 	tasks := make(chan toc.StatementWithType, len(statements))
 	for _, statement := range statements {
 		tasks <- statement
 	}
 	close(tasks)

 	if !executeInParallel {
 		connNum := connectionPool.ValidateConnNum(whichConn...)
 		executeStatementsForConn(tasks, &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
 	} else {
 		panicChan := make(chan error)
 		for i := 0; i < connectionPool.NumConns; i++ {
 			workerPool.Add(1)
 			go func(connNum int) {
 				defer func() {
 					if panicErr := recover(); panicErr != nil {
 						panicChan <- fmt.Errorf("%v", panicErr)
 					}
 				}()
 				defer workerPool.Done()
 				connNum = connectionPool.ValidateConnNum(connNum)
 				executeStatementsForConn(tasks, &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
 			}(i)
 		}
 		workerPool.Wait()
 		// Allow panics to crash from the main process, invoking DoCleanup
 		select {
 		case err := <-panicChan:
 			gplog.Fatal(err, "")
 		default:
 			// no panic, nothing to do
 		}
 	}
 	if fatalErr != nil {
 		fmt.Println("")
 		gplog.Fatal(fatalErr, "")
 	} else if numErrors > 0 {
 		fmt.Println("")
 		gplog.Error("Encountered %d errors during metadata restore; see log file %s for a list of failed statements.", numErrors, gplog.GetLogFilePath())
 	}

 	return numErrors
 }

 /*
  *   There is an existing bug in Greenplum where creating indexes in parallel
  *   on an AO table that didn't have any indexes previously can cause
  *   deadlock.
  *
  *   We work around this issue by restoring post data objects in
  *   two batches. The first batch takes one index from each table and
  *   restores them in parallel (which has no possibility of deadlock) and
  *   then the second restores all other postdata objects in parallel. After
  *   each table has at least one index, there is no more risk of deadlock.
  *
  *   A third batch is created specifically for postdata metadata
  *   (e.g. ALTER INDEX, ALTER EVENT TRIGGER, COMMENT ON). These
  *   statements cannot be concurrently run with batch two since that
  *   is where the dependent postdata objects are being created.
  */
 func BatchPostdataStatements(statements []toc.StatementWithType) ([]toc.StatementWithType, []toc.StatementWithType, []toc.StatementWithType) {
 	indexMap := make(map[string]bool)
 	firstBatch := make([]toc.StatementWithType, 0)
 	secondBatch := make([]toc.StatementWithType, 0)
 	thirdBatch := make([]toc.StatementWithType, 0)
 	for _, statement := range statements {
 		_, tableIndexPresent := indexMap[statement.ReferenceObject]
 		if statement.ObjectType == toc.OBJ_INDEX && !tableIndexPresent {
 			indexMap[statement.ReferenceObject] = true
 			firstBatch = append(firstBatch, statement)
 		} else if strings.Contains(statement.ObjectType, " METADATA") {
 			thirdBatch = append(thirdBatch, statement)
 		} else {
 			secondBatch = append(secondBatch, statement)
 		}
 	}
 	return firstBatch, secondBatch, thirdBatch
 }

 func BatchPredataStatements(statements []toc.StatementWithType) ([]toc.StatementWithType, map[uint32][]toc.StatementWithType, []toc.StatementWithType) {
 	foundNumberedTier := false
 	firstTierZero := make([]toc.StatementWithType, 0)
 	numberedTiers := make(map[uint32][]toc.StatementWithType)
 	secondTierZero := make([]toc.StatementWithType, 0)

 	for _, statement := range statements {
 		// For backwards compatibility, assume restores without Tier are always 0,0.  This is also
 		// used for some object types like Schemas where we haven't yet gone through the work of
 		// tiering them out for parallel execution.
 		if len(statement.Tier) == 0 {
 			statement.Tier = []uint32{0, 0}
 		}

 		if statement.Tier[0] > 0 {
 			foundNumberedTier = true
 			numberedTiers[statement.Tier[0]] = append(numberedTiers[statement.Tier[0]], statement)
 		} else if statement.Tier[0] == 0 {
 			if !foundNumberedTier {
 				// insert into firstTierZero
 				firstTierZero = append(firstTierZero, statement)
 			} else {
 				// insert into secondTierZero
 				secondTierZero = append(secondTierZero, statement)
 			}
 		}
 	}
 	return firstTierZero, numberedTiers, secondTierZero
 }
	package restore

	/*
	* This file contains functions related to executing multiple SQL statements in parallel.
	*/

	import (
	"fmt"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/apache/cloudberry-backup/options"
	"github.com/apache/cloudberry-backup/toc"
	"github.com/apache/cloudberry-backup/utils"
	"github.com/apache/cloudberry-go-libs/gplog"
	)

	var (
	mutex = &sync.Mutex{}
	txMutex = &sync.Mutex{}
	)

	func executeStatementsForConn(statements chan toc.StatementWithType, fatalErr error, numErrors int32, progressBar utils.ProgressBar, whichConn int, executeInParallel bool) {
	for statement := range statements {
	if wasTerminated \|\| *fatalErr != nil {
	if executeInParallel {
	gplog.Error("Error detected on connection %d. Terminating transactions.", whichConn)
	txMutex.Lock()
	if connectionPool.Tx[whichConn] != nil {
	connectionPool.Rollback(whichConn)
	}
	txMutex.Unlock()
	}
	return
	}

	gplog.Debug("Executing statement: %s on connection: %d", strings.TrimSpace(statement.Statement), whichConn)
	_, err := connectionPool.Exec(statement.Statement, whichConn)
	if err != nil {
	gplog.Verbose("Error encountered when executing statement: %s Error was: %s", strings.TrimSpace(statement.Statement), err.Error())
	if MustGetFlagBool(options.ON_ERROR_CONTINUE) {
	if executeInParallel {
	atomic.AddInt32(numErrors, 1)
	if statement.ObjectType == toc.OBJ_TABLE {
	mutex.Lock()
	errorTablesMetadata[statement.Schema+"."+statement.Name] = Empty{}
	mutex.Unlock()
	}
	} else {
	numErrors = numErrors + 1
	if statement.ObjectType == toc.OBJ_TABLE {
	errorTablesMetadata[statement.Schema+"."+statement.Name] = Empty{}
	}
	}
	} else {
	*fatalErr = err
	}
	}
	if progressBar != nil {
	progressBar.Increment()
	}
	}

	if executeInParallel {
	txMutex.Lock()
	if connectionPool.Tx[whichConn] != nil {
	connectionPool.Commit(whichConn)
	}
	txMutex.Unlock()
	}
	}

	func findLeastBusyConn(connStatementCounts map[int]int) int {
	minStatements := connStatementCounts[1]
	leastBusyConn := 1
	for conn, count := range connStatementCounts {
	if count < minStatements {
	leastBusyConn = conn
	minStatements = count
	}
	}
	return leastBusyConn
	}

	func scheduleStatementsOnWorkers(statements []toc.StatementWithType, numConns int) map[int][]toc.StatementWithType {
	splitStatements := make(map[int][]toc.StatementWithType, numConns)
	cohortConnAssignments := make(map[uint32]int)

	connStatementCounts := make(map[int]int, numConns)
	for i := 0; i < numConns; i++ {
	connStatementCounts[i] = 0
	}
	// During backup, we track objects into "cohorts" that must be backed up in the same
	// transaction to avoid lock contention and/or deadlocks. Split up statements to
	// connections by cohort.
	for i := 0; i < len(statements); i++ {
	var cohort uint32
	if len(statements[i].Tier) == 0 {
	// older backups will not have a Tier value
	cohort = 0
	} else {
	cohort = statements[i].Tier[1]
	}
	keyConn, alreadyPinned := cohortConnAssignments[cohort]
	if alreadyPinned {
	splitStatements[keyConn] = append(splitStatements[keyConn], statements[i])
	connStatementCounts[keyConn] = connStatementCounts[keyConn] + 1
	} else {
	leastBusyConn := findLeastBusyConn(connStatementCounts)
	cohortConnAssignments[cohort] = leastBusyConn
	splitStatements[leastBusyConn] = append(splitStatements[leastBusyConn], statements[i])
	connStatementCounts[leastBusyConn] = connStatementCounts[leastBusyConn] + 1
	}
	}
	return splitStatements
	}

	/*
	* This function creates a worker pool of N goroutines to be able to execute up
	* to N statements in parallel.
	*/
	func ExecutePredataStatements(statements []toc.StatementWithType, progressBar utils.ProgressBar, executeInParallel bool, whichConn ...int) int32 {
	var workerPool sync.WaitGroup
	var fatalErr error
	var numErrors int32

	if len(statements) == 0 {
	return 0
	}
	if !executeInParallel {
	tasks := make(chan toc.StatementWithType, len(statements))
	for _, statement := range statements {
	tasks <- statement
	}
	close(tasks)
	connNum := connectionPool.ValidateConnNum(whichConn...)
	executeStatementsForConn(tasks, &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
	} else {
	panicChan := make(chan error)
	splitStatements := scheduleStatementsOnWorkers(statements, connectionPool.NumConns)
	chanMap := make(map[int]chan toc.StatementWithType, connectionPool.NumConns)
	// preinitialize channels to prevent concurrent read and write
	for i := 0; i < connectionPool.NumConns; i++ {
	chanMap[i] = make(chan toc.StatementWithType, len(splitStatements[i]))
	for _, statement := range splitStatements[i] {
	chanMap[i] <- statement
	}
	close(chanMap[i])
	}
	for i := 0; i < connectionPool.NumConns; i++ {
	workerPool.Add(1)
	go func(connNum int) {
	defer func() {
	if panicErr := recover(); panicErr != nil {
	panicChan <- fmt.Errorf("%v", panicErr)
	}
	}()
	defer workerPool.Done()
	connNum = connectionPool.ValidateConnNum(connNum)
	executeStatementsForConn(chanMap[connNum], &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
	}(i)
	}
	workerPool.Wait()
	// Allow panics to crash from the main process, invoking DoCleanup
	select {
	case err := <-panicChan:
	gplog.Fatal(err, "")
	default:
	// no panic, nothing to do
	}
	}
	if fatalErr != nil {
	fmt.Println("")
	gplog.Fatal(fatalErr, "")
	} else if numErrors > 0 {
	fmt.Println("")
	gplog.Error("Encountered %d errors during metadata restore; see log file %s for a list of failed statements.", numErrors, gplog.GetLogFilePath())
	}

	return numErrors
	}

	func ExecuteStatements(statements []toc.StatementWithType, progressBar utils.ProgressBar, executeInParallel bool, whichConn ...int) int32 {
	var workerPool sync.WaitGroup
	var fatalErr error
	var numErrors int32
	tasks := make(chan toc.StatementWithType, len(statements))
	for _, statement := range statements {
	tasks <- statement
	}
	close(tasks)

	if !executeInParallel {
	connNum := connectionPool.ValidateConnNum(whichConn...)
	executeStatementsForConn(tasks, &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
	} else {
	panicChan := make(chan error)
	for i := 0; i < connectionPool.NumConns; i++ {
	workerPool.Add(1)
	go func(connNum int) {
	defer func() {
	if panicErr := recover(); panicErr != nil {
	panicChan <- fmt.Errorf("%v", panicErr)
	}
	}()
	defer workerPool.Done()
	connNum = connectionPool.ValidateConnNum(connNum)
	executeStatementsForConn(tasks, &fatalErr, &numErrors, progressBar, connNum, executeInParallel)
	}(i)
	}
	workerPool.Wait()
	// Allow panics to crash from the main process, invoking DoCleanup
	select {
	case err := <-panicChan:
	gplog.Fatal(err, "")
	default:
	// no panic, nothing to do
	}
	}
	if fatalErr != nil {
	fmt.Println("")
	gplog.Fatal(fatalErr, "")
	} else if numErrors > 0 {
	fmt.Println("")
	gplog.Error("Encountered %d errors during metadata restore; see log file %s for a list of failed statements.", numErrors, gplog.GetLogFilePath())
	}

	return numErrors
	}

	/*
	* There is an existing bug in Greenplum where creating indexes in parallel
	* on an AO table that didn't have any indexes previously can cause
	* deadlock.
	*
	* We work around this issue by restoring post data objects in
	* two batches. The first batch takes one index from each table and
	* restores them in parallel (which has no possibility of deadlock) and
	* then the second restores all other postdata objects in parallel. After
	* each table has at least one index, there is no more risk of deadlock.
	*
	* A third batch is created specifically for postdata metadata
	* (e.g. ALTER INDEX, ALTER EVENT TRIGGER, COMMENT ON). These
	* statements cannot be concurrently run with batch two since that
	* is where the dependent postdata objects are being created.
	*/
	func BatchPostdataStatements(statements []toc.StatementWithType) ([]toc.StatementWithType, []toc.StatementWithType, []toc.StatementWithType) {
	indexMap := make(map[string]bool)
	firstBatch := make([]toc.StatementWithType, 0)
	secondBatch := make([]toc.StatementWithType, 0)
	thirdBatch := make([]toc.StatementWithType, 0)
	for _, statement := range statements {
	_, tableIndexPresent := indexMap[statement.ReferenceObject]
	if statement.ObjectType == toc.OBJ_INDEX && !tableIndexPresent {
	indexMap[statement.ReferenceObject] = true
	firstBatch = append(firstBatch, statement)
	} else if strings.Contains(statement.ObjectType, " METADATA") {
	thirdBatch = append(thirdBatch, statement)
	} else {
	secondBatch = append(secondBatch, statement)
	}
	}
	return firstBatch, secondBatch, thirdBatch
	}

	func BatchPredataStatements(statements []toc.StatementWithType) ([]toc.StatementWithType, map[uint32][]toc.StatementWithType, []toc.StatementWithType) {
	foundNumberedTier := false
	firstTierZero := make([]toc.StatementWithType, 0)
	numberedTiers := make(map[uint32][]toc.StatementWithType)
	secondTierZero := make([]toc.StatementWithType, 0)

	for _, statement := range statements {
	// For backwards compatibility, assume restores without Tier are always 0,0. This is also
	// used for some object types like Schemas where we haven't yet gone through the work of
	// tiering them out for parallel execution.
	if len(statement.Tier) == 0 {
	statement.Tier = []uint32{0, 0}
	}

	if statement.Tier[0] > 0 {
	foundNumberedTier = true
	numberedTiers[statement.Tier[0]] = append(numberedTiers[statement.Tier[0]], statement)
	} else if statement.Tier[0] == 0 {
	if !foundNumberedTier {
	// insert into firstTierZero
	firstTierZero = append(firstTierZero, statement)
	} else {
	// insert into secondTierZero
	secondTierZero = append(secondTierZero, statement)
	}
	}
	}
	return firstTierZero, numberedTiers, secondTierZero
	}