From af9a8b99079f1f047f3e1b715097d78883b34af6 Mon Sep 17 00:00:00 2001
From: Aman Mangal <aman@dgraph.io>
Date: Mon, 1 Apr 2019 21:04:40 +0530
Subject: [PATCH] Add timeout option while running queries

The endpoint /query now accepts a timeout key /query?timeout=100ms
If the query doesn't finish within the timeout time, query is
stopped and an error is returned.
---
 dgraph/cmd/alpha/http.go      | 14 ++++++++++++++
 dgraph/cmd/alpha/http_test.go | 36 ++++++++++++++++++++++++++++++-----
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/dgraph/cmd/alpha/http.go b/dgraph/cmd/alpha/http.go
index f804b159290..9c386b72584 100644
--- a/dgraph/cmd/alpha/http.go
+++ b/dgraph/cmd/alpha/http.go
@@ -159,6 +159,20 @@ func queryHandler(w http.ResponseWriter, r *http.Request) {
 	ctx := context.WithValue(context.Background(), query.DebugKey, d)
 	ctx = attachAccessJwt(ctx, r)
 
+	// Timeout is expected to be in millisecond
+	paramTimeout := r.URL.Query().Get("timeout")
+	if paramTimeout != "" {
+		timeout, err := time.ParseDuration(paramTimeout)
+		if err != nil {
+			x.SetStatusWithData(w, x.Error, err.Error())
+			return
+		}
+
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
+
 	if req.StartTs == 0 {
 		// If be is set, run this as a best-effort query.
 		be, _ := strconv.ParseBool(r.URL.Query().Get("be"))
diff --git a/dgraph/cmd/alpha/http_test.go b/dgraph/cmd/alpha/http_test.go
index f5b71387438..00b7815ec14 100644
--- a/dgraph/cmd/alpha/http_test.go
+++ b/dgraph/cmd/alpha/http_test.go
@@ -37,13 +37,24 @@ import (
 	"github.com/dgraph-io/dgraph/z"
 )
 
+type respError struct {
+	Code    string `json:"code"`
+	Message string `json:"message"`
+}
+
 type res struct {
 	Data       json.RawMessage   `json:"data"`
 	Extensions *query.Extensions `json:"extensions,omitempty"`
+	Errors     []respError       `json:"errors,omitempty"`
 }
 
-func queryWithGz(q string, gzReq bool, gzResp bool) (string, *http.Response, error) {
+func queryWithGz(q string, gzReq bool, gzResp bool, timeout string) (
+	string, *http.Response, error) {
+
 	url := addr + "/query"
+	if timeout != "" {
+		url = url + fmt.Sprintf("?timeout=%v", timeout)
+	}
 
 	var buf *bytes.Buffer
 	if gzReq {
@@ -99,6 +110,11 @@ func queryWithGz(q string, gzReq bool, gzResp bool) (string, *http.Response, err
 	var r res
 	x.Check(json.Unmarshal(body, &r))
 
+	// Check for errors
+	if len(r.Errors) != 0 {
+		return "", nil, errors.New(r.Errors[0].Message)
+	}
+
 	// Remove the extensions.
 	r2 := res{
 		Data: r.Data,
@@ -463,23 +479,33 @@ func TestHttpCompressionSupport(t *testing.T) {
 	err := runMutation(m1)
 	require.NoError(t, err)
 
-	data, resp, err := queryWithGz(q1, false, false)
+	data, resp, err := queryWithGz(q1, false, false, "")
 	require.NoError(t, err)
 	require.Equal(t, r1, data)
 	require.Empty(t, resp.Header.Get("Content-Encoding"))
 
-	data, resp, err = queryWithGz(q1, false, true)
+	data, resp, err = queryWithGz(q1, false, true, "")
 	require.NoError(t, err)
 	require.Equal(t, r1, data)
 	require.Equal(t, "gzip", resp.Header.Get("Content-Encoding"))
 
-	data, resp, err = queryWithGz(q1, true, false)
+	data, resp, err = queryWithGz(q1, true, false, "")
 	require.NoError(t, err)
 	require.Equal(t, r1, data)
 	require.Empty(t, resp.Header.Get("Content-Encoding"))
 
-	data, resp, err = queryWithGz(q1, true, true)
+	data, resp, err = queryWithGz(q1, true, true, "")
 	require.NoError(t, err)
 	require.Equal(t, r1, data)
 	require.Equal(t, "gzip", resp.Header.Get("Content-Encoding"))
+
+	// query with timeout
+	data, resp, err = queryWithGz(q1, false, false, "1ms")
+	require.EqualError(t, err, ": context deadline exceeded")
+	require.Equal(t, "", data)
+
+	data, resp, err = queryWithGz(q1, false, false, "1s")
+	require.NoError(t, err)
+	require.Equal(t, r1, data)
+	require.Empty(t, resp.Header.Get("Content-Encoding"))
 }