[BREAKING] feat(zero): Make zero lease out namespace IDs (#7341)

This change makes zero lease out the namespace IDs. It
changes the `AssignUids()` API to `AssignIds()`. The new
API takes in the lease type and accordingly leases out.

Author: ahsanbarkati

dgraph/cmd/alpha/run_test.go

@@ -1694,7 +1694,8 @@ func TestMain(m *testing.M) {
 		log.Fatal(err)
 	}
 	zc := pb.NewZeroClient(conn)
-	if _, err := zc.AssignUids(context.Background(), &pb.Num{Val: 1e6}); err != nil {
+	if _, err := zc.AssignIds(context.Background(),
+		&pb.Num{Val: 1e6, Type: pb.Num_UID}); err != nil {
 		log.Fatal(err)
 	}
 	httpToken := testutil.GrootHttpLogin(addr + "/admin")

dgraph/cmd/debug/run.go

@@ -853,8 +853,10 @@ func printZeroProposal(buf *bytes.Buffer, zpr *pb.ZeroProposal) {
 		fmt.Fprintf(buf, " Member: %+v .", zpr.Member)
 	case zpr.Tablet != nil:
 		fmt.Fprintf(buf, " Tablet: %+v .", zpr.Tablet)
-	case zpr.MaxLeaseId > 0:
-		fmt.Fprintf(buf, " MaxLeaseId: %d .", zpr.MaxLeaseId)
+	case zpr.MaxUID > 0:
+		fmt.Fprintf(buf, " MaxUID: %d .", zpr.MaxUID)
+	case zpr.MaxNsID > 0:
+		fmt.Fprintf(buf, " MaxNsID: %d .", zpr.MaxNsID)
 	case zpr.MaxRaftId > 0:
 		fmt.Fprintf(buf, " MaxRaftId: %d .", zpr.MaxRaftId)
 	case zpr.MaxTxnTs > 0:

dgraph/cmd/zero/assign.go

@@ -36,24 +36,30 @@ const (
 func (s *Server) updateLeases() {
 	var startTs uint64
 	s.Lock()
-	s.nextLeaseId = s.state.MaxLeaseId + 1
-	s.nextTxnTs = s.state.MaxTxnTs + 1
-	startTs = s.nextTxnTs
-	glog.Infof("Updated Lease id: %d. Txn Ts: %d", s.nextLeaseId, s.nextTxnTs)
+	s.nextLease[pb.Num_UID] = s.state.MaxUID + 1
+	s.nextLease[pb.Num_TXN_TS] = s.state.MaxTxnTs + 1
+	s.nextLease[pb.Num_NS_ID] = s.state.MaxNsID + 1
+
+	startTs = s.nextLease[pb.Num_TXN_TS]
+	glog.Infof("Updated UID: %d. Txn Ts: %d. NsID: %d.",
+		s.nextLease[pb.Num_UID], s.nextLease[pb.Num_TXN_TS], s.nextLease[pb.Num_NS_ID])
 	s.Unlock()
 	s.orc.updateStartTxnTs(startTs)
 }
 
-func (s *Server) maxLeaseId() uint64 {
-	s.RLock()
-	defer s.RUnlock()
-	return s.state.MaxLeaseId
-}
-
-func (s *Server) maxTxnTs() uint64 {
+func (s *Server) maxLease(typ pb.NumLeaseType) uint64 {
 	s.RLock()
 	defer s.RUnlock()
-	return s.state.MaxTxnTs
+	var maxlease uint64
+	switch typ {
+	case pb.Num_UID:
+		maxlease = s.state.MaxUID
+	case pb.Num_TXN_TS:
+		maxlease = s.state.MaxTxnTs
+	case pb.Num_NS_ID:
+		maxlease = s.state.MaxNsID
+	}
+	return maxlease
 }
 
 var errServedFromMemory = errors.New("Lease was served from memory")
@@ -62,7 +68,8 @@ var errServedFromMemory = errors.New("Lease was served from memory")
 // This function is triggered by an RPC call. We ensure that only leader can assign new UIDs,
 // so we can tackle any collisions that might happen with the leasemanager
 // In essence, we just want one server to be handing out new uids.
-func (s *Server) lease(ctx context.Context, num *pb.Num, txn bool) (*pb.AssignedIds, error) {
+func (s *Server) lease(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+	typ := num.GetType()
 	node := s.Node
 	// TODO: Fix when we move to linearizable reads, need to check if we are the leader, might be
 	// based on leader leases. If this node gets partitioned and unless checkquorum is enabled, this
@@ -75,21 +82,21 @@ func (s *Server) lease(ctx context.Context, num *pb.Num, txn bool) (*pb.Assigned
 		return &emptyAssignedIds, errors.Errorf("Nothing to be leased")
 	}
 	if glog.V(3) {
-		glog.Infof("Got lease request for txn: %v. Num: %+v\n", txn, num)
+		glog.Infof("Got lease request for Type: %v. Num: %+v\n", typ, num)
 	}
 
 	s.leaseLock.Lock()
 	defer s.leaseLock.Unlock()
 
-	if txn {
+	if typ == pb.Num_TXN_TS {
 		if num.Val == 0 && num.ReadOnly {
 			// If we're only asking for a readonly timestamp, we can potentially
 			// service it directly.
 			if glog.V(3) {
 				glog.Infof("Attempting to serve read only txn ts [%d, %d]",
-					s.readOnlyTs, s.nextTxnTs)
+					s.readOnlyTs, s.nextLease[pb.Num_TXN_TS])
 			}
-			if s.readOnlyTs > 0 && s.readOnlyTs == s.nextTxnTs-1 {
+			if s.readOnlyTs > 0 && s.readOnlyTs == s.nextLease[pb.Num_TXN_TS]-1 {
 				return &pb.AssignedIds{ReadOnly: s.readOnlyTs}, errServedFromMemory
 			}
 		}
@@ -106,23 +113,24 @@ func (s *Server) lease(ctx context.Context, num *pb.Num, txn bool) (*pb.Assigned
 		howMany = num.Val + leaseBandwidth
 	}
 
-	if s.nextLeaseId == 0 || s.nextTxnTs == 0 {
+	if s.nextLease[pb.Num_UID] == 0 || s.nextLease[pb.Num_TXN_TS] == 0 ||
+		s.nextLease[pb.Num_NS_ID] == 0 {
 		return nil, errors.New("Server not initialized")
 	}
 
-	var maxLease, available uint64
 	var proposal pb.ZeroProposal
 
 	// Calculate how many ids do we have available in memory, before we need to
 	// renew our lease.
-	if txn {
-		maxLease = s.maxTxnTs()
-		available = maxLease - s.nextTxnTs + 1
+	maxLease := s.maxLease(typ)
+	available := maxLease - s.nextLease[typ] + 1
+	switch typ {
+	case pb.Num_TXN_TS:
 		proposal.MaxTxnTs = maxLease + howMany
-	} else {
-		maxLease = s.maxLeaseId()
-		available = maxLease - s.nextLeaseId + 1
-		proposal.MaxLeaseId = maxLease + howMany
+	case pb.Num_UID:
+		proposal.MaxUID = maxLease + howMany
+	case pb.Num_NS_ID:
+		proposal.MaxNsID = maxLease + howMany
 	}
 
 	// If we have less available than what we need, we need to renew our lease.
@@ -134,48 +142,55 @@ func (s *Server) lease(ctx context.Context, num *pb.Num, txn bool) (*pb.Assigned
 	}
 
 	out := &pb.AssignedIds{}
-	if txn {
+	if typ == pb.Num_TXN_TS {
 		if num.Val > 0 {
-			out.StartId = s.nextTxnTs
+			out.StartId = s.nextLease[pb.Num_TXN_TS]
 			out.EndId = out.StartId + num.Val - 1
-			s.nextTxnTs = out.EndId + 1
+			s.nextLease[pb.Num_TXN_TS] = out.EndId + 1
 		}
 		if num.ReadOnly {
-			s.readOnlyTs = s.nextTxnTs
-			s.nextTxnTs++
+			s.readOnlyTs = s.nextLease[pb.Num_TXN_TS]
+			s.nextLease[pb.Num_TXN_TS]++
 			out.ReadOnly = s.readOnlyTs
 		}
 		s.orc.doneUntil.Begin(x.Max(out.EndId, out.ReadOnly))
-	} else {
-		out.StartId = s.nextLeaseId
+	} else if typ == pb.Num_UID {
+		out.StartId = s.nextLease[pb.Num_UID]
+		out.EndId = out.StartId + num.Val - 1
+		s.nextLease[pb.Num_UID] = out.EndId + 1
+	} else if typ == pb.Num_NS_ID {
+		out.StartId = s.nextLease[pb.Num_NS_ID]
 		out.EndId = out.StartId + num.Val - 1
-		s.nextLeaseId = out.EndId + 1
+		s.nextLease[pb.Num_NS_ID] = out.EndId + 1
+
+	} else {
+		return out, errors.Errorf("Unknown lease type: %v\n", typ)
 	}
 	return out, nil
 }
 
-// AssignUids is used to assign new uids by communicating with the leader of the RAFT group
-// responsible for handing out uids.
-func (s *Server) AssignUids(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.
+func (s *Server) AssignIds(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
 	if ctx.Err() != nil {
 		return &emptyAssignedIds, ctx.Err()
 	}
-	ctx, span := otrace.StartSpan(ctx, "Zero.AssignUids")
+	ctx, span := otrace.StartSpan(ctx, "Zero.AssignIds")
 	defer span.End()
 
 	reply := &emptyAssignedIds
 	lease := func() error {
 		var err error
 		if s.Node.AmLeader() {
 			span.Annotatef(nil, "Zero leader leasing %d ids", num.GetVal())
-			reply, err = s.lease(ctx, num, false)
+			reply, err = s.lease(ctx, num)
 			return err
 		}
 		span.Annotate(nil, "Not Zero leader")
 		// I'm not the leader and this request was forwarded to me by a peer, who thought I'm the
 		// leader.
 		if num.Forwarded {
-			return errors.Errorf("Invalid Zero received AssignUids request forward. Please retry")
+			return errors.Errorf("Invalid Zero received AssignIds request forward. Please retry")
 		}
 		// This is an original request. Forward it to the leader.
 		pl := s.Leader(0)
@@ -185,7 +200,7 @@ func (s *Server) AssignUids(ctx context.Context, num *pb.Num) (*pb.AssignedIds,
 		span.Annotatef(nil, "Sending request to %v", pl.Addr)
 		zc := pb.NewZeroClient(pl.Get())
 		num.Forwarded = true
-		reply, err = zc.AssignUids(ctx, num)
+		reply, err = zc.AssignIds(ctx, num)
 		return err
 	}
 

dgraph/cmd/zero/http.go

@@ -73,12 +73,17 @@ func (st *state) assign(w http.ResponseWriter, r *http.Request) {
 	what := r.URL.Query().Get("what")
 	switch what {
 	case "uids":
-		ids, err = st.zero.AssignUids(ctx, num)
+		num.Type = pb.Num_UID
+		ids, err = st.zero.AssignIds(ctx, num)
 	case "timestamps":
+		num.Type = pb.Num_TXN_TS
 		if num.Val == 0 {
 			num.ReadOnly = true
 		}
 		ids, err = st.zero.Timestamps(ctx, num)
+	case "nsids":
+		num.Type = pb.Num_NS_ID
+		ids, err = st.zero.AssignIds(ctx, num)
 	default:
 		x.SetStatus(w, x.Error,
 			fmt.Sprintf("Invalid what: [%s]. Must be one of uids or timestamps", what))

dgraph/cmd/zero/oracle.go

@@ -402,8 +402,8 @@ func (s *Server) commit(ctx context.Context, src *api.TxnContext) error {
 		return s.proposeTxn(ctx, src)
 	}
 
-	num := pb.Num{Val: 1}
-	assigned, err := s.lease(ctx, &num, true)
+	num := pb.Num{Val: 1, Type: pb.Num_TXN_TS}
+	assigned, err := s.lease(ctx, &num)
 	if err != nil {
 		return err
 	}
@@ -512,7 +512,8 @@ func (s *Server) Timestamps(ctx context.Context, num *pb.Num) (*pb.AssignedIds,
 		return &emptyAssignedIds, ctx.Err()
 	}
 
-	reply, err := s.lease(ctx, num, true)
+	num.Type = pb.Num_TXN_TS
+	reply, err := s.lease(ctx, num)
 	span.Annotatef(nil, "Response: %+v. Error: %v", reply, err)
 
 	switch err {

dgraph/cmd/zero/raft.go

@@ -418,15 +418,18 @@ func (n *node) applyProposal(e raftpb.Entry) (uint64, error) {
 	}
 
 	switch {
-	case p.MaxLeaseId > state.MaxLeaseId:
-		state.MaxLeaseId = p.MaxLeaseId
+	case p.MaxUID > state.MaxUID:
+		state.MaxUID = p.MaxUID
 	case p.MaxTxnTs > state.MaxTxnTs:
 		state.MaxTxnTs = p.MaxTxnTs
-	case p.MaxLeaseId != 0 || p.MaxTxnTs != 0:
+	case p.MaxNsID > state.MaxNsID:
+		state.MaxNsID = p.MaxNsID
+	case p.MaxUID != 0 || p.MaxTxnTs != 0 || p.MaxNsID != 0:
 		// Could happen after restart when some entries were there in WAL and did not get
 		// snapshotted.
-		glog.Infof("Could not apply proposal, ignoring: p.MaxLeaseId=%v, p.MaxTxnTs=%v maxLeaseId=%d"+
-			" maxTxnTs=%d\n", p.MaxLeaseId, p.MaxTxnTs, state.MaxLeaseId, state.MaxTxnTs)
+		glog.Infof("Could not apply proposal, ignoring: p.MaxUID=%v, p.MaxTxnTs=%v"+
+			"p.MaxNsID=%v, maxUID=%d maxTxnTs=%d maxNsID=%d\n",
+			p.MaxUID, p.MaxTxnTs, p.MaxNsID, state.MaxUID, state.MaxTxnTs, state.MaxNsID)
 	}
 	if p.Txn != nil {
 		n.server.orc.updateCommitStatus(e.Index, p.Txn)

dgraph/cmd/zero/zero.go

@@ -59,10 +59,9 @@ type Server struct {
 	state       *pb.MembershipState
 	nextRaftId  uint64
 
-	nextLeaseId uint64
-	nextTxnTs   uint64
-	readOnlyTs  uint64
-	leaseLock   sync.Mutex // protects nextLeaseId, nextTxnTs and corresponding proposals.
+	nextLease  map[pb.NumLeaseType]uint64
+	readOnlyTs uint64
+	leaseLock  sync.Mutex // protects nextUID, nextTxnTs, nextNsID and corresponding proposals.
 
 	// groupMap    map[uint32]*Group
 	nextGroup      uint32
@@ -90,9 +89,11 @@ func (s *Server) Init() {
 		Groups: make(map[uint32]*pb.Group),
 		Zeros:  make(map[uint64]*pb.Member),
 	}
+	s.nextLease = make(map[pb.NumLeaseType]uint64)
 	s.nextRaftId = 1
-	s.nextLeaseId = 1
-	s.nextTxnTs = 1
+	s.nextLease[pb.Num_UID] = 1
+	s.nextLease[pb.Num_TXN_TS] = 1
+	s.nextLease[pb.Num_NS_ID] = 1
 	s.nextGroup = 1
 	s.leaderChangeCh = make(chan struct{}, 1)
 	s.closer = z.NewCloser(2) // grpc and http

ee/backup/run.go

@@ -246,7 +246,7 @@ func runRestoreCmd() error {
 		if result.MaxLeaseUid > 0 {
 			ctx, cancelUid := context.WithTimeout(context.Background(), time.Minute)
 			defer cancelUid()
-			if _, err = zc.AssignUids(ctx, &pb.Num{Val: result.MaxLeaseUid}); err != nil {
+			if _, err = zc.AssignIds(ctx, &pb.Num{Val: result.MaxLeaseUid, Type: pb.Num_UID}); err != nil {
 				fmt.Printf("Failed to assign maxLeaseId %d in Zero: %v\n", result.MaxLeaseUid, err)
 				return err
 			}

graphql/admin/admin.go

@@ -140,7 +140,8 @@ const (
 		counter: Int
 		groups: [ClusterGroup]
 		zeros: [Member]
-		maxLeaseId: Int
+		maxUID: Int
+		maxNsID: Int
 		maxTxnTs: Int
 		maxRaftId: Int
 		removed: [Member]

graphql/admin/state.go

@@ -15,15 +15,16 @@ import (
 )
 
 type membershipState struct {
-	Counter    uint64         `json:"counter,omitempty"`
-	Groups     []clusterGroup `json:"groups,omitempty"`
-	Zeros      []*pb.Member   `json:"zeros,omitempty"`
-	MaxLeaseId uint64         `json:"maxLeaseId,omitempty"`
-	MaxTxnTs   uint64         `json:"maxTxnTs,omitempty"`
-	MaxRaftId  uint64         `json:"maxRaftId,omitempty"`
-	Removed    []*pb.Member   `json:"removed,omitempty"`
-	Cid        string         `json:"cid,omitempty"`
-	License    *pb.License    `json:"license,omitempty"`
+	Counter   uint64         `json:"counter,omitempty"`
+	Groups    []clusterGroup `json:"groups,omitempty"`
+	Zeros     []*pb.Member   `json:"zeros,omitempty"`
+	MaxUID    uint64         `json:"maxUID,omitempty"`
+	MaxNsID   uint64         `json:"maxNsID,omitempty"`
+	MaxTxnTs  uint64         `json:"maxTxnTs,omitempty"`
+	MaxRaftId uint64         `json:"maxRaftId,omitempty"`
+	Removed   []*pb.Member   `json:"removed,omitempty"`
+	Cid       string         `json:"cid,omitempty"`
+	License   *pb.License    `json:"license,omitempty"`
 }
 
 type clusterGroup struct {
@@ -98,8 +99,9 @@ func convertToGraphQLResp(ms pb.MembershipState) membershipState {
 	for _, v := range ms.Zeros {
 		state.Zeros = append(state.Zeros, v)
 	}
-	state.MaxLeaseId = ms.MaxLeaseId
+	state.MaxUID = ms.MaxUID
 	state.MaxTxnTs = ms.MaxTxnTs
+	state.MaxNsID = ms.MaxNsID
 	state.MaxRaftId = ms.MaxRaftId
 	state.Removed = ms.Removed
 	state.Cid = ms.Cid