feat: custom timeout and contextTimeout flag in response (#4022)

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)
 }
 

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()
 		}
 

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)
+		})
+	}
+}

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"

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 {