make timeouts configurable

Author: jimdaga

api/probes/v1/http.go

@@ -19,11 +19,14 @@ import (
 )
 
 type handlerOptions struct {
-	logger           log.Logger
-	tenantHeader     string
-	tlsOptions       *tls.UpstreamOptions
-	readMiddlewares  []func(http.Handler) http.Handler
-	writeMiddlewares []func(http.Handler) http.Handler
+	logger              log.Logger
+	tenantHeader        string
+	tlsOptions          *tls.UpstreamOptions
+	dialTimeout         time.Duration
+	keepAliveTimeout    time.Duration
+	tlsHandshakeTimeout time.Duration
+	readMiddlewares     []func(http.Handler) http.Handler
+	writeMiddlewares    []func(http.Handler) http.Handler
 }
 
 // HandlerOption is a function that configures the handler.
@@ -50,6 +53,27 @@ func WithUpstreamTLSOptions(opts *tls.UpstreamOptions) HandlerOption {
 	}
 }
 
+// WithDialTimeout sets the dial timeout for upstream connections.
+func WithDialTimeout(timeout time.Duration) HandlerOption {
+	return func(o *handlerOptions) {
+		o.dialTimeout = timeout
+	}
+}
+
+// WithKeepAliveTimeout sets the keep-alive timeout for upstream connections.
+func WithKeepAliveTimeout(timeout time.Duration) HandlerOption {
+	return func(o *handlerOptions) {
+		o.keepAliveTimeout = timeout
+	}
+}
+
+// WithTLSHandshakeTimeout sets the TLS handshake timeout for upstream connections.
+func WithTLSHandshakeTimeout(timeout time.Duration) HandlerOption {
+	return func(o *handlerOptions) {
+		o.tlsHandshakeTimeout = timeout
+	}
+}
+
 // WithReadMiddleware adds a middleware for read operations.
 func WithReadMiddleware(m func(http.Handler) http.Handler) HandlerOption {
 	return func(o *handlerOptions) {
@@ -67,7 +91,10 @@ func WithWriteMiddleware(m func(http.Handler) http.Handler) HandlerOption {
 // NewHandler creates a new handler for the probes API.
 func NewHandler(downstream *url.URL, opts ...HandlerOption) (http.Handler, error) {
 	options := &handlerOptions{
-		logger: log.NewNopLogger(),
+		logger:              log.NewNopLogger(),
+		dialTimeout:         30 * time.Second,
+		keepAliveTimeout:    30 * time.Second,
+		tlsHandshakeTimeout: 10 * time.Second,
 	}
 	for _, o := range opts {
 		o(options)
@@ -79,10 +106,10 @@ func NewHandler(downstream *url.URL, opts ...HandlerOption) (http.Handler, error
 	proxy.Transport = &http.Transport{
 		Proxy: http.ProxyFromEnvironment,
 		DialContext: (&net.Dialer{
-			Timeout:   30 * time.Second,
-			KeepAlive: 30 * time.Second,
+			Timeout:   options.dialTimeout,
+			KeepAlive: options.keepAliveTimeout,
 		}).DialContext,
-		TLSHandshakeTimeout: 10 * time.Second,
+		TLSHandshakeTimeout: options.tlsHandshakeTimeout,
 		TLSClientConfig:     options.tlsOptions.NewClientConfig(),
 	}
 

main.go

@@ -200,6 +200,9 @@ type tracesConfig struct {
 type probesConfig struct {
 	endpoint             *url.URL
 	upstreamWriteTimeout time.Duration
+	dialTimeout          time.Duration
+	keepAliveTimeout     time.Duration
+	tlsHandshakeTimeout  time.Duration
 	upstreamCAFile       string
 	upstreamCertFile     string
 	upstreamKeyFile      string
@@ -710,6 +713,9 @@ func main() {
 							probesv1.WithLogger(logger),
 							probesv1.WithUpstreamTLSOptions(probesUpstreamClientOptions),
 							probesv1.WithTenantHeader(cfg.probes.tenantHeader),
+							probesv1.WithDialTimeout(cfg.probes.dialTimeout),
+							probesv1.WithKeepAliveTimeout(cfg.probes.keepAliveTimeout),
+							probesv1.WithTLSHandshakeTimeout(cfg.probes.tlsHandshakeTimeout),
 							probesv1.WithReadMiddleware(authentication.WithTenantMiddlewares(pm.Middlewares)),
 							probesv1.WithReadMiddleware(rateLimitMiddleware),
 							probesv1.WithReadMiddleware(authorization.WithAuthorizers(authorizers, rbac.Read, "probes")),
@@ -1240,6 +1246,12 @@ func parseFlags() (config, error) {
 		"Watch for certificate changes and reload")
 	flag.StringVar(&cfg.probes.tenantHeader, "probes.tenant-header", "X-Tenant",
 		"The name of the HTTP header containing the tenant ID to forward to the probes upstream.")
+	flag.DurationVar(&cfg.probes.dialTimeout, "probes.dial-timeout", 30*time.Second,
+		"The timeout for establishing connections to the probes upstream.")
+	flag.DurationVar(&cfg.probes.keepAliveTimeout, "probes.keep-alive-timeout", 30*time.Second,
+		"The keep-alive timeout for connections to the probes upstream.")
+	flag.DurationVar(&cfg.probes.tlsHandshakeTimeout, "probes.tls-handshake-timeout", 10*time.Second,
+		"The TLS handshake timeout for connections to the probes upstream.")
 	flag.StringVar(&cfg.tls.serverCertFile, "tls.server.cert-file", "",
 		"File containing the default x509 Certificate for HTTPS. Leave blank to disable TLS.")
 	flag.StringVar(&cfg.tls.serverKeyFile, "tls.server.key-file", "",