fix: use common timeout middleware (#6866)

* fix: use common timeout middleware

* fix: use apiserver factory for config

* fix: add backward compatibility for old variables

* fix: remove apiserver provider and use config directly

* fix: remove apiserver interface

* fix: address comments

* fix: address minor comments

* fix: address minor comments

conf/example.yaml

@@ -68,3 +68,13 @@ sqlstore:
   sqlite:
     # The path to the SQLite database file.
     path: /var/lib/signoz/signoz.db
+
+
+##################### APIServer #####################
+apiserver:
+  timeout:
+    default: 60s
+    max: 600s
+    excluded_routes:
+      - /api/v1/logs/tail
+      - /api/v3/logs/livetail

ee/query-service/app/server.go

@@ -64,6 +64,7 @@ import (
 const AppDbEngine = "sqlite"
 
 type ServerOptions struct {
+	Config            signoz.Config
 	SigNoz            *signoz.SigNoz
 	PromConfigPath    string
 	SkipTopLvlOpsPath string
@@ -316,7 +317,11 @@ func (s *Server) createPrivateServer(apiHandler *api.APIHandler) (*http.Server,
 
 	r := baseapp.NewRouter()
 
-	r.Use(setTimeoutMiddleware)
+	r.Use(middleware.NewTimeout(zap.L(),
+		s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes,
+		s.serverOptions.Config.APIServer.Timeout.Default,
+		s.serverOptions.Config.APIServer.Timeout.Max,
+	).Wrap)
 	r.Use(s.analyticsMiddleware)
 	r.Use(middleware.NewLogging(zap.L()).Wrap)
 	r.Use(baseapp.LogCommentEnricher)
@@ -359,7 +364,11 @@ func (s *Server) createPublicServer(apiHandler *api.APIHandler, web web.Web) (*h
 	}
 	am := baseapp.NewAuthMiddleware(getUserFromRequest)
 
-	r.Use(setTimeoutMiddleware)
+	r.Use(middleware.NewTimeout(zap.L(),
+		s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes,
+		s.serverOptions.Config.APIServer.Timeout.Default,
+		s.serverOptions.Config.APIServer.Timeout.Max,
+	).Wrap)
 	r.Use(s.analyticsMiddleware)
 	r.Use(middleware.NewLogging(zap.L()).Wrap)
 	r.Use(baseapp.LogCommentEnricher)
@@ -562,23 +571,6 @@ func (s *Server) analyticsMiddleware(next http.Handler) http.Handler {
 	})
 }
 
-// TODO(remove): Implemented at pkg/http/middleware/timeout.go
-func setTimeoutMiddleware(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := r.Context()
-		var cancel context.CancelFunc
-		// check if route is not excluded
-		url := r.URL.Path
-		if _, ok := baseconst.TimeoutExcludedRoutes[url]; !ok {
-			ctx, cancel = context.WithTimeout(r.Context(), baseconst.ContextTimeout)
-			defer cancel()
-		}
-
-		r = r.WithContext(ctx)
-		next.ServeHTTP(w, r)
-	})
-}
-
 // initListeners initialises listeners of the server
 func (s *Server) initListeners() error {
 	// listen on public port

ee/query-service/main.go

@@ -151,6 +151,7 @@ func main() {
 	}
 
 	serverOptions := &app.ServerOptions{
+		Config:            config,
 		SigNoz:            signoz,
 		HTTPHostPort:      baseconst.HTTPHostPort,
 		PromConfigPath:    promConfigPath,

pkg/apiserver/config.go

@@ -0,0 +1,42 @@
+package apiserver
+
+import (
+	"time"
+
+	"go.signoz.io/signoz/pkg/factory"
+)
+
+// Config holds the configuration for config.
+type Config struct {
+	Timeout Timeout `mapstructure:"timeout"`
+}
+
+type Timeout struct {
+	// The default context timeout that can be overridden by the request
+	Default time.Duration `mapstructure:"default"`
+	// The maximum allowed context timeout
+	Max time.Duration `mapstructure:"max"`
+	// The list of routes that are excluded from the timeout
+	ExcludedRoutes []string `mapstructure:"excluded_routes"`
+}
+
+func NewConfigFactory() factory.ConfigFactory {
+	return factory.NewConfigFactory(factory.MustNewName("apiserver"), newConfig)
+}
+
+func newConfig() factory.Config {
+	return &Config{
+		Timeout: Timeout{
+			Default: 60 * time.Second,
+			Max:     600 * time.Second,
+			ExcludedRoutes: []string{
+				"/api/v1/logs/tail",
+				"/api/v3/logs/livetail",
+			},
+		},
+	}
+}
+
+func (c Config) Validate() error {
+	return nil
+}

pkg/apiserver/config_test.go

@@ -0,0 +1,51 @@
+package apiserver
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.signoz.io/signoz/pkg/config"
+	"go.signoz.io/signoz/pkg/config/envprovider"
+	"go.signoz.io/signoz/pkg/factory"
+)
+
+func TestNewWithEnvProvider(t *testing.T) {
+	t.Setenv("SIGNOZ_APISERVER_TIMEOUT_DEFAULT", "70s")
+	t.Setenv("SIGNOZ_APISERVER_TIMEOUT_MAX", "700s")
+	t.Setenv("SIGNOZ_APISERVER_TIMEOUT_EXCLUDED__ROUTES", "/excluded1,/excluded2")
+
+	conf, err := config.New(
+		context.Background(),
+		config.ResolverConfig{
+			Uris: []string{"env:"},
+			ProviderFactories: []config.ProviderFactory{
+				envprovider.NewFactory(),
+			},
+		},
+		[]factory.ConfigFactory{
+			NewConfigFactory(),
+		},
+	)
+	require.NoError(t, err)
+
+	actual := &Config{}
+	err = conf.Unmarshal("apiserver", actual)
+
+	require.NoError(t, err)
+
+	expected := &Config{
+		Timeout: Timeout{
+			Default: 70 * time.Second,
+			Max:     700 * time.Second,
+			ExcludedRoutes: []string{
+				"/excluded1",
+				"/excluded2",
+			},
+		},
+	}
+
+	assert.Equal(t, expected, actual)
+}

pkg/http/middleware/timeout.go

@@ -22,13 +22,14 @@ type Timeout struct {
 	maxTimeout time.Duration
 }
 
-func NewTimeout(logger *zap.Logger, excluded map[string]struct{}, defaultTimeout time.Duration, maxTimeout time.Duration) *Timeout {
+func NewTimeout(logger *zap.Logger, excludedRoutes []string, defaultTimeout time.Duration, maxTimeout time.Duration) *Timeout {
 	if logger == nil {
 		panic("cannot build timeout, logger is empty")
 	}
 
-	if excluded == nil {
+	excluded := make(map[string]struct{}, len(excludedRoutes))
-		excluded = make(map[string]struct{})
+	for _, route := range excludedRoutes {
+		excluded[route] = struct{}{}
 	}
 
 	if defaultTimeout.Seconds() == 0 {

pkg/http/middleware/timeout_test.go

@@ -16,7 +16,7 @@ func TestTimeout(t *testing.T) {
 	writeTimeout := 6 * time.Second
 	defaultTimeout := 2 * time.Second
 	maxTimeout := 4 * time.Second
-	m := NewTimeout(zap.NewNop(), map[string]struct{}{"/excluded": {}}, defaultTimeout, maxTimeout)
+	m := NewTimeout(zap.NewNop(), []string{"/excluded"}, defaultTimeout, maxTimeout)
 
 	listener, err := net.Listen("tcp", "localhost:0")
 	require.NoError(t, err)

pkg/query-service/app/server.go

@@ -56,6 +56,7 @@ import (
 )
 
 type ServerOptions struct {
+	Config            signoz.Config
 	PromConfigPath    string
 	SkipTopLvlOpsPath string
 	HTTPHostPort      string
@@ -262,7 +263,11 @@ func (s *Server) createPrivateServer(api *APIHandler) (*http.Server, error) {
 
 	r := NewRouter()
 
-	r.Use(setTimeoutMiddleware)
+	r.Use(middleware.NewTimeout(zap.L(),
+		s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes,
+		s.serverOptions.Config.APIServer.Timeout.Default,
+		s.serverOptions.Config.APIServer.Timeout.Max,
+	).Wrap)
 	r.Use(s.analyticsMiddleware)
 	r.Use(middleware.NewLogging(zap.L()).Wrap)
 
@@ -288,7 +293,11 @@ func (s *Server) createPublicServer(api *APIHandler, web web.Web) (*http.Server,
 
 	r := NewRouter()
 
-	r.Use(setTimeoutMiddleware)
+	r.Use(middleware.NewTimeout(zap.L(),
+		s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes,
+		s.serverOptions.Config.APIServer.Timeout.Default,
+		s.serverOptions.Config.APIServer.Timeout.Max,
+	).Wrap)
 	r.Use(s.analyticsMiddleware)
 	r.Use(middleware.NewLogging(zap.L()).Wrap)
 	r.Use(LogCommentEnricher)
@@ -571,40 +580,6 @@ func (s *Server) analyticsMiddleware(next http.Handler) http.Handler {
 	})
 }
 
-// TODO(remove): Implemented at pkg/http/middleware/timeout.go
-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
-}
-
-// TODO(remove): Implemented at pkg/http/middleware/timeout.go
-func setTimeoutMiddleware(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := r.Context()
-		var cancel context.CancelFunc
-		// check if route is not excluded
-		url := r.URL.Path
-		if _, ok := constants.TimeoutExcludedRoutes[url]; !ok {
-			ctx, cancel = context.WithTimeout(r.Context(), getRouteContextTimeout(r.Header.Get("timeout")))
-			defer cancel()
-		}
-
-		r = r.WithContext(ctx)
-		next.ServeHTTP(w, r)
-	})
-}
-
 // initListeners initialises listeners of the server
 func (s *Server) initListeners() error {
 	// listen on public port

pkg/query-service/app/server_test.go

@@ -1,42 +0,0 @@
-package app
-
-import (
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
-)
-
-// TODO(remove): Implemented at pkg/http/middleware/timeout_test.go
-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

@@ -153,26 +153,6 @@ var DEFAULT_FEATURE_SET = model.FeatureSet{
 	},
 }
 
-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()
-
-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
-}
-
 func GetEvalDelay() time.Duration {
 	evalDelayStr := GetOrDefaultEnv("RULES_EVAL_DELAY", "2m")
 	evalDelayDuration, err := time.ParseDuration(evalDelayStr)
@@ -182,8 +162,6 @@ func GetEvalDelay() time.Duration {
 	return evalDelayDuration
 }
 
-var ContextTimeoutMaxAllowed = GetContextTimeoutMaxAllowed()
-
 const (
 	TraceID                        = "traceID"
 	ServiceName                    = "serviceName"
@@ -255,11 +233,6 @@ const (
 	SIGNOZ_TOP_LEVEL_OPERATIONS_TABLENAME      = "distributed_top_level_operations"
 )
 
-var TimeoutExcludedRoutes = map[string]bool{
-	"/api/v1/logs/tail":     true,
-	"/api/v3/logs/livetail": true,
-}
-
 // alert related constants
 const (
 	// AlertHelpPage is used in case default alert repo url is not set

pkg/query-service/constants/constants_test.go

@@ -3,7 +3,6 @@ package constants
 import (
 	"os"
 	"testing"
-	"time"
 
 	. "github.com/smartystreets/goconvey/convey"
 )
@@ -20,16 +19,3 @@ 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))
-		})
-	})
-}

