fix(restore): Bump uid and namespace after restore (#7790) (#7800)

Bump UID and namespace ID after reducephase in restore. To do
this bump we introduce an option to bump to a particular ID in zero.

(cherry picked from commit 99f4ca3)

dgraph/cmd/zero/assign.go

@@ -175,7 +175,9 @@ func (s *Server) lease(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error
 }
 
 // AssignIds is used to assign new ids (UIDs, NsIDs) by communicating with the leader of the
-// RAFT group responsible for handing out ids.
+// RAFT group responsible for handing out ids. If bump is set to true in the request then the
+// lease for the given id type is bumped to num.Val and {startId, endId} of the newly leased ids
+// in the process of bump is returned.
 func (s *Server) AssignIds(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
 	if ctx.Err() != nil {
 		return &emptyAssignedIds, ctx.Err()
@@ -246,6 +248,25 @@ func (s *Server) AssignIds(ctx context.Context, num *pb.Num) (*pb.AssignedIds, e
 		return err
 	}
 
+	// If this is a bump request and the current node is the leader then we create a normal lease
+	// request based on the number of required ids to reach the asked bump value. If the current
+	// node is not the leader then the bump request will be forwarded to the leader by lease().
+	if num.GetBump() && s.Node.AmLeader() {
+		s.leaseLock.Lock()
+		cur := s.nextLease[num.GetType()] - 1
+		s.leaseLock.Unlock()
+
+		// We need to lease more UIDs if bump request is more than current max lease.
+		req := num.GetVal()
+		if cur >= req {
+			return &emptyAssignedIds, errors.Errorf("Nothing to be leased")
+		}
+		num.Val = req - cur
+
+		// Set bump to false because we want to lease the required ids in the following request.
+		num.Bump = false
+	}
+
 	c := make(chan error, 1)
 	go func() {
 		c <- lease()

dgraph/cmd/zero/zero_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/testutil"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
 )
 
 func TestRemoveNode(t *testing.T) {
@@ -44,3 +45,41 @@ func TestIdLeaseOverflow(t *testing.T) {
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "limit has reached")
 }
+
+func TestIdBump(t *testing.T) {
+	dialOpts := []grpc.DialOption{
+		grpc.WithBlock(),
+		grpc.WithInsecure(),
+	}
+	ctx := context.Background()
+	con, err := grpc.DialContext(ctx, testutil.SockAddrZero, dialOpts...)
+	require.NoError(t, err)
+
+	zc := pb.NewZeroClient(con)
+
+	res, err := zc.AssignIds(ctx, &pb.Num{Val: 10, Type: pb.Num_UID})
+	require.NoError(t, err)
+	require.Equal(t, uint64(10), res.GetEndId()-res.GetStartId()+1)
+
+	// Next assignemnt's startId should be greater than 10.
+	res, err = zc.AssignIds(ctx, &pb.Num{Val: 50, Type: pb.Num_UID})
+	require.NoError(t, err)
+	require.Greater(t, res.GetStartId(), uint64(10))
+	require.Equal(t, uint64(50), res.GetEndId()-res.GetStartId()+1)
+
+	bumpTo := res.GetEndId() + 100000
+
+	// Bump the lease to (last result + 100000).
+	res, err = zc.AssignIds(ctx, &pb.Num{Val: bumpTo, Type: pb.Num_UID, Bump: true})
+	require.NoError(t, err)
+
+	// Next assignemnt's startId should be greater than bumpTo.
+	res, err = zc.AssignIds(ctx, &pb.Num{Val: 10, Type: pb.Num_UID})
+	require.NoError(t, err)
+	require.Greater(t, res.GetStartId(), bumpTo)
+	require.Equal(t, uint64(10), res.GetEndId()-res.GetStartId()+1)
+
+	// If bump request is less than maxLease, then it should result in no-op.
+	res, err = zc.AssignIds(ctx, &pb.Num{Val: 10, Type: pb.Num_UID, Bump: true})
+	require.Contains(t, err.Error(), "Nothing to be leased")
+}

ee/backup/run.go

@@ -266,7 +266,7 @@ func runExportBackup() error {
 			EncryptionKeyFile: ExportBackup.Conf.GetString("encryption_key_file"),
 			RestoreTs:         1,
 		}
-		if err := worker.RunMapper(req, mapDir); err != nil {
+		if _, err := worker.RunMapper(req, mapDir); err != nil {
 			return errors.Wrap(err, "Failed to map the backups")
 		}
 		in := &pb.ExportRequest{

protos/pb.proto

@@ -618,6 +618,9 @@ message Num {
   uint64 val = 1;
   bool read_only = 2;
   bool forwarded = 3;  // True if this request was forwarded by a peer.
+  // If bump is set to true then we bump the lease to val. If false, we assign new ids with count
+  // equal to val.
+  bool bump = 5;
   enum leaseType {
     NS_ID = 0;
     UID = 1;