backup/data.go

package backup

/*
 * This file contains structs and functions related to backing up data on the segments.
 */

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

	"github.com/cloudberrydb/gp-common-go-libs/dbconn"
	"github.com/cloudberrydb/gp-common-go-libs/gplog"
	"github.com/cloudberrydb/gpbackup/options"
	"github.com/cloudberrydb/gpbackup/utils"
	"github.com/jackc/pgconn"
	"gopkg.in/cheggaaa/pb.v1"
)

var (
	tableDelim = ","
)

func ConstructTableAttributesList(columnDefs []ColumnDefinition) string {
	// this attribute list used ONLY by CopyTableIn on the restore side
	// columns where data should not be copied out and back in are excluded from this string.
	names := make([]string, 0)
	for _, col := range columnDefs {
		// data in generated columns should not be backed up or restored.
		if col.AttGenerated == "" {
			names = append(names, col.Name)
		}
	}
	if len(names) > 0 {
		return fmt.Sprintf("(%s)", strings.Join(names, ","))
	}
	return ""
}

func AddTableDataEntriesToTOC(tables []Table, rowsCopiedMaps []map[uint32]int64) {
	for _, table := range tables {
		if !table.SkipDataBackup() {
			var rowsCopied int64
			for _, rowsCopiedMap := range rowsCopiedMaps {
				if val, ok := rowsCopiedMap[table.Oid]; ok {
					rowsCopied = val
					break
				}
			}
			attributes := ConstructTableAttributesList(table.ColumnDefs)
			globalTOC.AddCoordinatorDataEntry(table.Schema, table.Name, table.Oid, attributes, rowsCopied, table.PartitionLevelInfo.RootName, table.DistPolicy)
		}
	}
}

type BackupProgressCounters struct {
	NumRegTables   int64
	TotalRegTables int64
	ProgressBar    utils.ProgressBar
}

func CopyTableOut(connectionPool *dbconn.DBConn, table Table, destinationToWrite string, connNum int) (int64, error) {
	checkPipeExistsCommand := ""
	customPipeThroughCommand := utils.GetPipeThroughProgram().OutputCommand
	sendToDestinationCommand := ">"
	if MustGetFlagBool(options.SINGLE_DATA_FILE) {
		/*
		 * The segment TOC files are always written to the segment data directory for
		 * performance reasons, in case the user-specified directory is on a mounted
		 * drive.  It will be copied to a user-specified directory, if any, once all
		 * of the data is backed up.
		 */
		checkPipeExistsCommand = fmt.Sprintf("(test -p \"%s\" || (echo \"Pipe not found %s\">&2; exit 1)) && ", destinationToWrite, destinationToWrite)
		customPipeThroughCommand = "cat -"
	} else if MustGetFlagString(options.PLUGIN_CONFIG) != "" {
		sendToDestinationCommand = fmt.Sprintf("| %s backup_data %s", pluginConfig.ExecutablePath, pluginConfig.ConfigPath)
	}

	copyCommand := fmt.Sprintf("PROGRAM '%s%s %s %s'", checkPipeExistsCommand, customPipeThroughCommand, sendToDestinationCommand, destinationToWrite)

	columnNames := ""
		// process column names to exclude generated columns from data copy out
	columnNames = ConstructTableAttributesList(table.ColumnDefs)

	query := fmt.Sprintf("COPY %s%s TO %s WITH CSV DELIMITER '%s' ON SEGMENT IGNORE EXTERNAL PARTITIONS;", table.FQN(), columnNames, copyCommand, tableDelim)
	gplog.Verbose("Worker %d: %s", connNum, query)
	result, err := connectionPool.Exec(query, connNum)
	if err != nil {
		return 0, err
	}
	numRows, _ := result.RowsAffected()
	return numRows, nil
}

func BackupSingleTableData(table Table, rowsCopiedMap map[uint32]int64, counters *BackupProgressCounters, whichConn int) error {
	logMessage := fmt.Sprintf("Worker %d: Writing data for table %s to file", whichConn, table.FQN())
	// Avoid race condition by incrementing counters in call to sprintf
	tableCount := fmt.Sprintf(" (table %d of %d)", atomic.AddInt64(&counters.NumRegTables, 1), counters.TotalRegTables)
	if gplog.GetVerbosity() > gplog.LOGINFO {
		// No progress bar at this log level, so we note table count here
		gplog.Verbose(logMessage + tableCount)
	} else {
		gplog.Verbose(logMessage)
	}

	destinationToWrite := ""
	if MustGetFlagBool(options.SINGLE_DATA_FILE) {
		destinationToWrite = fmt.Sprintf("%s_%d", globalFPInfo.GetSegmentPipePathForCopyCommand(), table.Oid)
	} else {
		destinationToWrite = globalFPInfo.GetTableBackupFilePathForCopyCommand(table.Oid, utils.GetPipeThroughProgram().Extension, false)
	}
	rowsCopied, err := CopyTableOut(connectionPool, table, destinationToWrite, whichConn)
	if err != nil {
		return err
	}
	rowsCopiedMap[table.Oid] = rowsCopied
	counters.ProgressBar.Increment()
	return nil
}