pkg/query-service/main.go

@@ -95,6 +95,7 @@ func main() {
 	}
 
 	serverOptions := &app.ServerOptions{
+		Config:            config,
 		HTTPHostPort:      constants.HTTPHostPort,
 		PromConfigPath:    promConfigPath,
 		SkipTopLvlOpsPath: skipTopLvlOpsPath,

pkg/signoz/config.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"time"
 
+	"go.signoz.io/signoz/pkg/apiserver"
 	"go.signoz.io/signoz/pkg/cache"
 	"go.signoz.io/signoz/pkg/config"
 	"go.signoz.io/signoz/pkg/factory"
@@ -30,6 +32,9 @@ type Config struct {
 
 	// SQLMigrator config
 	SQLMigrator sqlmigrator.Config `mapstructure:"sqlmigrator"`
+
+	// API Server config
+	APIServer apiserver.Config `mapstructure:"apiserver"`
 }
 
 func NewConfig(ctx context.Context, resolverConfig config.ResolverConfig) (Config, error) {
@@ -39,6 +44,7 @@ func NewConfig(ctx context.Context, resolverConfig config.ResolverConfig) (Confi
 		cache.NewConfigFactory(),
 		sqlstore.NewConfigFactory(),
 		sqlmigrator.NewConfigFactory(),
+		apiserver.NewConfigFactory(),
 	}
 
 	conf, err := config.New(ctx, resolverConfig, configFactories)
@@ -62,4 +68,23 @@ func mergeAndEnsureBackwardCompatibility(config *Config) {
 		fmt.Println("[Deprecated] env SIGNOZ_LOCAL_DB_PATH is deprecated and scheduled for removal. Please use SIGNOZ_SQLSTORE_SQLITE_PATH instead.")
 		config.SQLStore.Sqlite.Path = os.Getenv("SIGNOZ_LOCAL_DB_PATH")
 	}
+	if os.Getenv("CONTEXT_TIMEOUT") != "" {
+		fmt.Println("[Deprecated] env CONTEXT_TIMEOUT is deprecated and scheduled for removal. Please use SIGNOZ_APISERVER_TIMEOUT_DEFAULT instead.")
+		contextTimeoutDuration, err := time.ParseDuration(os.Getenv("CONTEXT_TIMEOUT") + "s")
+		if err == nil {
+			config.APIServer.Timeout.Default = contextTimeoutDuration
+		} else {
+			fmt.Println("Error parsing CONTEXT_TIMEOUT, using default value of 60s")
+		}
+	}
+	if os.Getenv("CONTEXT_TIMEOUT_MAX_ALLOWED") != "" {
+		fmt.Println("[Deprecated] env CONTEXT_TIMEOUT_MAX_ALLOWED is deprecated and scheduled for removal. Please use SIGNOZ_APISERVER_TIMEOUT_MAX instead.")
+
+		contextTimeoutDuration, err := time.ParseDuration(os.Getenv("CONTEXT_TIMEOUT_MAX_ALLOWED") + "s")
+		if err == nil {
+			config.APIServer.Timeout.Max = contextTimeoutDuration
+		} else {
+			fmt.Println("Error parsing CONTEXT_TIMEOUT_MAX_ALLOWED, using default value of 600s")
+		}
+	}
 }