feat(flags): Add query timeout as a limit config (#7599)

Fixes DGRAPH-3184

dgraph/cmd/alpha/run.go

@@ -204,6 +204,9 @@ 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("query-timeout",
+			"Maximum time after which a query execution will fail. If set to"+
+				" 0, the timeout is infinite.").
 		Flag("max-pending-queries",
 			"Number of maximum pending queries before we reject them as too many requests.").
 		String())
@@ -733,6 +736,7 @@ func run() {
 	x.Config.LimitMutationsNquad = int(x.Config.Limit.GetInt64("mutations-nquad"))
 	x.Config.LimitQueryEdge = x.Config.Limit.GetUint64("query-edge")
 	x.Config.BlockClusterWideDrop = x.Config.Limit.GetBool("disallow-drop")
+	x.Config.QueryTimeout = x.Config.Limit.GetDuration("query-timeout")
 
 	x.Config.GraphQL = z.NewSuperFlag(Alpha.Conf.GetString("graphql")).MergeAndCheckDefault(
 		worker.GraphQLDefaults)

edgraph/server.go

@@ -1114,6 +1114,16 @@ func (s *Server) QueryGraphQL(ctx context.Context, req *api.Request,
 // Query handles queries or mutations
 func (s *Server) Query(ctx context.Context, req *api.Request) (*api.Response, error) {
 	ctx = x.AttachJWTNamespace(ctx)
+	// Add a timeout for queries which don't have a deadline set. We don't want to
+	// apply a timeout if it's a mutation, that's currently handled by flag
+	// "abort_older_than".
+	if req.GetMutations() == nil && x.Config.QueryTimeout != 0 {
+		if d, _ := ctx.Deadline(); d.IsZero() {
+			var cancel context.CancelFunc
+			ctx, cancel = context.WithTimeout(ctx, x.Config.QueryTimeout)
+			defer cancel()
+		}
+	}
 	return s.doQuery(ctx, &Request{req: req, doAuth: getAuthMode(ctx)})
 }
 

worker/server_state.go

@@ -47,7 +47,8 @@ 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; max-pending-queries=10000`
+		`mutations-nquad=1000000; disallow-drop=false; max-pending-queries=10000; query-timeout
+		=0ms;`
 	GraphQLDefaults = `introspection=true; debug=false; extensions=true; poll-interval=1s; ` +
 		`lambda-url=;`
 )

x/config.go

@@ -36,10 +36,12 @@ 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.
+	// query-timeout duration - Maximum time after which a query execution will fail.
 	Limit                *z.SuperFlag
 	LimitMutationsNquad  int
 	LimitQueryEdge       uint64
 	BlockClusterWideDrop bool
+	QueryTimeout         time.Duration
 
 	// GraphQL options:
 	//