diff --git a/internal/auth/xai/xai.go b/internal/auth/xai/xai.go index 65d988c3..a0b76ea8 100644 --- a/internal/auth/xai/xai.go +++ b/internal/auth/xai/xai.go @@ -19,7 +19,8 @@ import ( // XAIAuth performs xAI OAuth discovery, device-code login, and refresh. type XAIAuth struct { - httpClient *http.Client + httpClient *http.Client + minPollInterval time.Duration } var xaiRefreshGroup singleflight.Group @@ -211,9 +212,16 @@ func (a *XAIAuth) PollForToken(ctx context.Context, deviceCode *DeviceCodeRespon tokenEndpoint = discovery.TokenEndpoint } + minInterval := defaultPollInterval + if a != nil && a.minPollInterval > 0 { + minInterval = a.minPollInterval + } + interval := time.Duration(deviceCode.Interval) * time.Second - if interval < defaultPollInterval { - interval = defaultPollInterval + if a != nil && a.minPollInterval > 0 && deviceCode.Interval <= 0 { + interval = a.minPollInterval + } else if interval < minInterval { + interval = minInterval } deadline := time.Now().Add(MaxPollDuration) @@ -301,7 +309,11 @@ func (a *XAIAuth) exchangeDeviceCode(ctx context.Context, tokenEndpoint, deviceC case "authorization_pending": return nil, nil, interval, true case "slow_down": - nextInterval := interval + defaultPollInterval + step := defaultPollInterval + if a != nil && a.minPollInterval > 0 { + step = a.minPollInterval + } + nextInterval := interval + step return nil, nil, nextInterval, true case "expired_token": return nil, fmt.Errorf("xai device code expired"), interval, false diff --git a/internal/auth/xai/xai_auth_test.go b/internal/auth/xai/xai_auth_test.go index 9554f8c2..5773fd90 100644 --- a/internal/auth/xai/xai_auth_test.go +++ b/internal/auth/xai/xai_auth_test.go @@ -116,11 +116,12 @@ func TestPollForTokenExchangesDeviceCode(t *testing.T) { defer server.Close() auth := NewXAIAuth(nil) + auth.minPollInterval = time.Millisecond tokenData, err := auth.PollForToken(context.Background(), &DeviceCodeResponse{ DeviceCode: "device-abc", UserCode: "ABCD-1234", ExpiresIn: 60, - Interval: 1, + Interval: 0, TokenEndpoint: server.URL, }) if err != nil { @@ -187,11 +188,12 @@ func TestPollForTokenSlowDownContinuesPolling(t *testing.T) { defer server.Close() auth := NewXAIAuth(nil) + auth.minPollInterval = time.Millisecond tokenData, err := auth.PollForToken(context.Background(), &DeviceCodeResponse{ DeviceCode: "device-abc", UserCode: "ABCD-1234", ExpiresIn: 60, - Interval: 5, + Interval: 0, TokenEndpoint: server.URL, }) if err != nil { @@ -320,6 +322,16 @@ func TestRefreshTokens_DeduplicatesConcurrentRefresh(t *testing.T) { } } +func TestXAIAuthDefaultPollInterval(t *testing.T) { + auth := NewXAIAuth(nil) + if auth.minPollInterval != 0 { + t.Fatalf("auth.minPollInterval = %v, want 0 (defaults to %v in production)", auth.minPollInterval, defaultPollInterval) + } + if defaultPollInterval != 5*time.Second { + t.Fatalf("defaultPollInterval = %v, want 5s", defaultPollInterval) + } +} + func fakeJWTWithEmail(email, subject string) string { header := base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString([]byte(`{"alg":"none","typ":"JWT"}`)) payload := base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString([]byte(`{"email":"` + email + `","sub":"` + subject + `"}`)) diff --git a/internal/home/client.go b/internal/home/client.go index 00e9316c..dc299b8c 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -200,12 +200,13 @@ type Client struct { // lets a Home upgrade take effect on the next reconnect instead of // requiring a CPA restart. The probe costs one round trip that returns an // error without performing any write. - casUnsupported atomic.Bool - recoveryState atomic.Uint32 - instanceID string - legacyMembership bool - clusterNodes []clusterNode - reconnectFailures int + casUnsupported atomic.Bool + testOperationTimeout time.Duration + recoveryState atomic.Uint32 + instanceID string + legacyMembership bool + clusterNodes []clusterNode + reconnectFailures int } func New(homeCfg config.HomeConfig) *Client { @@ -225,13 +226,14 @@ func (c *Client) NewLifetime() *Client { c.mu.Lock() defer c.mu.Unlock() next := &Client{ - homeCfg: c.homeCfg, - seedHost: c.seedHost, - seedPort: c.seedPort, - clusterNodes: append([]clusterNode(nil), c.clusterNodes...), - reconnectFailures: c.reconnectFailures, - instanceID: c.instanceID, - legacyMembership: c.legacyMembership, + homeCfg: c.homeCfg, + seedHost: c.seedHost, + seedPort: c.seedPort, + clusterNodes: append([]clusterNode(nil), c.clusterNodes...), + reconnectFailures: c.reconnectFailures, + testOperationTimeout: c.testOperationTimeout, + instanceID: c.instanceID, + legacyMembership: c.legacyMembership, } next.recoveryState.Store(c.recoveryState.Load()) return next @@ -517,12 +519,30 @@ func (c *Client) redisOptionsLocked(addr string) (*redis.Options, error) { if errTLS != nil { return nil, errTLS } + dialTimeout := homeRedisOperationTimeout + readTimeout := homeRedisOperationTimeout + writeTimeout := homeRedisOperationTimeout + if c.testOperationTimeout > 0 { + dialTimeout = c.testOperationTimeout + readTimeout = c.testOperationTimeout + writeTimeout = c.testOperationTimeout + } else if c.cmdOptions != nil { + if c.cmdOptions.DialTimeout > 0 { + dialTimeout = c.cmdOptions.DialTimeout + } + if c.cmdOptions.ReadTimeout > 0 { + readTimeout = c.cmdOptions.ReadTimeout + } + if c.cmdOptions.WriteTimeout > 0 { + writeTimeout = c.cmdOptions.WriteTimeout + } + } options := &redis.Options{ Addr: addr, TLSConfig: tlsConfig, - DialTimeout: homeRedisOperationTimeout, - ReadTimeout: homeRedisOperationTimeout, - WriteTimeout: homeRedisOperationTimeout, + DialTimeout: dialTimeout, + ReadTimeout: readTimeout, + WriteTimeout: writeTimeout, MaxRetries: -1, DialerRetries: 1, ContextTimeoutEnabled: true, @@ -1766,13 +1786,19 @@ func (c *Client) subscriptionParameters() ([]string, time.Duration) { cfg := c.lifecycle.WithDefaults() instanceID := c.instanceID legacyMembership := c.legacyMembership + testTimeout := c.testOperationTimeout c.mu.Unlock() + timeout := cfg.CPAHeartbeatTimeout + if testTimeout > 0 && cfg.LifecycleConfigRevision == 0 { + timeout = testTimeout + } + args := []string{redisChannelConfig} if cfg.LifecycleConfigRevision > 0 { args = append(args, strconv.FormatInt(cfg.LifecycleConfigRevision, 10)) if legacyMembership { - return args, cfg.CPAHeartbeatTimeout + return args, timeout } state := recoveryState(c.recoveryState.Load()) if state == recoveryStateTakeoverEligible || state == recoveryStateSwitchingTakeover { @@ -1780,7 +1806,7 @@ func (c *Client) subscriptionParameters() ([]string, time.Duration) { } args = append(args, instanceID) } - return args, cfg.CPAHeartbeatTimeout + return args, timeout } func (c *Client) markMembershipTakeoverEligible() { diff --git a/internal/home/client_test.go b/internal/home/client_test.go index 351b1ab5..effe8b43 100644 --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -1233,6 +1233,7 @@ func newRedisCommandTestClient(t *testing.T, handler func([]string) string) (*Cl Port: port, DisableClusterDiscovery: true, }) + client.testOperationTimeout = homeRedisTestOperationTimeout options := &redis.Options{ Addr: listener.Addr().String(), Protocol: 2, @@ -1245,6 +1246,7 @@ func newRedisCommandTestClient(t *testing.T, handler func([]string) string) (*Cl } client.cmdOptions = cloneRedisOptions(options) client.cmd = redis.NewClient(options) + client.sub = redis.NewClient(cloneRedisOptions(options)) t.Cleanup(func() { client.Close() }) @@ -1635,8 +1637,10 @@ func TestRunConfigSubscriberLifetimeRejectsInvalidSubscriptionACK(t *testing.T) switch { case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "CLUSTER") && strings.EqualFold(args[1], "NODES"): + return "-ERR unknown command 'CLUSTER'\r\n" case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: - return "$16\r\nhost: 127.0.0.1\r\n" + return "$15\r\nhost: 127.0.0.1\r\n" case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: return ack default: @@ -2254,6 +2258,30 @@ func TestRunConfigSubscriberLifetimePreservesTakeoverWhenFreshCommandProbeFails( } } +func TestDefaultProductionClientTimeouts(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 6379}) + if client.testOperationTimeout != 0 { + t.Fatalf("default testOperationTimeout = %v, want 0", client.testOperationTimeout) + } + options, err := client.redisOptionsLocked("127.0.0.1:6379") + if err != nil { + t.Fatalf("redisOptionsLocked() error = %v", err) + } + if options.DialTimeout != homeRedisOperationTimeout { + t.Fatalf("DialTimeout = %v, want %v", options.DialTimeout, homeRedisOperationTimeout) + } + if options.ReadTimeout != homeRedisOperationTimeout { + t.Fatalf("ReadTimeout = %v, want %v", options.ReadTimeout, homeRedisOperationTimeout) + } + if options.WriteTimeout != homeRedisOperationTimeout { + t.Fatalf("WriteTimeout = %v, want %v", options.WriteTimeout, homeRedisOperationTimeout) + } + _, timeout := client.subscriptionParameters() + if timeout != 3*time.Second { + t.Fatalf("subscriptionParameters timeout = %v, want 3s", timeout) + } +} + func findRedisCommand(commands [][]string, commandName string) []string { for _, command := range commands { if len(command) > 0 && strings.EqualFold(command[0], commandName) {