l-ai-chat-client-go/options.go

package chatclient

import (
	"context"
	"crypto/tls"
	"fmt"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
)

const (
	defaultSubscriptionBufferSize = 100
	defaultReconnectInitial       = 200 * time.Millisecond
	defaultReconnectMax           = 5 * time.Second
	defaultCloseTimeout           = 5 * time.Second
)

type Sleeper func(ctx context.Context, d time.Duration) error
type Jitter func(d time.Duration) time.Duration

type Option func(*options)

type options struct {
	tokenProvider          TokenProvider
	insecureTransport      bool
	dialOptions            []grpc.DialOption
	conn                   *grpc.ClientConn
	subscriptionBufferSize int
	reconnectInitial       time.Duration
	reconnectMax           time.Duration
	closeTimeout           time.Duration
	sleeper                Sleeper
	jitter                 Jitter
}

func defaultOptions() options {
	return options{
		subscriptionBufferSize: defaultSubscriptionBufferSize,
		reconnectInitial:       defaultReconnectInitial,
		reconnectMax:           defaultReconnectMax,
		closeTimeout:           defaultCloseTimeout,
		sleeper:                sleepContext,
		jitter:                 defaultJitter,
	}
}

func WithTokenProvider(provider TokenProvider) Option {
	return func(o *options) { o.tokenProvider = provider }
}

func WithStaticToken(token string) Option {
	return WithTokenProvider(StaticToken(token))
}

func WithInsecureTransport() Option {
	return func(o *options) { o.insecureTransport = true }
}

func WithGRPCDialOptions(opts ...grpc.DialOption) Option {
	return func(o *options) { o.dialOptions = append(o.dialOptions, opts...) }
}

func WithClientConn(conn *grpc.ClientConn) Option {
	return func(o *options) { o.conn = conn }
}

func WithSubscriptionBufferSize(size int) Option {
	return func(o *options) { o.subscriptionBufferSize = size }
}

func WithReconnectBackoff(initial, max time.Duration) Option {
	return func(o *options) {
		o.reconnectInitial = initial
		o.reconnectMax = max
	}
}

func withSleeperForTest(s Sleeper) Option {
	return func(o *options) { o.sleeper = s }
}

func withJitterForTest(j Jitter) Option {
	return func(o *options) { o.jitter = j }
}

func buildDialOptions(o options) ([]grpc.DialOption, error) {
	if o.tokenProvider == nil {
		return nil, fmt.Errorf("%w: token provider is required", ErrInvalidConfig)
	}

	var transport credentials.TransportCredentials
	if o.insecureTransport {
		transport = insecure.NewCredentials()
	} else {
		transport = credentials.NewTLS(&tls.Config{MinVersion: tls.VersionTLS12})
	}

	creds := tokenCredentials{
		provider:                 o.tokenProvider,
		requireTransportSecurity: !o.insecureTransport,
	}
	out := []grpc.DialOption{
		grpc.WithTransportCredentials(transport),
		grpc.WithPerRPCCredentials(creds),
	}
	out = append(out, o.dialOptions...)
	return out, nil
}

func validateOptions(o options) error {
	if o.tokenProvider == nil {
		return fmt.Errorf("%w: token provider is required", ErrInvalidConfig)
	}
	if o.subscriptionBufferSize <= 0 {
		return fmt.Errorf("%w: subscription buffer size must be positive", ErrInvalidConfig)
	}
	if o.reconnectInitial <= 0 {
		return fmt.Errorf("%w: reconnect initial backoff must be positive", ErrInvalidConfig)
	}
	if o.reconnectMax < o.reconnectInitial {
		return fmt.Errorf("%w: reconnect max backoff must be >= initial backoff", ErrInvalidConfig)
	}
	if o.closeTimeout <= 0 {
		return fmt.Errorf("%w: close timeout must be positive", ErrInvalidConfig)
	}
	if o.sleeper == nil || o.jitter == nil {
		return fmt.Errorf("%w: backoff hooks must be configured", ErrInvalidConfig)
	}
	return nil
}

func sleepContext(ctx context.Context, d time.Duration) error {
	timer := time.NewTimer(d)
	defer timer.Stop()
	select {
	case <-timer.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func defaultJitter(d time.Duration) time.Duration {
	// Deterministic half-jitter keeps reconnects bounded without pulling in a
	// retry framework. Tests replace this hook.
	if d <= 1 {
		return d
	}
	return d - d/4
}