REP-6465: Fix reliability of dotted-key sharding test.

This fixes intermittent failures in PR #127’s new dotted-key sharding test.

Author: FGasper

internal/util/sharding.go

@@ -7,6 +7,8 @@ import (
 	"github.com/pkg/errors"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/mongo/writeconcern"
 )
 
 const (
@@ -52,3 +54,31 @@ func GetShardKey(
 
 	return option.Some(key), nil
 }
+
+// Used in tests:
+func DisableBalancing(ctx context.Context, coll *mongo.Collection) error {
+	client := coll.Database().Client()
+	configDB := client.Database("config")
+
+	ns := FullName(coll)
+
+	_, err := configDB.
+		Collection(
+			"collections",
+			options.Collection().SetWriteConcern(writeconcern.Majority()),
+		).
+		UpdateOne(
+			ctx,
+			bson.D{{"_id", ns}},
+			bson.D{
+				{"$set", bson.D{
+					{"noBalance", true},
+				}},
+			},
+		)
+
+	if err != nil {
+		return errors.Wrapf(err, "disabling %#q’s shard balancing", ns)
+	}
+	return nil
+}

internal/verifier/migration_verifier_test.go

@@ -162,70 +162,115 @@ func (suite *IntegrationTestSuite) TestVerifier_Dotted_Shard_Key() {
 	dbName := suite.DBNameForTest()
 	collName := "coll"
 
-	verifier := suite.BuildVerifier()
-	srcColl := verifier.srcClient.Database(dbName).Collection(collName)
-	dstColl := verifier.dstClient.Database(dbName).Collection(collName)
-
 	docs := []bson.D{
-		// NB: It’s not possible to test with /foo.bar.baz or /foo/bar.baz
-		// because the server (as of v8) will route those to the same shard.
-		// Only /foo/bar/baz actually gets routed as per the shard key, so
-		// only that value can create a duplicate ID.
-		{{"_id", 33}, {"foo", bson.D{{"bar", bson.D{{"baz", 100}}}}}},
-		{{"_id", 33}, {"foo", bson.D{{"bar", bson.D{{"baz", 200}}}}}},
+		{{"_id", 33}, {"foo", bson.D{{"bar", 100}}}},
+		{{"_id", 33}, {"foo", bson.D{{"bar", 200}}}},
 	}
 
 	shardKey := bson.D{
-		{"foo.bar.baz", 1},
+		{"foo.bar", 1},
 	}
 
-	for _, coll := range mslices.Of(srcColl, dstColl) {
-		db := coll.Database()
-		client := db.Client()
+	keyField := shardKey[0].Key
+	splitKey := bson.D{{keyField, 150}}
 
-		// For sharded, pre-v8 clusters we need to create the collection first.
-		require.NoError(db.CreateCollection(ctx, coll.Name()))
-		require.NoError(client.Database("admin").RunCommand(
-			ctx,
-			bson.D{{"enableSharding", db.Name()}},
-		).Err())
+	clientsWithLabels := []struct {
+		label  string
+		client *mongo.Client
+	}{
+		{"source", suite.srcMongoClient},
+		{"destination", suite.dstMongoClient},
+	}
 
-		require.NoError(client.Database("admin").RunCommand(
-			ctx,
-			bson.D{
-				{"shardCollection", db.Name() + "." + coll.Name()},
-				{"key", shardKey},
-			},
-		).Err(),
-		)
+	for _, clientWithLabel := range clientsWithLabels {
+		client := clientWithLabel.client
+		clientLabel := clientWithLabel.label
+
+		db := client.Database(dbName)
+		coll := db.Collection(collName)
+
+		// The DB shouldn’t exist anyway, but just in case.
+		require.NoError(db.Drop(ctx), "should drop database")
 
 		shardIds := getShardIds(suite.T(), client)
 
 		admin := client.Database("admin")
-		keyField := "foo.bar.baz"
-		splitKey := bson.D{{keyField, 150}}
+
+		require.NoError(admin.RunCommand(
+			ctx,
+			bson.D{
+				{"enableSharding", db.Name()},
+			},
+		).Err())
 
 		require.NoError(
-			admin.RunCommand(ctx, bson.D{
-				{"split", db.Name() + "." + coll.Name()},
-				{"middle", splitKey},
-			}).Err(),
+			admin.RunCommand(
+				ctx,
+				bson.D{
+					{"shardCollection", FullName(coll)},
+					{"key", shardKey},
+				},
+			).Err(),
+			"should shard collection on %s",
+			clientLabel,
 		)
 
 		require.NoError(
-			admin.RunCommand(ctx, bson.D{
-				{"moveChunk", db.Name() + "." + coll.Name()},
-				{"to", shardIds[0]},
-				{"find", bson.D{{keyField, 149}}},
-			}).Err(),
+			util.DisableBalancing(ctx, coll),
+			"should disable %#q’s balancing on %s",
+			FullName(coll),
+			clientLabel,
 		)
 
 		require.NoError(
 			admin.RunCommand(ctx, bson.D{
-				{"moveChunk", db.Name() + "." + coll.Name()},
-				{"to", shardIds[1]},
-				{"find", bson.D{{keyField, 151}}},
+				{"split", FullName(coll)},
+				{"middle", splitKey},
 			}).Err(),
+			"should split on %s",
+			clientLabel,
+		)
+
+		require.Eventually(
+			func() bool {
+				err := admin.RunCommand(ctx, bson.D{
+					{"moveChunk", FullName(coll)},
+					{"find", bson.D{{keyField, 149}}},
+					{"to", shardIds[0]},
+					{"_waitForDelete", true},
+				}).Err()
+
+				if err != nil {
+					suite.T().Logf("Failed to move %s’s lower chunk to shard %#q: %v", clientLabel, shardIds[0], err)
+					return false
+				}
+
+				return true
+			},
+			5*time.Minute,
+			time.Second,
+			"Should move lower chunk to the 1st shard",
+		)
+
+		require.Eventually(
+			func() bool {
+				err := admin.RunCommand(ctx, bson.D{
+					{"moveChunk", FullName(coll)},
+					{"find", bson.D{{keyField, 151}}},
+					{"to", shardIds[1]},
+					{"_waitForDelete", true},
+				}).Err()
+
+				if err != nil {
+					suite.T().Logf("Failed to move %s’s upper chunk to shard %#q: %v", clientLabel, shardIds[1], err)
+					return false
+				}
+
+				return true
+			},
+			5*time.Minute,
+			time.Second,
+			"Should move upper chunk to the 2nd shard",
 		)
 
 		_, err := coll.InsertMany(ctx, lo.ToAnySlice(lo.Shuffle(docs)))
@@ -252,6 +297,7 @@ func (suite *IntegrationTestSuite) TestVerifier_Dotted_Shard_Key() {
 		},
 	}
 
+	verifier := suite.BuildVerifier()
 	results, docCount, _, err := verifier.FetchAndCompareDocuments(ctx, 0, task)
 	require.NoError(err, "should fetch & compare")
 	assert.EqualValues(suite.T(), len(docs), docCount, "expected # of docs")