fix(zero): add a flag to limit the uid lease per namespace #7568

NamanJain8 · 2021-03-15T09:56:41Z

Fixes DGRAPH-3124

There might be a scenario where in a multi-tenant cluster a malicious user leases out a large number of UIDs or bump it to a large number, causing an issue for the other tenants.
This PR adds a super flag --limit to zero. This flag limits the number of UIDs (uid-lease=100000;) that can be leased by a namespace within an interval specified by refill-interval=30s;. This PR also changes the xidmap's New() API as we need to pass in the namespace and dgo client.

Along with it, it does some minor cleaning where flag parsing was in the critical path.

This change is

worker/zero_proxy.go

manishrjain

Looks like there're a lot of changes. So, reduce the scope.

Reviewed 10 of 18 files at r2, 3 of 3 files at r3, 5 of 6 files at r4.
Reviewable status: all files reviewed, 14 unresolved discussions (waiting on @NamanJain8 and @vvbalaji-dgraph)

dgraph/cmd/bulk/loader.go, line 192 at r4 (raw file):

	for {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		ctx = x.AttachNamespaceOutgoing(ctx, x.GalaxyNamespace)

We don't need to do this. If Zero doesn't find a namespace, then just assume GalaxyNs.

dgraph/cmd/live/run.go, line 634 at r4 (raw file):

	x.Checkf(err, "Unable to connect to zero, Is it running at %s?", opt.zero)

	creds := z.NewSuperFlag(Live.Conf.GetString("creds")).MergeAndCheckDefault(x.DefaultCreds)

same over here?

dgraph/cmd/live/run.go, line 636 at r4 (raw file):

	creds := z.NewSuperFlag(Live.Conf.GetString("creds")).MergeAndCheckDefault(x.DefaultCreds)
	alloc := xidmap.New(xidmap.XidMapOptions{
		Zero:      connzero,

Might want to rename, if this could be a conn to Alpha.

dgraph/cmd/zero/assign.go, line 184 at r4 (raw file):

	defer span.End()

	validateAndGetToken := func() error {

rateLimit

dgraph/cmd/zero/assign.go, line 186 at r4 (raw file):

	validateAndGetToken := func() error {
		if num.GetType() != pb.Num_UID {
			// We only rate limit lease of UIDs.

same thing would apply to timestamp too.

dgraph/cmd/zero/assign.go, line 194 at r4 (raw file):

			return err
		}
		if num.Val > opts.limit.GetUint64("uid-lease") {

Parse it once, and use it. Don't parse it again and again.

dgraph/cmd/zero/assign.go, line 201 at r4 (raw file):

		if !s.rateLimiter.Allow(ns, num.Val) {
			// Return error after random delay.
			delay := rand.Intn(int(opts.limit.GetDuration("refill-interval")))

parse once. For rand, use a locally generated rand. rand.NewSource rand.NewRand or something.

dgraph/cmd/zero/assign.go, line 203 at r4 (raw file):

			delay := rand.Intn(int(opts.limit.GetDuration("refill-interval")))
			time.Sleep(time.Duration(delay) * time.Second)
			return errors.New("Cannot lease UID because UID lease for the namespace is exhausted." +

for namespace: %#x is exhausted.

dgraph/cmd/zero/assign.go, line 236 at r4 (raw file):

		// pass on the incoming metadata to the zero leader.
		if md, ok := metadata.FromIncomingContext(ctx); ok {
			ctx = metadata.NewOutgoingContext(ctx, md)

Do you need to do this?

dgraph/cmd/zero/http.go, line 71 at r4 (raw file):

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	// Consider all the request coming through /assign belong to GalaxyNamespace.
	ctx = x.AttachNamespace(ctx, x.GalaxyNamespace)

can remove.

dgraph/cmd/zero/run.go, line 109 at r4 (raw file):

			in an interval specified by refill-interval.`).
		Flag("refill-interval",
			"The interval after which the tokens for UID lease are replenished,").

comma at the end.

x/x.go, line 1413 at r4 (raw file):

// RateLimiter implements a basic rate limiter.
type RateLimiter struct {

There are good libraries which exist already.

x/x.go, line 1425 at r4 (raw file):

	r := &RateLimiter{
		limiter: lockedMap{
			mp: make(map[uint64]uint64),

Use sync.Map for the map, you store the key and the pointer. And use atomics for the values.

And then also try maybe current approach. Do some parallel benchmarking, see which one is faster. In your benchmarks assume that all the namespaces that need to be accessed are already in the map.

Also benchmark with 3rd party solutions.

NamanJain8

Dismissed @manishrjain from 7 discussions.
Reviewable status: 2 of 17 files reviewed, 7 unresolved discussions (waiting on @manishrjain, @NamanJain8, @pawanrawal, and @vvbalaji-dgraph)

dgraph/cmd/bulk/loader.go, line 192 at r4 (raw file):