Add an option to pass numInsertionWorkersPerCol from command line and from config file

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, 0)
 	err = rr.ReplayOplog(ctx, r, opID, l)
 	if err != nil {
 		if errors.Is(err, restore.ErrNoDataForShard) {

cmd/pbm-agent/restore.go

@@ -16,6 +16,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 {
@@ -128,7 +132,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,8 @@ 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.").
+		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
 	}
@@ -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
+}

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,
@@ -79,14 +80,7 @@ type oplogRange struct {
 type restoreUsersAndRolesOption bool
 
 // New creates a new restore object
-func New(
-	leadConn connect.Client,
-	nodeConn *mongo.Client,
-	brief topo.NodeBrief,
-	cfg *config.Config,
-	rsMap map[string]string,
-	numParallelColls int,
-) *Restore {
+func New(leadConn connect.Client, nodeConn *mongo.Client, brief topo.NodeBrief, cfg *config.Config, rsMap map[string]string, numParallelColls, numInsertionWorkersPerCol int) *Restore {
 	if rsMap == nil {
 		rsMap = make(map[string]string)
 	}
@@ -99,9 +93,9 @@ func New(
 
 		cfg: cfg,
 
-		numParallelColls: numParallelColls,
-
-		indexCatalog: idx.NewIndexCatalog(),
+		numParallelColls:          numParallelColls,
+		numInsertionWorkersPerCol: numInsertionWorkersPerCol,
+		indexCatalog:              idx.NewIndexCatalog(),
 	}
 }
 
@@ -1330,7 +1324,7 @@ 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 = 1
+	batchSizeDefault = 500
 )
 
 var ExcludeFromRestore = []string{
@@ -54,13 +53,7 @@ func (c *CloneNS) IsSpecified() bool {
 	return c.FromNS != "" && c.ToNS != ""
 }
 
-func NewRestore(
-	uri string,
-	cfg *config.Config,
-	cloneNS CloneNS,
-	numParallelColls int,
-	excludeRouterCollections bool,
-) (io.ReaderFrom, error) {
+func NewRestore(uri string, cfg *config.Config, cloneNS CloneNS, numParallelColls int, numInsertionWorkersPerCol int, excludeRouterCollections bool) (io.ReaderFrom, error) {
 	topts := options.New("mongorestore",
 		"0.0.1",
 		"none",
@@ -90,10 +83,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 +109,7 @@ func NewRestore(
 		BulkBufferSize:           batchSize,
 		BypassDocumentValidation: true,
 		Drop:                     true,
-		NumInsertionWorkers:      numInsertionWorkers,
+		NumInsertionWorkers:      numInsertionWorkersPerCol,
 		NumParallelCollections:   numParallelColls,
 		PreserveUUID:             preserveUUID,
 		StopOnError:              true,