fix: Add grpc-connect-timeout parameter

This controls the timeout for the GRPC connection. It defaults to 0,
which means there's no timeout.

Signed-off-by: Marcelo E. Magallon <marcelo.magallon@grafana.com>

Author: mem

cmd/synthetic-monitoring-agent/main.go

@@ -67,6 +67,7 @@ func run(args []string, stdout io.Writer) error {
 			ReportVersion         bool
 			GrpcApiServerAddr     string
 			GrpcInsecure          bool
+			GrpcConnectTimeout    time.Duration
 			ApiToken              Secret
 			EnableChangeLogLevel  bool
 			EnableDisconnect      bool
@@ -106,6 +107,7 @@ func run(args []string, stdout io.Writer) error {
 	flags.BoolVar(&config.ReportVersion, "version", config.ReportVersion, "report version and exit")
 	flags.StringVar(&config.GrpcApiServerAddr, "api-server-address", config.GrpcApiServerAddr, "GRPC API server address")
 	flags.BoolVar(&config.GrpcInsecure, "api-insecure", config.GrpcInsecure, "Don't use TLS with connections to GRPC API")
+	flags.DurationVar(&config.GrpcConnectTimeout, "grpc-connect-timeout", config.GrpcConnectTimeout, "timeout for initial GRPC connection (0 = no timeout)")
 	flags.Var(&config.ApiToken, "api-token", `synthetic monitoring probe authentication token (default "")`)
 	flags.BoolVar(&config.EnableChangeLogLevel, "enable-change-log-level", config.EnableChangeLogLevel, "enable changing the log level at runtime")
 	flags.BoolVar(&config.EnableDisconnect, "enable-disconnect", config.EnableDisconnect, "enable HTTP /disconnect endpoint")
@@ -385,6 +387,7 @@ func run(args []string, stdout io.Writer) error {
 		UsageReporter:           usageReporter,
 		CostAttributionLabels:   cals,
 		SupportsProtocolSecrets: config.EnableProtocolSecrets,
+		GrpcConnectTimeout:      config.GrpcConnectTimeout,
 	})
 
 	if err != nil {
@@ -406,6 +409,7 @@ func run(args []string, stdout io.Writer) error {
 		K6Runner:                k6Runner,
 		SecretProvider:          secretProvider,
 		SupportsProtocolSecrets: config.EnableProtocolSecrets,
+		GrpcConnectTimeout:      config.GrpcConnectTimeout,
 	})
 	if err != nil {
 		return fmt.Errorf("cannot create ad-hoc checks handler: %w", err)

internal/adhoc/adhoc.go

@@ -46,6 +46,7 @@ type Handler struct {
 	grpcAdhocChecksClientFactory func(conn ClientConn) (sm.AdHocChecksClient, error)
 	proberFactory                prober.ProberFactory
 	supportsProtocolSecrets      bool
+	connectTimeout               time.Duration
 }
 
 // Error represents errors returned from this package.
@@ -120,6 +121,7 @@ type HandlerOpts struct {
 	K6Runner                k6runner.Runner
 	SecretProvider          secrets.SecretProvider
 	SupportsProtocolSecrets bool
+	GrpcConnectTimeout      time.Duration
 
 	// these two fields exists so that tests can pass alternate
 	// implementations, they are unexported so that clients of this
@@ -154,6 +156,7 @@ func NewHandler(opts HandlerOpts) (*Handler, error) {
 		grpcAdhocChecksClientFactory: opts.grpcAdhocChecksClientFactory,
 		proberFactory:                prober.NewProberFactory(opts.K6Runner, 0, opts.Features, opts.SecretProvider),
 		supportsProtocolSecrets:      opts.SupportsProtocolSecrets,
+		connectTimeout:               opts.GrpcConnectTimeout,
 		api: apiInfo{
 			conn: opts.Conn,
 		},
@@ -352,15 +355,28 @@ func (h *Handler) loop(ctx context.Context) error {
 		}
 	}
 
+	// Create a timeout context for RegisterProbe if configured.
+	// This ensures we don't wait indefinitely if the server is unreachable.
+	registerCtx := ctx
+	var cancel context.CancelFunc
+
+	if h.connectTimeout > 0 {
+		registerCtx, cancel = context.WithTimeout(ctx, h.connectTimeout)
+		defer cancel()
+		h.logger.Info().
+			Dur("timeout", h.connectTimeout).
+			Msg("using explicit connection timeout for RegisterProbe")
+	}
+
 	result, err := client.RegisterProbe(
-		ctx,
+		registerCtx,
 		&sm.ProbeInfo{
 			Version:                 version.Short(),
 			Commit:                  version.Commit(),
 			Buildstamp:              version.Buildstamp(),
 			SupportsProtocolSecrets: h.supportsProtocolSecrets,
 		},
-		grpc.WaitForReady(true), // Wait for connection on critical startup RPC
+		grpc.WaitForReady(true), // Wait for connection on critical startup RPC (respects context timeout)
 	)
 	if err != nil {
 		return grpcErrorHandler(

internal/checks/checks.go

@@ -90,6 +90,7 @@ type Updater struct {
 	usageReporter           usage.Reporter
 	tenantCals              *cals.CostAttributionLabels
 	supportsProtocolSecrets bool
+	connectTimeout          time.Duration
 }
 
 type apiInfo struct {
@@ -128,6 +129,7 @@ type UpdaterOptions struct {
 	UsageReporter           usage.Reporter
 	CostAttributionLabels   *cals.CostAttributionLabels
 	SupportsProtocolSecrets bool
+	GrpcConnectTimeout      time.Duration
 }
 
 func NewUpdater(opts UpdaterOptions) (*Updater, error) {
@@ -252,6 +254,7 @@ func NewUpdater(opts UpdaterOptions) (*Updater, error) {
 		tenantSecrets:           opts.SecretProvider,
 		telemeter:               opts.Telemeter,
 		supportsProtocolSecrets: opts.SupportsProtocolSecrets,
+		connectTimeout:          opts.GrpcConnectTimeout,
 		metrics: metrics{
 			changeErrorsCounter: changeErrorsCounter,
 			changesCounter:      changesCounter,
@@ -345,9 +348,11 @@ func handleError(ctx context.Context, logger zerolog.Logger, backoff Backoffer,
 		return true, nil
 
 	case errors.As(err, &transientErr):
-		logger.Warn().Err(err).Msg("transient error, trying to reconnect")
+		dur := backoff.Duration()
 
-		if err := sleepCtx(ctx, backoff.Duration()); err != nil {
+		logger.Warn().Err(err).Dur("backoff", dur).Msg("transient error, trying to reconnect")
+
+		if err := sleepCtx(ctx, dur); err != nil {
 			return true, err
 		}
 
@@ -408,7 +413,7 @@ func (c *Updater) loop(ctx context.Context) (bool, error) {
 			// Context was cancelled
 			return context.Canceled
 
-		case codes.Unavailable:
+		case codes.Unavailable, codes.DeadlineExceeded:
 			// Network errors, connection resets, transport closing, etc.
 			// All these are transient and should trigger retry logic
 			return TransientError(fmt.Sprintf("%s: %s", action, st.Message()))
@@ -427,12 +432,25 @@ func (c *Updater) loop(ctx context.Context) (bool, error) {
 		}
 	}
 
-	result, err := client.RegisterProbe(ctx, &sm.ProbeInfo{
+	// Create a timeout context for RegisterProbe if configured.
+	// This ensures we don't wait indefinitely if the server is unreachable.
+	registerCtx := ctx
+	var cancel context.CancelFunc
+
+	if c.connectTimeout > 0 {
+		registerCtx, cancel = context.WithTimeout(ctx, c.connectTimeout)
+		defer cancel()
+		c.logger.Info().
+			Dur("timeout", c.connectTimeout).
+			Msg("using explicit connection timeout for RegisterProbe")
+	}
+
+	result, err := client.RegisterProbe(registerCtx, &sm.ProbeInfo{
 		Version:                 version.Short(),
 		Commit:                  version.Commit(),
 		Buildstamp:              version.Buildstamp(),
 		SupportsProtocolSecrets: c.supportsProtocolSecrets,
-	}, grpc.WaitForReady(true)) // Wait for connection on critical startup RPC
+	}, grpc.WaitForReady(true)) // Wait for connection on critical startup RPC (respects context timeout)
 	if err != nil {
 		return connected, grpcErrorHandler("registering probe with synthetic-monitoring-api", err)
 	}