From dcad77746ae68d411319bdd60d29e1f1b4cddff1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl?=
 <130993189+daniel-t4e@users.noreply.github.com>
Date: Tue, 25 Apr 2023 12:59:49 +0200
Subject: [PATCH] feat: added configuration via env for context timeout (#2585)

---
 ee/query-service/app/server.go                |  2 +-
 pkg/query-service/app/server.go               |  2 +-
 pkg/query-service/constants/constants.go      | 13 ++++++++++++-
 pkg/query-service/constants/constants_test.go | 14 ++++++++++++++
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/ee/query-service/app/server.go b/ee/query-service/app/server.go
index 72091cc1e3..30daa090b6 100644
--- a/ee/query-service/app/server.go
+++ b/ee/query-service/app/server.go
@@ -402,7 +402,7 @@ func setTimeoutMiddleware(next http.Handler) http.Handler {
 		// check if route is not excluded
 		url := r.URL.Path
 		if _, ok := baseconst.TimeoutExcludedRoutes[url]; !ok {
-			ctx, cancel = context.WithTimeout(r.Context(), baseconst.ContextTimeout*time.Second)
+			ctx, cancel = context.WithTimeout(r.Context(), baseconst.ContextTimeout)
 			defer cancel()
 		}
 
diff --git a/pkg/query-service/app/server.go b/pkg/query-service/app/server.go
index c3142643f0..73b1baed1e 100644
--- a/pkg/query-service/app/server.go
+++ b/pkg/query-service/app/server.go
@@ -349,7 +349,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*time.Second)
+			ctx, cancel = context.WithTimeout(r.Context(), constants.ContextTimeout)
 			defer cancel()
 		}
 
diff --git a/pkg/query-service/constants/constants.go b/pkg/query-service/constants/constants.go
index d0ded00bbe..3d00c25e78 100644
--- a/pkg/query-service/constants/constants.go
+++ b/pkg/query-service/constants/constants.go
@@ -3,6 +3,7 @@ package constants
 import (
 	"os"
 	"strconv"
+	"time"
 
 	"go.signoz.io/signoz/pkg/query-service/model"
 	v3 "go.signoz.io/signoz/pkg/query-service/model/v3"
@@ -74,6 +75,17 @@ var DEFAULT_FEATURE_SET = model.FeatureSet{
 	TimestampSort: IsTimestampSortFeatureEnabled(),
 }
 
+func GetContextTimeout() time.Duration {
+	contextTimeoutStr := GetOrDefaultEnv("CONTEXT_TIMEOUT", "60")
+	contextTimeoutDuration, err := time.ParseDuration(contextTimeoutStr + "s")
+	if err != nil {
+		return time.Minute
+	}
+	return contextTimeoutDuration
+}
+
+var ContextTimeout = GetContextTimeout()
+
 const (
 	TraceID                        = "traceID"
 	ServiceName                    = "serviceName"
@@ -97,7 +109,6 @@ const (
 	ResponseStatusCode             = "responseStatusCode"
 	Descending                     = "descending"
 	Ascending                      = "ascending"
-	ContextTimeout                 = 60 // seconds
 	StatusPending                  = "pending"
 	StatusFailed                   = "failed"
 	StatusSuccess                  = "success"
diff --git a/pkg/query-service/constants/constants_test.go b/pkg/query-service/constants/constants_test.go
index 59c35eae12..1edf07d219 100644
--- a/pkg/query-service/constants/constants_test.go
+++ b/pkg/query-service/constants/constants_test.go
@@ -3,6 +3,7 @@ package constants
 import (
 	"os"
 	"testing"
+	"time"
 
 	. "github.com/smartystreets/goconvey/convey"
 )
@@ -19,3 +20,16 @@ func TestGetAlertManagerApiPrefix(t *testing.T) {
 		})
 	})
 }
+
+func TestGetContextTimeout(t *testing.T) {
+	Convey("TestGetContextTimeout", t, func() {
+		res := GetContextTimeout()
+		So(res, ShouldEqual, time.Duration(60000000000))
+
+		Convey("WithEnvSet", func() {
+			os.Setenv("CONTEXT_TIMEOUT", "120")
+			res = GetContextTimeout()
+			So(res, ShouldEqual, time.Duration(120000000000))
+		})
+	})
+}