diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go index 7b5d863f4e1..7695600610a 100644 --- a/dgraph/cmd/zero/raft.go +++ b/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 } diff --git a/dgraph/cmd/zero/zero_test.go b/dgraph/cmd/zero/zero_test.go index 9f4e18f113f..140b407b21e 100644 --- a/dgraph/cmd/zero/zero_test.go +++ b/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") +} diff --git a/worker/proposal.go b/worker/proposal.go index f218b9dc876..816d2cd23e1 100644 --- a/worker/proposal.go +++ b/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