client.go

// Package bitclient is a Go client SDK for the bit_tracker service.
//
// A Client is bound to a single project. On Connect it queries the project's
// current QoS (claims/second budget) and installs a token-bucket rate limiter
// that throttles Claim calls. The QoS is re-fetched from the server once per
// minute by a background goroutine.
//
// Every request carries a caller-supplied User-Agent header.
package bitclient

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

// DefaultBaseURL is the address the server listens on (see server compose.yml).
const DefaultBaseURL = "http://127.0.0.1:8835"

// DefaultQoSRefreshInterval is how often the background goroutine re-fetches QoS.
const DefaultQoSRefreshInterval = time.Minute

// Client is a bit_tracker client bound to a single project.
//
// Create one with NewClient and start it with Connect. When done, call Close to
// stop the QoS refresh goroutine.
type Client struct {
	baseURL string
	project string
	userA   string // User-Agent sent with every request
	http    *http.Client
	logger  *slog.Logger

	refreshInterval time.Duration

	limiter   *rate.Limiter
	connected bool

	stopOnce sync.Once
	stop     chan struct{}
	done     chan struct{}
}

// Option configures a Client.
type Option func(*config)

type config struct {
	baseURL         string
	project         string
	userAgent       string
	httpClient      *http.Client
	logger          *slog.Logger
	refreshInterval time.Duration
}

// WithBaseURL overrides the server base URL (default DefaultBaseURL).
func WithBaseURL(url string) Option {
	return func(c *config) { c.baseURL = strings.TrimRight(url, "/") }
}

// WithProject sets the project this client is bound to (required).
func WithProject(project string) Option {
	return func(c *config) { c.project = project }
}

// WithUserAgent sets the User-Agent header sent with every request (required).
func WithUserAgent(ua string) Option {
	return func(c *config) { c.userAgent = ua }
}

// WithHTTPClient supplies a custom *http.Client. Defaults to one with a 30s timeout.
func WithHTTPClient(h *http.Client) Option {
	return func(c *config) { c.httpClient = h }
}

// WithLogger supplies a custom *slog.Logger. Defaults to slog.Default().
func WithLogger(l *slog.Logger) Option {
	return func(c *config) { c.logger = l }
}

// WithQoSRefreshInterval overrides how often QoS is re-fetched (default 1m).
func WithQoSRefreshInterval(d time.Duration) Option {
	return func(c *config) { c.refreshInterval = d }
}

// NewClient builds a configured but unconnected Client. Call Connect to fetch
// the initial QoS and start the refresh goroutine.
//
// project and userAgent are required; omitting them returns an error.
func NewClient(opts ...Option) (*Client, error) {
	cfg := config{
		baseURL:         DefaultBaseURL,
		httpClient:      &http.Client{Timeout: 30 * time.Second},
		logger:          slog.Default(),
		refreshInterval: DefaultQoSRefreshInterval,
	}
	for _, o := range opts {
		o(&cfg)
	}
	if cfg.project == "" {
		return nil, fmt.Errorf("bitclient: WithProject is required")
	}
	if cfg.userAgent == "" {
		return nil, fmt.Errorf("bitclient: WithUserAgent is required")
	}
	if cfg.refreshInterval <= 0 {
		return nil, fmt.Errorf("bitclient: refresh interval must be positive")
	}

	return &Client{
		baseURL:         cfg.baseURL,
		project:         cfg.project,
		userA:           cfg.userAgent,
		http:            cfg.httpClient,
		logger:          cfg.logger,
		refreshInterval: cfg.refreshInterval,
		stop:            make(chan struct{}),
		done:            make(chan struct{}),
	}, nil
}

// Connect queries the current QoS from the server, installs the rate limiter
// used to throttle Claim calls, and starts a goroutine that re-fetches the QoS
// once per refresh interval.
func (c *Client) Connect(ctx context.Context) error {
	qos, err := c.fetchQoS(ctx)
	if err != nil {
		return fmt.Errorf("bitclient: initial QoS fetch failed: %w", err)
	}
	c.limiter = rate.NewLimiter(rate.Limit(qos), burstFor(qos))
	c.connected = true
	c.logger.Info("bitclient connected", "project", c.project, "qos", qos)

	go c.refreshLoop()
	return nil
}

// Close stops the background QoS refresh goroutine. It is safe to call multiple
// times and is a no-op if Connect was never called.
func (c *Client) Close() {
	c.stopOnce.Do(func() {
		close(c.stop)
	})
	<-c.done
}

// burstFor returns the token-bucket burst. We allow a burst of one second's
// worth of claims (min 1) so a healthy rate isn't artificially starved while
// still smoothing the average to `qos` claims/second.
func burstFor(qos float64) int {
	b := int(qos)
	if b < 1 {
		b = 1
	}
	return b
}