Merge pull request #1053 from percona/PBM-1414-set-new-default-for-number-of-parallel-collections

PBM-1414: Set new default for number of parallel collections

Author: radoslawszulgo

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) "+

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/restore/logical.go

@@ -48,7 +48,8 @@ type Restore struct {
 	bcpStg   storage.Storage
 	oplogStg storage.Storage
 
-	numParallelColls int
+	numParallelColls          int
+	numInsertionWorkersPerCol int
 	// Shards to participate in restore. Num of shards in bcp could
 	// be less than in the cluster and this is ok. Only these shards
 	// would be expected to run restore (distributed transactions sync,
@@ -85,7 +86,8 @@ func New(
 	brief topo.NodeBrief,
 	cfg *config.Config,
 	rsMap map[string]string,
-	numParallelColls int,
+	numParallelColls,
+	numInsertionWorkersPerCol int,
 ) *Restore {
 	if rsMap == nil {
 		rsMap = make(map[string]string)
@@ -99,9 +101,9 @@ func New(
 
 		cfg: cfg,
 
-		numParallelColls: numParallelColls,
-
-		indexCatalog: idx.NewIndexCatalog(),
+		numParallelColls:          numParallelColls,
+		numInsertionWorkersPerCol: numInsertionWorkersPerCol,
+		indexCatalog:              idx.NewIndexCatalog(),
 	}
 }
 
@@ -1333,7 +1335,12 @@ func (r *Restore) applyOplog(ctx context.Context, ranges []oplogRange, options *
 }
 
 func (r *Restore) snapshot(input io.Reader, cloneNS snapshot.CloneNS, excludeRouterCollections bool) error {
-	rf, err := snapshot.NewRestore(r.brief.URI, r.cfg, cloneNS, r.numParallelColls, excludeRouterCollections)
+	rf, err := snapshot.NewRestore(
+		r.brief.URI,
+		r.cfg, cloneNS,
+		r.numParallelColls,
+		r.numInsertionWorkersPerCol,
+		excludeRouterCollections)
 	if err != nil {
 		return err
 	}

pbm/snapshot/restore.go

@@ -16,8 +16,7 @@ import (
 const (
 	preserveUUID = true
 
-	batchSizeDefault           = 500
-	numInsertionWorkersDefault = 10
+	batchSizeDefault = 500
 )
 
 var ExcludeFromRestore = []string{
@@ -54,11 +53,11 @@ func (c *CloneNS) IsSpecified() bool {
 	return c.FromNS != "" && c.ToNS != ""
 }
 
-func NewRestore(
-	uri string,
+func NewRestore(uri string,
 	cfg *config.Config,
 	cloneNS CloneNS,
-	numParallelColls int,
+	numParallelColls,
+	numInsertionWorkersPerCol int,
 	excludeRouterCollections bool,
 ) (io.ReaderFrom, error) {
 	topts := options.New("mongorestore",
@@ -90,10 +89,7 @@ func NewRestore(
 	if cfg.Restore.BatchSize > 0 {
 		batchSize = cfg.Restore.BatchSize
 	}
-	numInsertionWorkers := numInsertionWorkersDefault
-	if cfg.Restore.NumInsertionWorkers > 0 {
-		numInsertionWorkers = cfg.Restore.NumInsertionWorkers
-	}
+
 	if numParallelColls < 1 {
 		numParallelColls = 1
 	}
@@ -119,7 +115,7 @@ func NewRestore(
 		BulkBufferSize:           batchSize,
 		BypassDocumentValidation: true,
 		Drop:                     true,
-		NumInsertionWorkers:      numInsertionWorkers,
+		NumInsertionWorkers:      numInsertionWorkersPerCol,
 		NumParallelCollections:   numParallelColls,
 		PreserveUUID:             preserveUUID,
 		StopOnError:              true,

pbm/topo/cluster.go

@@ -58,7 +58,7 @@ func GetClusterTime(ctx context.Context, m connect.Client) (primitive.Timestamp,
 
 func ClusterTimeFromNodeInfo(info *NodeInfo) (primitive.Timestamp, error) {
 	if info.ClusterTime == nil {
-		return primitive.Timestamp{}, errors.New("no clusterTime in response")
+		return primitive.Timestamp{}, errors.Errorf("No clusterTime in response. Received: %+v", info)
 	}
 
 	return info.ClusterTime.ClusterTime, nil