fix(proposalKey): Proposal key initialisation with correct byte order…

…ing (#8622)

Proposal key initialisation has a bug where we want to reserve first 2
bytes for node id but do not do because we read random bytes in all 8
bytes and do a logical OR. This PR adds a test case and fixes the logic
of initialising the proposal key.

---------

Co-authored-by: Ahsan Barkati <[email protected]>

dgraph/cmd/zero/raft.go

@@ -88,11 +88,16 @@ func (n *node) AmLeader() bool {
 // {2 bytes Node ID} {4 bytes for random} {2 bytes zero}
 func (n *node) initProposalKey(id uint64) error {
 	x.AssertTrue(id != 0)
-	b := make([]byte, 8)
-	if _, err := rand.Read(b); err != nil {
+	random4Bytes := make([]byte, 4)
+	if _, err := rand.Read(random4Bytes); err != nil {
 		return err
 	}
-	proposalKey = n.Id<<48 | binary.BigEndian.Uint64(b)<<16
+	proposalKey = n.Id<<48 | uint64(binary.BigEndian.Uint32(random4Bytes))<<16
+	// We want to avoid spillage to node id in case of overflow. For instance, if the
+	// random bytes end up being [xx,xx, 255, 255, 255, 255, 0 , 0] (xx, xx being the node id)
+	// we would spill to node id after 65535 calls to unique key.
+	// So by setting 48th bit to 0 we ensure that we never spill out to node ids.
+	proposalKey &= ^(uint64(1) << 47)
 	return nil
 }
 

dgraph/cmd/zero/zero_test.go

@@ -21,11 +21,13 @@ import (
 	"math"
 	"testing"
 
-	"github.com/stretchr/testify/require"
-	"google.golang.org/grpc"
-
+	"github.com/dgraph-io/dgraph/conn"
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/testutil"
+	"github.com/dgraph-io/ristretto/z"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
 )
 
 func TestRemoveNode(t *testing.T) {
@@ -84,3 +86,27 @@ func TestIdBump(t *testing.T) {
 	_, err = zc.AssignIds(ctx, &pb.Num{Val: 10, Type: pb.Num_UID, Bump: true})
 	require.Contains(t, err.Error(), "Nothing to be leased")
 }
+
+func TestProposalKey(t *testing.T) {
+
+	id := uint64(2)
+	node := &node{Node: &conn.Node{Id: id}, ctx: context.Background(), closer: z.NewCloser(1)}
+	node.initProposalKey(node.Id)
+
+	pkey := proposalKey
+	nodeIdFromKey := proposalKey >> 48
+	require.Equal(t, id, nodeIdFromKey, "id extracted from proposal key is not equal to initial value")
+
+	valueOf48thBit := int(pkey & (1 << 48))
+	require.Equal(t, 0, valueOf48thBit, "48th bit is not set to zero on initialisation")
+
+	node.uniqueKey()
+	require.Equal(t, pkey+1, proposalKey, "proposal key should increment by 1 at each call of unique key")
+
+	uniqueKeys := make(map[uint64]struct{})
+	for i := 0; i < 10; i++ {
+		node.uniqueKey()
+		uniqueKeys[proposalKey] = struct{}{}
+	}
+	require.Equal(t, len(uniqueKeys), 10, "each iteration should create unique key")
+}

worker/proposal.go

@@ -24,15 +24,15 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/pkg/errors"
-	ostats "go.opencensus.io/stats"
-	tag "go.opencensus.io/tag"
-	otrace "go.opencensus.io/trace"
-
 	"github.com/dgraph-io/dgraph/conn"
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/schema"
 	"github.com/dgraph-io/dgraph/x"
+
+	"github.com/pkg/errors"
+	ostats "go.opencensus.io/stats"
+	tag "go.opencensus.io/tag"
+	otrace "go.opencensus.io/trace"
 )
 
 const baseTimeout time.Duration = 4 * time.Second