From 0433d6043f7a05aeaa4621b91899c873ac8d60da Mon Sep 17 00:00:00 2001
From: NamanJain8 <jnaman806@gmail.com>
Date: Tue, 20 Apr 2021 13:30:27 +0530
Subject: [PATCH 1/2] update xidmap lease

---
 dgraph/cmd/zero/assign.go |  3 ++-
 xidmap/xidmap.go          | 27 ++++++++++++++++++++++++---
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/dgraph/cmd/zero/assign.go b/dgraph/cmd/zero/assign.go
index 87151347765..5756de3f1eb 100644
--- a/dgraph/cmd/zero/assign.go
+++ b/dgraph/cmd/zero/assign.go
@@ -127,7 +127,8 @@ func (s *Server) lease(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error
 			howMany = num.Val + leaseBandwidth
 		}
 		if howMany < num.Val || maxLease+howMany < maxLease { // check for overflow.
-			return &emptyAssignedIds, errors.Errorf("Cannot lease %s as the limit has reached", typ)
+			return &emptyAssignedIds, errors.Errorf("Cannot lease %s as the limit has reached."+
+				" currMax:%d", typ, s.nextLease[typ]-1)
 		}
 
 		var proposal pb.ZeroProposal
diff --git a/xidmap/xidmap.go b/xidmap/xidmap.go
index d1bc428e276..1d8c1f24fc4 100644
--- a/xidmap/xidmap.go
+++ b/xidmap/xidmap.go
@@ -20,6 +20,9 @@ import (
 	"context"
 	"encoding/binary"
 	"math/rand"
+	"regexp"
+	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -36,6 +39,8 @@ import (
 	"github.com/golang/glog"
 )
 
+var maxLeaseRegex = regexp.MustCompile(`currMax:([0-9]+)`)
+
 // XidMapOptions specifies the options for creating a new xidmap.
 type XidMapOptions struct {
 	UidAssigner *grpc.ClientConn
@@ -290,12 +295,27 @@ func (m *XidMap) updateMaxSeen(max uint64) {
 // BumpTo can be used to make Zero allocate UIDs up to this given number. Attempts are made to
 // ensure all future allocations of UIDs be higher than this one, but results are not guaranteed.
 func (m *XidMap) BumpTo(uid uint64) {
-	curMax := atomic.LoadUint64(&m.maxUidSeen)
-	if uid <= curMax {
-		return
+	updateLease := func(msg string) {
+		if !strings.Contains(msg, "limit has reached. currMax:") {
+			return
+		}
+		matches := maxLeaseRegex.FindAllStringSubmatch(msg, 1)
+		if len(matches) == 0 {
+			return
+		}
+		maxUidLeased, err := strconv.ParseUint(matches[0][1], 10, 64)
+		if err != nil {
+			glog.Error("While parsing currMax %+v", err)
+			return
+		}
+		m.updateMaxSeen(maxUidLeased)
 	}
 
 	for {
+		curMax := atomic.LoadUint64(&m.maxUidSeen)
+		if uid <= curMax {
+			return
+		}
 		glog.V(1).Infof("Bumping up to %v", uid)
 		num := x.Max(uid-curMax, 1e4)
 		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
@@ -307,6 +327,7 @@ func (m *XidMap) BumpTo(uid uint64) {
 			m.updateMaxSeen(assigned.EndId)
 			return
 		}
+		updateLease(err.Error())
 		glog.Errorf("While requesting AssignUids(%d): %v", num, err)
 		if x.IsJwtExpired(err) {
 			if err := m.relogin(); err != nil {

From 3e4236107d0f78cdc859c909648cece826e2fc44 Mon Sep 17 00:00:00 2001
From: NamanJain8 <jnaman806@gmail.com>
Date: Wed, 21 Apr 2021 18:55:14 +0530
Subject: [PATCH 2/2] address comments

---
 xidmap/xidmap.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/xidmap/xidmap.go b/xidmap/xidmap.go
index 1d8c1f24fc4..1e1b6f1d27a 100644
--- a/xidmap/xidmap.go
+++ b/xidmap/xidmap.go
@@ -295,6 +295,9 @@ func (m *XidMap) updateMaxSeen(max uint64) {
 // BumpTo can be used to make Zero allocate UIDs up to this given number. Attempts are made to
 // ensure all future allocations of UIDs be higher than this one, but results are not guaranteed.
 func (m *XidMap) BumpTo(uid uint64) {
+	// If we have a cluster that cannot lease out new UIDs because it has already leased upto its
+	// max limit. Now, we try to live load the data with the given UIDs and the AssignIds complains
+	// that the limit has reached. Hence, update the xidmap's maxSeenUid and make progress.
 	updateLease := func(msg string) {
 		if !strings.Contains(msg, "limit has reached. currMax:") {
 			return
@@ -305,7 +308,7 @@ func (m *XidMap) BumpTo(uid uint64) {
 		}
 		maxUidLeased, err := strconv.ParseUint(matches[0][1], 10, 64)
 		if err != nil {
-			glog.Error("While parsing currMax %+v", err)
+			glog.Errorf("While parsing currMax %+v", err)
 			return
 		}
 		m.updateMaxSeen(maxUidLeased)