Add mechanism to have a limit on number of pending queries

dgraph/cmd/alpha/run.go

@@ -204,6 +204,8 @@ they form a Raft group and provide synchronous replication.
 		Flag("disallow-drop",
 			"Set disallow-drop to true to block drop-all and drop-data operation. It still"+
 				" allows dropping attributes and types.").
+		Flag("max-pending-queries",
+			"Number of maximum pending queries before we reject them as too many requests.").
 		String())
 
 	flag.String("ludicrous", worker.LudicrousDefaults, z.NewSuperFlagHelp(worker.LudicrousDefaults).
@@ -737,6 +739,7 @@ func run() {
 			return
 		}
 	}
+	edgraph.Init()
 
 	x.PrintVersion()
 	glog.Infof("x.Config: %+v", x.Config)

dgraph/cmd/increment/increment.go

@@ -23,6 +23,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/dgraph-io/dgo/v200"
@@ -56,9 +57,10 @@ func init() {
 
 	flag.String("cloud", "", "addr: xxx; jwt: xxx")
 	flag.String("alpha", "localhost:9080", "Address of Dgraph Alpha.")
-	flag.Int("num", 1, "How many times to run.")
+	flag.Int("num", 1, "How many times to run per goroutine.")
 	flag.Int("retries", 10, "How many times to retry setting up the connection.")
 	flag.Duration("wait", 0*time.Second, "How long to wait.")
+	flag.Int("conc", 1, "How many goroutines to run.")
 
 	flag.String("creds", "",
 		`Various login credentials if login is required.
@@ -178,6 +180,7 @@ func run(conf *viper.Viper) {
 
 	waitDur := conf.GetDuration("wait")
 	num := conf.GetInt("num")
+	conc := int(conf.GetInt("conc"))
 	format := "0102 03:04:05.999"
 
 	// Do a sanity check on the passed credentials.
@@ -196,20 +199,40 @@ func run(conf *viper.Viper) {
 		dg = dgTmp
 	}
 
-	for num > 0 {
-		txnStart := time.Now() // Start time of transaction
-		cnt, err := process(dg, conf)
-		now := time.Now().UTC().Format(format)
-		if err != nil {
-			fmt.Printf("%-17s While trying to process counter: %v. Retrying...\n", now, err)
-			time.Sleep(time.Second)
-			continue
-		}
-		serverLat := cnt.qLatency + cnt.mLatency
-		clientLat := time.Since(txnStart).Round(time.Millisecond)
-		fmt.Printf("%-17s Counter VAL: %d   [ Ts: %d ] Latency: Q %s M %s S %s C %s D %s\n", now, cnt.Val,
-			cnt.startTs, cnt.qLatency, cnt.mLatency, serverLat, clientLat, clientLat-serverLat)
+	// Run things serially first.
+	for i := 0; i < conc; i++ {
+		_, err := process(dg, conf)
+		x.Check(err)
 		num--
-		time.Sleep(waitDur)
 	}
+
+	var wg sync.WaitGroup
+	f := func(i int) {
+		defer wg.Done()
+		count := 0
+		for count < num {
+			txnStart := time.Now() // Start time of transaction
+			cnt, err := process(dg, conf)
+			now := time.Now().UTC().Format(format)
+			if err != nil {
+				fmt.Printf("%-17s While trying to process counter: %v. Retrying...\n", now, err)
+				time.Sleep(time.Second)
+				continue
+			}
+			serverLat := cnt.qLatency + cnt.mLatency
+			clientLat := time.Since(txnStart).Round(time.Millisecond)
+			fmt.Printf(
+				"[%d] %-17s Counter VAL: %d   [ Ts: %d ] Latency: Q %s M %s S %s C %s D %s\n",
+				i, now, cnt.Val, cnt.startTs, cnt.qLatency, cnt.mLatency,
+				serverLat, clientLat, clientLat-serverLat)
+			time.Sleep(waitDur)
+			count++
+		}
+	}
+
+	for i := 0; i < conc; i++ {
+		wg.Add(1)
+		go f(i)
+	}
+	wg.Wait()
 }

edgraph/server.go

@@ -1117,22 +1117,35 @@ func (s *Server) Query(ctx context.Context, req *api.Request) (*api.Response, er
 	return s.doQuery(ctx, &Request{req: req, doAuth: getAuthMode(ctx)})
 }
 
+var pendingQueries int64
+var maxPendingQueries int64
+var serverOverloadErr = errors.New("429 Too Many Requests. Please throttle your requests")
+
+func Init() {
+	maxPendingQueries = x.Config.Limit.GetInt64("max-pending-queries")
+}
+
 func (s *Server) doQuery(ctx context.Context, req *Request) (
 	resp *api.Response, rerr error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+	defer atomic.AddInt64(&pendingQueries, -1)
+	if val := atomic.AddInt64(&pendingQueries, 1); val > maxPendingQueries {
+		return nil, serverOverloadErr
+	}
+
 	if bool(glog.V(3)) || worker.LogRequestEnabled() {
 		glog.Infof("Got a query: %+v", req.req)
 	}
+
 	isGraphQL, _ := ctx.Value(IsGraphql).(bool)
 	if isGraphQL {
 		atomic.AddUint64(&numGraphQL, 1)
 	} else {
 		atomic.AddUint64(&numGraphQLPM, 1)
 	}
 
-	if ctx.Err() != nil {
-		return nil, ctx.Err()
-	}
-
 	l := &query.Latency{}
 	l.Start = time.Now()
 

worker/server_state.go

@@ -47,7 +47,7 @@ const (
 	CDCDefaults       = `file=; kafka=; sasl_user=; sasl_password=; ca_cert=; client_cert=; ` +
 		`client_key=;`
 	LimitDefaults = `mutations=allow; query-edge=1000000; normalize-node=10000; ` +
-		`mutations-nquad=1000000; disallow-drop=false;`
+		`mutations-nquad=1000000; disallow-drop=false; max-pending-queries=10000`
 	GraphQLDefaults = `introspection=true; debug=false; extensions=true; poll-interval=1s; ` +
 		`lambda-url=;`
 )

x/config.go

@@ -36,10 +36,10 @@ type Options struct {
 	//                      normalize directive
 	// mutations-nquad int - maximum number of nquads that can be inserted in a mutation request
 	// BlockDropAll bool - if set to true, the drop all operation will be rejected by the server.
-	Limit               *z.SuperFlag
-	LimitMutationsNquad int
-	LimitQueryEdge      uint64
-	BlockClusterWideDrop        bool
+	Limit                *z.SuperFlag
+	LimitMutationsNquad  int
+	LimitQueryEdge       uint64
+	BlockClusterWideDrop bool
 
 	// GraphQL options:
 	//