Merge branch 'dev' into PBM-1514-Improve-resync-performance

Author: veceraj

e2e-tests/cmd/pbm-test/run.go

@@ -109,6 +109,9 @@ func run(t *sharded.Cluster, typ testTyp) {
 		runTest("Distributed Transactions PITR",
 			t.DistributedTrxPITR)
 
+		runTest("Cleaning up sharded database for full restore",
+			t.CleanupFullRestore)
+
 		// disttxnconf := "/etc/pbm/fs-disttxn-4x.yaml"
 		// tsTo := primitive.Timestamp{1644410656, 8}
 

e2e-tests/pkg/tests/sharded/cmd.go

@@ -0,0 +1,34 @@
+package sharded
+
+import (
+	"context"
+	"log"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+func (c *Cluster) stopBalancer(ctx context.Context, conn *mongo.Client) {
+	err := conn.Database("admin").RunCommand(
+		ctx,
+		bson.D{{"balancerStop", 1}},
+	).Err()
+	if err != nil {
+		log.Fatalf("ERROR: stopping balancer: %v", err)
+	}
+}
+
+func (c *Cluster) moveChunk(ctx context.Context, db, col string, idx int, to string) {
+	log.Println("move chunk", idx, "to", to)
+	err := c.mongos.Conn().Database("admin").RunCommand(
+		ctx,
+		bson.D{
+			{"moveChunk", db + "." + col},
+			{"find", bson.M{"idx": idx}},
+			{"to", to},
+		},
+	).Err()
+	if err != nil {
+		log.Printf("ERROR: moveChunk %s.%s/idx:2000: %v", db, col, err)
+	}
+}

e2e-tests/pkg/tests/sharded/test_cleanup_full_restore.go

@@ -0,0 +1,106 @@
+package sharded
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+func (c *Cluster) CleanupFullRestore() {
+	ctx := context.Background()
+	db, col := "cleanupdb", "cleanupcol"
+	var numOfDocs int64 = 2000
+
+	conn := c.mongos.Conn()
+	c.stopBalancer(ctx, conn)
+	c.setupShardedCollection(ctx, conn, db, col)
+
+	err := c.mongos.GenData(db, col, 0, int64(numOfDocs))
+	if err != nil {
+		log.Fatalf("error: generate data: %v", err)
+	}
+
+	bcpName := c.LogicalBackup()
+	c.BackupWaitDone(context.Background(), bcpName)
+
+	c.splitChunkAt(ctx, conn, db, col, 1000)
+	c.moveChunk(ctx, db, col, 500, "rsx")
+
+	c.LogicalRestore(context.Background(), bcpName)
+	c.flushRouterConfig(ctx)
+
+	cnt, err := conn.Database(db).
+		Collection(col).
+		CountDocuments(ctx, bson.D{})
+	if err != nil {
+		log.Fatalf("ERROR: querying count: %v", err)
+	} else if cnt != numOfDocs {
+		log.Fatalf("ERROR: wrong number of docs: want=%d, got=%d", numOfDocs, cnt)
+	}
+}
+
+func (c *Cluster) setupShardedCollection(
+	ctx context.Context,
+	conn *mongo.Client,
+	db, col string,
+) {
+	err := conn.Database(db).
+		Collection(col).
+		Drop(ctx)
+	if err != nil {
+		log.Fatalf("ERROR: droping %s.%s collection: %v", db, col, err)
+	}
+
+	err = conn.Database("admin").RunCommand(
+		ctx,
+		bson.D{
+			{"enableSharding", db},
+			{"primaryShard", "rs1"},
+		},
+	).Err()
+	if err != nil {
+		log.Fatalf("ERROR: enableSharding on for %s primary shard rs: %v", db, err)
+	}
+
+	ns := fmt.Sprintf("%s.%s", db, col)
+	err = conn.Database(db).RunCommand(
+		ctx,
+		bson.D{{"create", col}},
+	).Err()
+	if err != nil {
+		log.Fatalf("ERROR: create %s collection: %v", ns, err)
+	}
+
+	err = conn.Database("admin").RunCommand(
+		ctx,
+		bson.D{
+			{"shardCollection", ns},
+			{"key", bson.M{"idx": 1}},
+		},
+	).Err()
+	if err != nil {
+		log.Fatalf("ERROR: shard %s collection: %v", ns, err)
+	}
+}
+
+func (c *Cluster) splitChunkAt(
+	ctx context.Context,
+	conn *mongo.Client,
+	db, col string,
+	id int,
+) {
+	ns := fmt.Sprintf("%s.%s", db, col)
+	err := conn.Database("admin").RunCommand(
+		ctx,
+		bson.D{
+			{"split", ns},
+			{"find", bson.M{"idx": id}},
+		},
+	).Err()
+	if err != nil {
+		log.Fatalf("ERROR: split %s collection: %v", ns, err)
+	}
+}

e2e-tests/pkg/tests/sharded/trx.go

@@ -47,22 +47,22 @@ func (c *Cluster) DistributedTransactions(bcp Backuper, col string) {
 
 	c.setupTrxCollection(ctx, col)
 
-	c.moveChunk(ctx, col, 0, "rs1")
-	c.moveChunk(ctx, col, 30, "rs1")
-	c.moveChunk(ctx, col, 89, "rs1")
-	c.moveChunk(ctx, col, 99, "rs1")
-	c.moveChunk(ctx, col, 110, "rs1")
-	c.moveChunk(ctx, col, 130, "rs1")
-	c.moveChunk(ctx, col, 131, "rs1")
-	c.moveChunk(ctx, col, 630, "rsx")
-	c.moveChunk(ctx, col, 530, "rsx")
-	c.moveChunk(ctx, col, 631, "rsx")
-	c.moveChunk(ctx, col, 730, "rsx")
-	c.moveChunk(ctx, col, 3000, "rsx")
-	c.moveChunk(ctx, col, 3001, "rsx")
-	c.moveChunk(ctx, col, 180, "rsx")
-	c.moveChunk(ctx, col, 199, "rsx")
-	c.moveChunk(ctx, col, 2001, "rsx")
+	c.moveChunk(ctx, trxdb, col, 0, "rs1")
+	c.moveChunk(ctx, trxdb, col, 30, "rs1")
+	c.moveChunk(ctx, trxdb, col, 89, "rs1")
+	c.moveChunk(ctx, trxdb, col, 99, "rs1")
+	c.moveChunk(ctx, trxdb, col, 110, "rs1")
+	c.moveChunk(ctx, trxdb, col, 130, "rs1")
+	c.moveChunk(ctx, trxdb, col, 131, "rs1")
+	c.moveChunk(ctx, trxdb, col, 630, "rsx")
+	c.moveChunk(ctx, trxdb, col, 530, "rsx")
+	c.moveChunk(ctx, trxdb, col, 631, "rsx")
+	c.moveChunk(ctx, trxdb, col, 730, "rsx")
+	c.moveChunk(ctx, trxdb, col, 3000, "rsx")
+	c.moveChunk(ctx, trxdb, col, 3001, "rsx")
+	c.moveChunk(ctx, trxdb, col, 180, "rsx")
+	c.moveChunk(ctx, trxdb, col, 199, "rsx")
+	c.moveChunk(ctx, trxdb, col, 2001, "rsx")
 
 	_, err = conn.Database(trxdb).Collection(col).DeleteMany(ctx, bson.M{})
 	if err != nil {
@@ -322,21 +322,6 @@ func (c *Cluster) setupTrxCollection(ctx context.Context, col string) {
 	}
 }
 
-func (c *Cluster) moveChunk(ctx context.Context, col string, idx int, to string) {
-	log.Println("move chunk", idx, "to", to)
-	err := c.mongos.Conn().Database("admin").RunCommand(
-		ctx,
-		bson.D{
-			{"moveChunk", trxdb + "." + col},
-			{"find", bson.M{"idx": idx}},
-			{"to", to},
-		},
-	).Err()
-	if err != nil {
-		log.Printf("ERROR: moveChunk %s.%s/idx:2000: %v", trxdb, col, err)
-	}
-}
-
 func (c *Cluster) checkTrxCollection(ctx context.Context, col string, bcp Backuper) {
 	log.Println("Checking restored data")
 

pbm/defs/defs.go

@@ -97,15 +97,16 @@ const (
 	// for phys restore, to indicate shards have been stopped
 	StatusDown Status = "down"
 
-	StatusStarting   Status = "starting"
-	StatusRunning    Status = "running"
-	StatusDumpDone   Status = "dumpDone"
-	StatusCopyReady  Status = "copyReady"
-	StatusCopyDone   Status = "copyDone"
-	StatusPartlyDone Status = "partlyDone"
-	StatusDone       Status = "done"
-	StatusCancelled  Status = "canceled"
-	StatusError      Status = "error"
+	StatusStarting       Status = "starting"
+	StatusCleanupCluster Status = "cleanupCluster"
+	StatusRunning        Status = "running"
+	StatusDumpDone       Status = "dumpDone"
+	StatusCopyReady      Status = "copyReady"
+	StatusCopyDone       Status = "copyDone"
+	StatusPartlyDone     Status = "partlyDone"
+	StatusDone           Status = "done"
+	StatusCancelled      Status = "canceled"
+	StatusError          Status = "error"
 
 	// status to communicate last op timestamp if it's not set
 	// during external restore

pbm/restore/logical.go

@@ -17,6 +17,7 @@ import (
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/writeconcern"
 
 	"github.com/percona/percona-backup-mongodb/pbm/archive"
 	"github.com/percona/percona-backup-mongodb/pbm/backup"
@@ -57,11 +58,11 @@ type Restore struct {
 	//
 	// Only the restore leader would have this info.
 	shards []topo.Shard
-	// rsMap is mapping between old and new replset names. used for data restore.
-	// empty if all replset names are the same
+	// rsMap is mapping between old and new replset names, used for data restore.
+	// It's empty if all replset names are the same.
 	rsMap map[string]string
-	// sMap is mapping between old and new shard names. used for router config update.
-	// empty if all shard names are the same
+	// sMap is mapping between old and new shard names, used for router config update.
+	// It's empty if all shard names are the same.
 	sMap map[string]string
 
 	log  log.LogEvent
@@ -76,6 +77,13 @@ type oplogRange struct {
 	storage storage.Storage
 }
 
+// configDatabasesDoc represents document in config.databases collection
+type configDatabasesDoc struct {
+	ID      string `bson:"_id"`
+	Primary string `bson:"primary"`
+	Version bson.D `bson:"version"`
+}
+
 // PBM restore from temp collections (pbmRUsers/pbmRRoles)should be used
 type restoreUsersAndRolesOption bool
 
@@ -240,7 +248,20 @@ func (r *Restore) Snapshot(
 		return err
 	}
 
-	err = r.toState(ctx, defs.StatusRunning, &defs.WaitActionStart)
+	err = r.toState(ctx, defs.StatusCleanupCluster, &defs.WaitActionStart)
+	if err != nil {
+		return err
+	}
+
+	// drop sharded dbs on sharded cluster, on each shard (not CSRS), only for full restore
+	if r.nodeInfo.IsSharded() && !r.nodeInfo.IsConfigSrv() && !util.IsSelective(nss) {
+		err = r.dropShardedDBs(ctx, bcp)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = r.toState(ctx, defs.StatusRunning, nil)
 	if err != nil {
 		return err
 	}
@@ -413,7 +434,20 @@ func (r *Restore) PITR(
 		return err
 	}
 
-	err = r.toState(ctx, defs.StatusRunning, &defs.WaitActionStart)
+	err = r.toState(ctx, defs.StatusCleanupCluster, &defs.WaitActionStart)
+	if err != nil {
+		return err
+	}
+
+	// drop sharded dbs on sharded cluster, on each shard (not CSRS), only for full restore
+	if r.nodeInfo.IsSharded() && !r.nodeInfo.IsConfigSrv() && !util.IsSelective(nss) {
+		err = r.dropShardedDBs(ctx, bcp)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = r.toState(ctx, defs.StatusRunning, nil)
 	if err != nil {
 		return err
 	}
@@ -828,6 +862,78 @@ func (r *Restore) toState(ctx context.Context, status defs.Status, wait *time.Du
 	return toState(ctx, r.leadConn, status, r.name, r.nodeInfo, r.reconcileStatus, wait)
 }
 
+// dropShardedDBs drop all sharded databases present in the backup.
+// Backup is specified with bcp parameter.
+// For each sharded database present in the backup _shardsvrDropDatabase command
+// is used to drop the database from the config srv and all shards.
+func (r *Restore) dropShardedDBs(ctx context.Context, bcp *backup.BackupMeta) error {
+	dbsInBcp, err := r.getDBsFromBackup(bcp)
+	if err != nil {
+		return errors.Wrap(err, "get dbs from backup")
+	}
+
+	// make cluster-wide drop for each db from the backup
+	for _, db := range dbsInBcp {
+		var configDBDoc configDatabasesDoc
+		err := r.leadConn.ConfigDatabase().
+			Collection("databases").
+			FindOne(ctx, bson.D{{"_id", db}}).
+			Decode(&configDBDoc)
+		if err != nil {
+			if errors.Is(err, mongo.ErrNoDocuments) {
+				continue
+			}
+			return errors.Wrapf(err, "get config.databases for %q", db)
+		}
+
+		if configDBDoc.Primary != r.nodeInfo.SetName {
+			// this shard is not primary shard for this db, so ignore it
+			continue
+		}
+
+		cmd := bson.D{
+			{"_shardsvrDropDatabase", 1},
+			{"databaseVersion", configDBDoc.Version},
+			{"writeConcern", writeconcern.Majority()},
+		}
+		res := r.nodeConn.Database(db).RunCommand(ctx, cmd)
+		if err := res.Err(); err != nil {
+			return errors.Wrapf(err, "_shardsvrDropDatabase for %q", db)
+		}
+		r.log.Debug("drop %q", db)
+	}
+
+	return nil
+}
+
+// getDBsFromBackup returns all databases present in backup metadata file
+// for each replicaset.
+func (r *Restore) getDBsFromBackup(bcp *backup.BackupMeta) ([]string, error) {
+	rsName := util.MakeReverseRSMapFunc(r.rsMap)(r.brief.SetName)
+	filepath := path.Join(bcp.Name, rsName, archive.MetaFile)
+	rdr, err := r.bcpStg.SourceReader(filepath)
+	if err != nil {
+		return nil, errors.Wrap(err, "read metadata file")
+	}
+	defer rdr.Close()
+
+	meta, err := archive.ReadMetadata(rdr)
+	if err != nil {
+		return nil, errors.Wrap(err, "get metadata")
+	}
+
+	uniqueDbs := map[string]bool{}
+	for _, ns := range meta.Namespaces {
+		uniqueDbs[ns.Database] = true
+	}
+	dbs := []string{}
+	for db := range uniqueDbs {
+		dbs = append(dbs, db)
+	}
+
+	return dbs, nil
+}
+
 func (r *Restore) RunSnapshot(
 	ctx context.Context,
 	dump string,