/*
* Iterate through tables, backup data from each table in set.
* If supported by the database, a synchronized database snapshot is used to sync
* all parallel workers' view of the database and ensure data consistency across workers.
* In each worker, we try to acquire a ACCESS SHARE lock on each table in NOWAIT mode,
* to avoid potential deadlocks with queued DDL operations that requested ACCESS EXCLUSIVE
* locks on the same tables (for eg concurrent ALTER TABLE operations). If worker is unable to
* acquire a lock, defer table to worker 0, which holds all ACCESS SHARE locks for the backup set.
* If synchronized snapshot is not supported and worker is unable to acquire a lock, the
* worker must be terminated because the session no longer has a valid distributed snapshot
*
* FIXME: Simplify backupDataForAllTables by having one function for snapshot workflow and
* another without, then extract common portions into their own functions.
 */
func backupDataForAllTables(tables []Table) []map[uint32]int64 {
	counters := BackupProgressCounters{NumRegTables: 0, TotalRegTables: int64(len(tables))}
	counters.ProgressBar = utils.NewProgressBar(int(counters.TotalRegTables), "Tables backed up: ", utils.PB_INFO)
	counters.ProgressBar.Start()
	rowsCopiedMaps := make([]map[uint32]int64, connectionPool.NumConns)
	/*
	 * We break when an interrupt is received and rely on
	 * TerminateHangingCopySessions to halt any COPY statements
	 * in progress if they don't finish on their own.
	 */
	tasks := make(chan Table, len(tables))
	var oidMap sync.Map
	var workerPool sync.WaitGroup
	var copyErr error
	// Record and track tables in a hashmap of oids and table states (preloaded with value Unknown).
	// The tables are loaded into the tasks channel for the subsequent goroutines to work on.
	for _, table := range tables {
		oidMap.Store(table.Oid, Unknown)
		tasks <- table
	}

	/*
	 * Worker 0 is a special database connection that
	 * 	1) Exports the database snapshot if the feature is supported
	 * 	2) Does not have tables pre-assigned to it.
	 * 	3) Processes tables only in the event that the other workers encounter locking issues.
	 * Worker 0 already has all locks on the tables so it will not run into locking issues.
	 */
	rowsCopiedMaps[0] = make(map[uint32]int64)
	for connNum := 1; connNum < connectionPool.NumConns; connNum++ {
		rowsCopiedMaps[connNum] = make(map[uint32]int64)
		workerPool.Add(1)
		go func(whichConn int) {
			defer workerPool.Done()
			/* If the --leaf-partition-data flag is not set, the parent and all leaf
			 * partition data are treated as a single table and will be assigned to a single worker.
			 * Large partition hierarchies result in a large number of locks being held until the
			 * transaction commits and the locks are released.
			 */
			for table := range tasks {
				if wasTerminated || copyErr != nil {
					counters.ProgressBar.(*pb.ProgressBar).NotPrint = true
					return
				}
				if backupSnapshot != "" && connectionPool.Tx[whichConn] == nil {
					err := SetSynchronizedSnapshot(connectionPool, whichConn, backupSnapshot)
					if err != nil {
						gplog.FatalOnError(err)
					}
				}
				// If a random external SQL command had queued an AccessExclusiveLock acquisition request
				// against this next table, the --job worker thread would deadlock on the COPY attempt.
				// To prevent gpbackup from hanging, we attempt to acquire an AccessShareLock on the
				// relation with the NOWAIT option before we run COPY. If the LOCK TABLE NOWAIT call
				// fails, we catch the error and defer the table to the main worker thread, worker 0.
				// Afterwards, we break early and terminate the worker since its transaction is now in an
				// aborted state. We do not need to do this with the main worker thread because it has
				// already acquired AccessShareLocks on all tables before the metadata dumping part.
				err := LockTableNoWait(table, whichConn)
				if err != nil {
					if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code != PG_LOCK_NOT_AVAILABLE {
						copyErr = err
						oidMap.Store(table.Oid, Deferred)
						err = connectionPool.Rollback(whichConn)
						if err != nil {
							gplog.Warn("Worker %d: %s", whichConn, err)
						}
						continue
					}
					if gplog.GetVerbosity() < gplog.LOGVERBOSE {
						// Add a newline to interrupt the progress bar so that
						// the following WARN message is nicely outputted.
						fmt.Printf("\n")
					}
					gplog.Warn("Worker %d could not acquire AccessShareLock for table %s.", whichConn, table.FQN())
					logTableLocks(table, whichConn)
					// rollback transaction and defer table
					err = connectionPool.Rollback(whichConn)
					if err != nil {
						gplog.Warn("Worker %d: %s", whichConn, err)
					}
					oidMap.Store(table.Oid, Deferred)
					// If have backup snapshot, continue to next table, else terminate worker.
					if backupSnapshot != "" {
						continue
					} else {
						gplog.Warn("Terminating worker %d and deferring table %s to main worker thread.", whichConn, table.FQN())
						break
					}
				}
				err = BackupSingleTableData(table, rowsCopiedMaps[whichConn], &counters, whichConn)
				if err != nil {
					copyErr = err
					break
				} else {
					oidMap.Store(table.Oid, Complete)
				}
				if backupSnapshot != "" {
					err = connectionPool.Commit(whichConn)
					if err != nil {
						gplog.Warn("Worker %d: %s", whichConn, err)
					}
				}
			}
		}(connNum)
	}
	// Special goroutine to handle deferred tables
	// Handle all tables deferred by the deadlock detection. This can only be
	// done with the main worker thread, worker 0, because it has
	// AccessShareLocks on all the tables already.
	deferredWorkerDone := make(chan bool)
	go func() {
		for _, table := range tables {
			for {
				state, _ := oidMap.Load(table.Oid)
				if state.(int) == Unknown {
					time.Sleep(time.Millisecond * 50)
				} else if state.(int) == Deferred {
					err := BackupSingleTableData(table, rowsCopiedMaps[0], &counters, 0)
					if err != nil {
						copyErr = err
					}
					oidMap.Store(table.Oid, Complete)
					break
				} else if state.(int) == Complete {
					break
				} else {
					gplog.Fatal(errors.New("Encountered unknown table state"), "")
				}
			}
		}
		deferredWorkerDone <- true
	}()
	close(tasks)
	workerPool.Wait()
	// If not using synchronized snapshots,
	// check if all workers were terminated due to lock issues.
	if backupSnapshot == "" {
		allWorkersTerminatedLogged := false
		for _, table := range tables {
			if wasTerminated || copyErr != nil {
				counters.ProgressBar.(*pb.ProgressBar).NotPrint = true
				break
			}
			state, _ := oidMap.Load(table.Oid)
			if state == Unknown {
				if !allWorkersTerminatedLogged {
					gplog.Warn("All workers terminated due to lock issues. Falling back to single main worker.")
					allWorkersTerminatedLogged = true
				}
				oidMap.Store(table.Oid, Deferred)
			}
		}
	}

	// Main goroutine waits for deferred worker 0 by waiting on this channel
	<-deferredWorkerDone
	agentErr := utils.CheckAgentErrorsOnSegments(globalCluster, globalFPInfo)
	if copyErr != nil && agentErr != nil {
		gplog.Error(agentErr.Error())
		gplog.Fatal(copyErr, "")
	} else if copyErr != nil {
		gplog.Fatal(copyErr, "")
	} else if agentErr != nil {
		gplog.Fatal(agentErr, "")
	}

	counters.ProgressBar.Finish()
	return rowsCopiedMaps
}

