Add timeout option while running queries over HTTP

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.

(cherry picked from commit 913c920)
- Renamed commit message to clarify this for queries over HTTP.

dgraph/cmd/alpha/http.go

@@ -157,6 +157,20 @@ func queryHandler(w http.ResponseWriter, r *http.Request) {
 	d := r.URL.Query().Get("debug")
 	ctx := context.WithValue(context.Background(), query.DebugKey, d)
 
+	// 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"))

dgraph/cmd/alpha/http_test.go

@@ -35,15 +35,26 @@ import (
 	"github.com/dgraph-io/dgraph/x"
 )
 
+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"`
 }
 
 var addr = "http://localhost:8180"
 
-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,
@@ -308,23 +324,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"))
 }