Merge remote-tracking branch 'origin/dev'

Author: boris-ilijic

cmd/pbm-agent/main.go

@@ -49,6 +49,20 @@ func main() {
 		versionFormat = versionCmd.Flag("format", "Output format <json or \"\">").
 				Default("").
 				String()
+
+		logPath = pbmCmd.Flag("log-path", "Path to file").
+			Envar("LOG_PATH").
+			Default("/dev/stderr").
+			String()
+		logJSON = pbmCmd.Flag("log-json", "Enable JSON output").
+			Envar("LOG_JSON").
+			Bool()
+		logLevel = pbmCmd.Flag(
+			"log-level",
+			"Minimal log level based on severity level: D, I, W, E or F, low to high. Choosing one includes higher levels too.").
+			Envar("LOG_LEVEL").
+			Default(log.D).
+			Enum(log.D, log.I, log.W, log.E, log.F)
 	)
 
 	cmd, err := pbmCmd.DefaultEnvars().Parse(os.Args[1:])
@@ -74,19 +88,24 @@ func main() {
 
 	hidecreds()
 
-	fmt.Print(perconaSquadNotice)
+	logOpts := &log.Opts{
+		LogPath:  *logPath,
+		LogLevel: *logLevel,
+		LogJSON:  *logJSON,
+	}
 
-	err = runAgent(url, *dumpConns)
+	err = runAgent(url, *dumpConns, logOpts)
 	stdlog.Println("Exit:", err)
 	if err != nil {
 		os.Exit(1)
 	}
 }
 