func printDataBackupWarnings(numExtTables int64) {
	if numExtTables > 0 {
		gplog.Info("Skipped data backup of %d external/foreign table(s).", numExtTables)
		gplog.Info("See %s for a complete list of skipped tables.", gplog.GetLogFilePath())
	}
}

// Remove external/foreign tables from the data backup set
func GetBackupDataSet(tables []Table) ([]Table, int64) {
	var backupDataSet []Table
	var numExtOrForeignTables int64

	if !backupReport.MetadataOnly {
		for _, table := range tables {
			if !table.SkipDataBackup() {
				backupDataSet = append(backupDataSet, table)
			} else {
				gplog.Verbose("Skipping data backup of table %s because it is either an external or foreign table.", table.FQN())
				numExtOrForeignTables++
			}
		}
	}
	return backupDataSet, numExtOrForeignTables
}

// Acquire AccessShareLock on a table with NOWAIT option. If we are unable to acquire
// the lock, the call will fail instead of block. Return the failure for handling.
func LockTableNoWait(dataTable Table, connNum int) error {
	var lockMode string
	lockMode = `IN ACCESS SHARE MODE`
	query := fmt.Sprintf("LOCK TABLE %s %s;", dataTable.FQN(), lockMode)
	gplog.Verbose("Worker %d: %s", connNum, query)
	_, err := connectionPool.Exec(query, connNum)
	if err != nil {
		return err
	}
	return nil
}