From 7451e885c3e91d46cc08946ca5a2cebea059fe05 Mon Sep 17 00:00:00 2001
From: Nityananda Gohain <nityanandagohain@gmail.com>
Date: Wed, 29 Nov 2023 09:10:30 +0530
Subject: [PATCH] feat: custom timeout and contextTimeout flag in response
 (#4022)

---
 pkg/query-service/app/http_handler.go    | 11 +++++++
 pkg/query-service/app/server.go          | 18 ++++++++++-
 pkg/query-service/app/server_test.go     | 41 ++++++++++++++++++++++++
 pkg/query-service/constants/constants.go | 11 +++++++
 pkg/query-service/model/v3/v3.go         |  6 ++--
 5 files changed, 84 insertions(+), 3 deletions(-)
 create mode 100644 pkg/query-service/app/server_test.go

diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index 526f12b28a..973e23328c 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -3042,6 +3042,17 @@ func (aH *APIHandler) queryRangeV3(ctx context.Context, queryRangeParams *v3.Que
 	resp := v3.QueryRangeResponse{
 		Result: result,
 	}
+
+	// This checks if the time for context to complete has exceeded.
+	// it adds flag to notify the user of incomplete respone
+	select {
+	case <-ctx.Done():
+		resp.ContextTimeout = true
+		resp.ContextTimeoutMessage = "result might contain incomplete data due to context timeout, for custom timeout set the timeout header eg:- timeout:120"
+	default:
+		break
+	}
+
 	aH.Respond(w, resp)
 }
 
diff --git a/pkg/query-service/app/server.go b/pkg/query-service/app/server.go
index 7af82b072a..e4e43f1648 100644
--- a/pkg/query-service/app/server.go
+++ b/pkg/query-service/app/server.go
@@ -428,6 +428,22 @@ func (s *Server) analyticsMiddleware(next http.Handler) http.Handler {
 	})
 }
 
+func getRouteContextTimeout(overrideTimeout string) time.Duration {
+	var timeout time.Duration
+	var err error
+	if overrideTimeout != "" {
+		timeout, err = time.ParseDuration(overrideTimeout + "s")
+		if err != nil {
+			timeout = constants.ContextTimeout
+		}
+		if timeout > constants.ContextTimeoutMaxAllowed {
+			timeout = constants.ContextTimeoutMaxAllowed
+		}
+		return timeout
+	}
+	return constants.ContextTimeout
+}
+
 func setTimeoutMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
@@ -435,7 +451,7 @@ func setTimeoutMiddleware(next http.Handler) http.Handler {
 		// check if route is not excluded
 		url := r.URL.Path
 		if _, ok := constants.TimeoutExcludedRoutes[url]; !ok {
-			ctx, cancel = context.WithTimeout(r.Context(), constants.ContextTimeout)
+			ctx, cancel = context.WithTimeout(r.Context(), getRouteContextTimeout(r.Header.Get("timeout")))
 			defer cancel()
 		}
 
diff --git a/pkg/query-service/app/server_test.go b/pkg/query-service/app/server_test.go
new file mode 100644
index 0000000000..afdc06fa33
--- /dev/null
+++ b/pkg/query-service/app/server_test.go
@@ -0,0 +1,41 @@
+package app
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetRouteContextTimeout(t *testing.T) {
+	var testGetRouteContextTimeoutData = []struct {
+		Name          string
+		OverrideValue string
+		timeout       time.Duration
+	}{
+		{
+			Name:          "default",
+			OverrideValue: "",
+			timeout:       60 * time.Second,
+		},
+		{
+			Name:          "override",
+			OverrideValue: "180",
+			timeout:       180 * time.Second,
+		},
+		{
+			Name:          "override more than max",
+			OverrideValue: "610",
+			timeout:       600 * time.Second,
+		},
+	}
+
+	t.Parallel()
+
+	for _, test := range testGetRouteContextTimeoutData {
+		t.Run(test.Name, func(t *testing.T) {
+			res := getRouteContextTimeout(test.OverrideValue)
+			assert.Equal(t, test.timeout, res)
+		})
+	}
+}
diff --git a/pkg/query-service/constants/constants.go b/pkg/query-service/constants/constants.go
index b27a85306c..7cdd502ac5 100644
--- a/pkg/query-service/constants/constants.go
+++ b/pkg/query-service/constants/constants.go
@@ -135,6 +135,17 @@ func GetContextTimeout() time.Duration {
 
 var ContextTimeout = GetContextTimeout()
 
+func GetContextTimeoutMaxAllowed() time.Duration {
+	contextTimeoutStr := GetOrDefaultEnv("CONTEXT_TIMEOUT_MAX_ALLOWED", "600")
+	contextTimeoutDuration, err := time.ParseDuration(contextTimeoutStr + "s")
+	if err != nil {
+		return time.Minute
+	}
+	return contextTimeoutDuration
+}
+
+var ContextTimeoutMaxAllowed = GetContextTimeoutMaxAllowed()
+
 const (
 	TraceID                        = "traceID"
 	ServiceName                    = "serviceName"
diff --git a/pkg/query-service/model/v3/v3.go b/pkg/query-service/model/v3/v3.go
index 09851a423d..e04be217cf 100644
--- a/pkg/query-service/model/v3/v3.go
+++ b/pkg/query-service/model/v3/v3.go
@@ -613,8 +613,10 @@ func (h *Having) CacheKey() string {
 }
 
 type QueryRangeResponse struct {
-	ResultType string    `json:"resultType"`
-	Result     []*Result `json:"result"`
+	ContextTimeout        bool      `json:"contextTimeout,omitempty"`
+	ContextTimeoutMessage string    `json:"contextTimeoutMessage,omitempty"`
+	ResultType            string    `json:"resultType"`
+	Result                []*Result `json:"result"`
 }
 
 type Result struct {