-func runAgent(mongoURI string, dumpConns int) error {
-	mtLog.SetDateFormat(log.LogTimeFormat)
-	mtLog.SetVerbosity(&options.Verbosity{VLevel: mtLog.DebugLow})
-
+func runAgent(
+	mongoURI string,
+	dumpConns int,
+	logOpts *log.Opts,
+) error {
 	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
 	defer cancel()
 
@@ -95,16 +114,34 @@ func runAgent(mongoURI string, dumpConns int) error {
 		return errors.Wrap(err, "connect to PBM")
 	}
 
+	err = setupNewDB(ctx, leadConn)
+	if err != nil {
+		return errors.Wrap(err, "setup pbm collections")
+	}
+
 	agent, err := newAgent(ctx, leadConn, mongoURI, dumpConns)
 	if err != nil {
 		return errors.Wrap(err, "connect to the node")
 	}
 
-	logger := log.New(agent.leadConn.LogCollection(), agent.brief.SetName, agent.brief.Me)
+	logger := log.NewWithOpts(
+		ctx,
+		agent.leadConn,
+		agent.brief.SetName,
+		agent.brief.Me,
+		logOpts)
 	defer logger.Close()
 
 	ctx = log.SetLoggerToContext(ctx, logger)
 
+	mtLog.SetDateFormat(log.LogTimeFormat)
+	mtLog.SetVerbosity(&options.Verbosity{VLevel: mtLog.DebugLow})
+	mtLog.SetWriter(logger)
+
+	logger.Printf(perconaSquadNotice)
+	logger.Printf("log options: log-path=%s, log-level:%s, log-json:%t",
+		logOpts.LogPath, logOpts.LogLevel, logOpts.LogJSON)
+
 	canRunSlicer := true
 	if err := agent.CanStart(ctx); err != nil {
 		if errors.Is(err, ErrArbiterNode) || errors.Is(err, ErrDelayedNode) {
@@ -114,11 +151,6 @@ func runAgent(mongoURI string, dumpConns int) error {
 		}
 	}
 
-	err = setupNewDB(ctx, agent.leadConn)
-	if err != nil {
-		return errors.Wrap(err, "setup pbm collections")
-	}
-
 	agent.showIncompatibilityWarning(ctx)
 
 	if canRunSlicer {

cmd/pbm-agent/oplog.go

@@ -77,7 +77,7 @@ func (a *Agent) OplogReplay(ctx context.Context, r *ctrl.ReplayCmd, opID ctrl.OP
 	}
 
 	l.Info("oplog replay started")
-	rr := restore.New(a.leadConn, a.nodeConn, a.brief, cfg, r.RSMap, 0)
+	rr := restore.New(a.leadConn, a.nodeConn, a.brief, cfg, r.RSMap, 0, 1)
 	err = rr.ReplayOplog(ctx, r, opID, l)
 	if err != nil {
 		if errors.Is(err, restore.ErrNoDataForShard) {

cmd/pbm-agent/restore.go

@@ -17,6 +17,10 @@ import (
 	"github.com/percona/percona-backup-mongodb/pbm/topo"
 )
 
+const (
+	numInsertionWorkersDefault = 1
+)
+
 func (a *Agent) Restore(ctx context.Context, r *ctrl.RestoreCmd, opid ctrl.OPID, ep config.Epoch) {
 	logger := log.FromContext(ctx)
 	if r == nil {
@@ -138,7 +142,14 @@ func (a *Agent) Restore(ctx context.Context, r *ctrl.RestoreCmd, opid ctrl.OPID,
 			numParallelColls = cfg.Restore.NumParallelCollections
 		}
 
-		rr := restore.New(a.leadConn, a.nodeConn, a.brief, cfg, r.RSMap, numParallelColls)
+		numInsertionWorkersPerCol := numInsertionWorkersDefault
+		if r.NumInsertionWorkers != nil && *r.NumInsertionWorkers > 0 {
+			numInsertionWorkersPerCol = int(*r.NumInsertionWorkers)
+		} else if cfg.Restore != nil && cfg.Restore.NumInsertionWorkers > 0 {
+			numInsertionWorkersPerCol = cfg.Restore.NumInsertionWorkers
+		}
+
+		rr := restore.New(a.leadConn, a.nodeConn, a.brief, cfg, r.RSMap, numParallelColls, numInsertionWorkersPerCol)
 		if r.OplogTS.IsZero() {
 			err = rr.Snapshot(ctx, r, opid, bcp)
 		} else {

cmd/pbm/main.go

@@ -244,6 +244,11 @@ func main() {
 		StringVar(&restore.pitrBase)
 	restoreCmd.Flag("num-parallel-collections", "Number of parallel collections").
 		Int32Var(&restore.numParallelColls)
+	restoreCmd.Flag("num-insertion-workers-per-collection",
+		"Specifies the number of insertion workers to run concurrently per collection. For large imports, "+
+			"increasing the number of insertion workers may increase the speed of the import.").
+		Default("1").
+		Int32Var(&restore.numInsertionWorkers)
 	restoreCmd.Flag("ns", `Namespaces to restore (e.g. "db1.*,db2.collection2"). If not set, restore all ("*.*")`).
 		StringVar(&restore.ns)
 	restoreCmd.Flag("ns-from", "Allows collection cloning (creating from the backup with different name) "+
@@ -495,7 +500,7 @@ func main() {
 		if err != nil {
 			exitErr(errors.Wrap(err, "connect to mongodb"), pbmOutF)
 		}
-		ctx = log.SetLoggerToContext(ctx, log.New(conn.LogCollection(), "", ""))
+		ctx = log.SetLoggerToContext(ctx, log.New(conn, "", ""))
 
 		ver, err := version.GetMongoVersion(ctx, conn.MongoClient())
 		if err != nil {

cmd/pbm/restore.go

@@ -53,7 +53,8 @@ type restoreOpts struct {
 	conf          string
 	ts            string
 
-	numParallelColls int32
+	numParallelColls    int32
+	numInsertionWorkers int32
 }
 
 type restoreRet struct {
@@ -124,6 +125,10 @@ func runRestore(
 	if err != nil {
 		return nil, errors.Wrap(err, "parse --num-parallel-collections option")
 	}
+	numInsertionWorkers, err := parseCLINumInsertionWorkersOption(o.numInsertionWorkers)
+	if err != nil {
+		return nil, errors.Wrap(err, "parse --num-insertion-workers option")
+	}
 	nss, err := parseCLINSOption(o.ns)
 	if err != nil {
 		return nil, errors.Wrap(err, "parse --ns option")
@@ -154,7 +159,7 @@ func runRestore(
 	}
 	tdiff := time.Now().Unix() - int64(clusterTime.T)
 
-	m, err := doRestore(ctx, conn, o, numParallelColls, nss, o.nsFrom, o.nsTo, rsMap, node, outf)
+	m, err := doRestore(ctx, conn, o, numParallelColls, numInsertionWorkers, nss, o.nsFrom, o.nsTo, rsMap, node, outf)
 	if err != nil {
 		return nil, err
 	}
@@ -353,12 +358,12 @@ func nsIsTaken(
 	ns string,
 ) error {
 	ns = strings.TrimSpace(ns)
-	db, coll, ok := strings.Cut(ns, ".")
+	dbName, coll, ok := strings.Cut(ns, ".")
 	if !ok {
 		return errors.Wrap(ErrInvalidNamespace, ns)
 	}
 
-	collNames, err := conn.MongoClient().Database(db).ListCollectionNames(ctx, bson.D{{"name", coll}})
+	collNames, err := conn.MongoClient().Database(dbName).ListCollectionNames(ctx, bson.D{{"name", coll}})
 	if err != nil {
 		return errors.Wrap(err, "list collection names for cloning target validation")
 	}
@@ -375,6 +380,7 @@ func doRestore(
 	conn connect.Client,
 	o *restoreOpts,
 	numParallelColls *int32,
+	numInsertionWorkers *int32,
 	nss []string,
 	nsFrom string,
 	nsTo string,
@@ -399,15 +405,16 @@ func doRestore(
 	cmd := ctrl.Cmd{
 		Cmd: ctrl.CmdRestore,
 		Restore: &ctrl.RestoreCmd{
-			Name:             name,
-			BackupName:       bcp,
-			NumParallelColls: numParallelColls,
-			Namespaces:       nss,
-			NamespaceFrom:    nsFrom,
-			NamespaceTo:      nsTo,
-			UsersAndRoles:    o.usersAndRoles,
-			RSMap:            rsMapping,
-			External:         o.extern,
+			Name:                name,
+			BackupName:          bcp,
+			NumParallelColls:    numParallelColls,
+			NumInsertionWorkers: numInsertionWorkers,
+			Namespaces:          nss,
+			NamespaceFrom:       nsFrom,
+			NamespaceTo:         nsTo,
+			UsersAndRoles:       o.usersAndRoles,
+			RSMap:               rsMapping,
+			External:            o.extern,
 		},
 	}
 	if o.pitr != "" {
@@ -804,3 +811,14 @@ func validateNSFromNSTo(o *restoreOpts) error {
 
 	return nil
 }
+
+func parseCLINumInsertionWorkersOption(value int32) (*int32, error) {
+	if value < 0 {
+		return nil, errors.New("Number of insertion workers has to be greater than zero.")
+	}
+	if value == 0 {
+		return nil, nil //nolint:nilnil
+	}
+
+	return &value, nil
+}

cmd/pbm/restore_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"errors"
+	"reflect"
 	"testing"
 )
 
@@ -101,3 +102,45 @@ func TestCloningValidation(t *testing.T) {
 		})
 	}
 }
+
+func TestParseCLINumInsertionWorkersOption(t *testing.T) {
+	var num int32 = 1
+
+	tests := []struct {
+		name    string
+		value   int32
+		want    *int32
+		wantErr bool
+	}{
+		{
+			name:    "valid number - no error",
+			value:   1,
+			want:    &num,
+			wantErr: false,
+		},
+		{
+			name:    "zero  - no error, but return nil",
+			value:   0,
+			want:    nil,
+			wantErr: false,
+		},
+		{
+			name:    "negative value  - error",
+			value:   -1,
+			want:    nil,
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := parseCLINumInsertionWorkersOption(tt.value)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("parseCLINumInsertionWorkersOption() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("parseCLINumInsertionWorkersOption() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

pbm/config/config.go

@@ -554,7 +554,7 @@ func confSetPITR(ctx context.Context, m connect.Client, value bool) error {
 // GetConfigVar returns value of given config vaiable
 func GetConfigVar(ctx context.Context, m connect.Client, key string) (interface{}, error) {
 	if !validateConfigKey(key) {
-		return nil, errors.New("invalid config key")
+		return nil, errors.Errorf("invalid config key: %s", key)
 	}
 
 	bts, err := m.ConfigCollection().

pbm/ctrl/cmd.go

@@ -157,7 +157,8 @@ type RestoreCmd struct {
 	UsersAndRoles bool              `bson:"usersAndRoles,omitempty"`
 	RSMap         map[string]string `bson:"rsMap,omitempty"`
 
-	NumParallelColls *int32 `bson:"numParallelColls,omitempty"`
+	NumParallelColls    *int32 `bson:"numParallelColls,omitempty"`
+	NumInsertionWorkers *int32 `bson:"numInsertionWorkers,omitempty"`
 
 	OplogTS primitive.Timestamp `bson:"oplogTS,omitempty"`
 

pbm/log/discard.go

@@ -28,6 +28,10 @@ func (discardLoggerImpl) PauseMgo() {
 func (discardLoggerImpl) ResumeMgo() {
 }
 
+func (discardLoggerImpl) Write([]byte) (int, error) {
+	return 0, nil
+}
+
 func (discardLoggerImpl) Printf(msg string, args ...any) {
 }
 

pbm/log/log.go

@@ -26,6 +26,8 @@ type Logger interface {
 	PauseMgo()
 	ResumeMgo()
 
+	Write(p []byte) (n int, err error)
+
 	Printf(msg string, args ...any)
 	Debug(event, obj, opid string, epoch primitive.Timestamp, msg string, args ...any)
 	Info(event, obj, opid string, epoch primitive.Timestamp, msg string, args ...any)
@@ -60,19 +62,44 @@ const (
 	Debug
 )
 
+const (
+	F = "F"
+	E = "E"
+	W = "W"
+	I = "I"
+	D = "D"
+)
+
 func (s Severity) String() string {
 	switch s {
 	case Fatal:
-		return "F"
+		return F
 	case Error:
-		return "E"
+		return E
 	case Warning:
-		return "W"
+		return W
 	case Info:
-		return "I"
+		return I
 	case Debug:
-		return "D"
+		return D
 	default:
 		return ""
 	}
 }
+
+func strToSeverity(s string) Severity {
+	switch s {
+	case F:
+		return Fatal
+	case E:
+		return Error
+	case W:
+		return Warning
+	case I:
+		return Info
+	case D:
+		return Debug
+	default:
+		return Debug
+	}
+}