runtime: publish load-balanced http client

go.mod

@@ -28,6 +28,7 @@ require (
 require (
 	github.com/antithesishq/antithesis-sdk-go v0.6.0-default-no-op // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/bufbuild/httplb v0.4.1 // indirect
 	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect

go.sum

@@ -18,6 +18,8 @@ github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
 github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
 github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
 github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
+github.com/bufbuild/httplb v0.4.1 h1:f8dMp7tx2aJfMX2UcOId1A58QDiBag7Dv6BA1OtV/YA=
+github.com/bufbuild/httplb v0.4.1/go.mod h1:9XDjl/3UvlkOQUKthLlKn92C1/1SuZ3UCiekxZbenck=
 github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM=
 github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=

httpclient/httplb.go

@@ -0,0 +1,135 @@
+package httpclient
+
+import (
+	"crypto/tls"
+	"net/http"
+	"time"
+
+	"github.com/bufbuild/httplb"
+)
+
+const (
+	// DefaultRequestTimeout bounds each outbound request through the load-balanced client.
+	DefaultRequestTimeout = 30 * time.Second
+	// DefaultIdleConnTimeout controls how long idle pooled connections are retained.
+	DefaultIdleConnTimeout = 90 * time.Second
+	// DefaultMaxIdleConns caps idle connections across all hosts.
+	DefaultMaxIdleConns = 64
+	// DefaultMaxIdleConnsPerHost caps idle connections retained per host.
+	DefaultMaxIdleConnsPerHost = 16
+	// DefaultMaxConnsPerHost caps concurrent connections to a single host.
+	DefaultMaxConnsPerHost = 32
+)
+
+// LoadBalancedOptions controls EvalOps' default outbound HTTP client.
+type LoadBalancedOptions struct {
+	RequestTimeout        time.Duration
+	IdleConnTimeout       time.Duration
+	MaxIdleConns          int
+	MaxIdleConnsPerHost   int
+	MaxConnsPerHost       int
+	MaxResponseHeaderSize int
+	TLSClientConfig       *tls.Config
+	TLSHandshakeTimeout   time.Duration
+	DisableCompression    bool
+}
+
+// LoadBalancedClient owns an httplb client and exposes its standard *http.Client.
+type LoadBalancedClient struct {
+	*http.Client
+
+	client *httplb.Client
+}
+
+var sharedDefaultClient = NewLoadBalanced(LoadBalancedOptions{})
+
+// DefaultClient returns a shared load-balanced client for nil-client fallbacks.
+func DefaultClient() *http.Client {
+	return sharedDefaultClient.Client
+}
+
+// NewLoadBalanced creates an HTTP client with DNS-backed client-side balancing.
+func NewLoadBalanced(options LoadBalancedOptions) *LoadBalancedClient {
+	options = options.withDefaults()
+	lbOptions := []httplb.ClientOption{
+		httplb.WithRequestTimeout(options.RequestTimeout),
+		httplb.WithIdleConnectionTimeout(options.IdleConnTimeout),
+		httplb.WithTransport("http", loadBalancedTransport{options: options}),
+		httplb.WithTransport("https", loadBalancedTransport{options: options}),
+	}
+	if options.TLSClientConfig != nil {
+		lbOptions = append(
+			lbOptions,
+			httplb.WithTLSConfig(options.TLSClientConfig, options.TLSHandshakeTimeout),
+		)
+	}
+	if options.MaxResponseHeaderSize > 0 {
+		lbOptions = append(lbOptions, httplb.WithMaxResponseHeaderBytes(options.MaxResponseHeaderSize))
+	}
+	if options.DisableCompression {
+		lbOptions = append(lbOptions, httplb.WithDisableCompression(true))
+	}
+
+	client := httplb.NewClient(lbOptions...)
+	return &LoadBalancedClient{
+		Client: client.Client,
+		client: client,
+	}
+}
+
+// Close releases resolver and transport resources owned by the httplb client.
+func (c *LoadBalancedClient) Close() error {
+	if c == nil || c.client == nil {
+		return nil
+	}
+	return c.client.Close()
+}
+
+func (o LoadBalancedOptions) withDefaults() LoadBalancedOptions {
+	if o.RequestTimeout <= 0 {
+		o.RequestTimeout = DefaultRequestTimeout
+	}
+	if o.IdleConnTimeout <= 0 {
+		o.IdleConnTimeout = DefaultIdleConnTimeout
+	}
+	if o.MaxIdleConns <= 0 {
+		o.MaxIdleConns = DefaultMaxIdleConns
+	}
+	if o.MaxIdleConnsPerHost <= 0 {
+		o.MaxIdleConnsPerHost = DefaultMaxIdleConnsPerHost
+	}
+	if o.MaxConnsPerHost <= 0 {
+		o.MaxConnsPerHost = DefaultMaxConnsPerHost
+	}
+	return o
+}
+
+type loadBalancedTransport struct {
+	options LoadBalancedOptions
+}
+
+func (t loadBalancedTransport) NewRoundTripper(
+	_ string,
+	_ string,
+	config httplb.TransportConfig,
+) httplb.RoundTripperResult {
+	transport := &http.Transport{
+		Proxy:                  config.ProxyFunc,
+		GetProxyConnectHeader:  config.ProxyConnectHeadersFunc,
+		DialContext:            config.DialFunc,
+		ForceAttemptHTTP2:      true,
+		MaxIdleConns:           t.options.MaxIdleConns,
+		MaxIdleConnsPerHost:    t.options.MaxIdleConnsPerHost,
+		MaxConnsPerHost:        t.options.MaxConnsPerHost,
+		IdleConnTimeout:        config.IdleConnTimeout,
+		TLSHandshakeTimeout:    config.TLSHandshakeTimeout,
+		TLSClientConfig:        config.TLSClientConfig,
+		MaxResponseHeaderBytes: config.MaxResponseHeaderBytes,
+		ExpectContinueTimeout:  time.Second,
+		DisableCompression:     config.DisableCompression,
+	}
+	return httplb.RoundTripperResult{
+		RoundTripper: transport,
+		Close:        transport.CloseIdleConnections,
+	}
+}

httpclient/httplb_test.go

@@ -0,0 +1,82 @@
+package httpclient
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+)
+
+func TestNewLoadBalancedUsesDefaults(t *testing.T) {
+	t.Parallel()
+
+	client := NewLoadBalanced(LoadBalancedOptions{})
+	t.Cleanup(func() {
+		if err := client.Close(); err != nil {
+			t.Fatalf("Close() error = %v", err)
+		}
+	})
+
+	if client.Client == nil {
+		t.Fatal("expected embedded http client")
+	}
+	if client.Transport == nil {
+		t.Fatal("expected httplb transport")
+	}
+	if client.Timeout != 0 {
+		t.Fatalf("Timeout = %s, want httplb request timeout only", client.Timeout)
+	}
+}
+
+func TestNewLoadBalancedHonorsCustomOptions(t *testing.T) {
+	t.Parallel()
+
+	client := NewLoadBalanced(LoadBalancedOptions{
+		RequestTimeout:      2 * time.Second,
+		IdleConnTimeout:     time.Second,
+		MaxIdleConns:        8,
+		MaxIdleConnsPerHost: 4,
+		MaxConnsPerHost:     6,
+	})
+	t.Cleanup(func() {
+		if err := client.Close(); err != nil {
+			t.Fatalf("Close() error = %v", err)
+		}
+	})
+
+	if client.Timeout != 0 {
+		t.Fatalf("Timeout = %s, want httplb request timeout only", client.Timeout)
+	}
+}
+
+func TestLoadBalancedClientSendsRequests(t *testing.T) {
+	t.Parallel()
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if got := r.URL.Path; got != "/healthz" {
+			t.Errorf("path = %q, want /healthz", got)
+		}
+		_, _ = fmt.Fprint(w, "ok")
+	}))
+	defer server.Close()
+
+	client := NewLoadBalanced(LoadBalancedOptions{RequestTimeout: time.Second})
+	t.Cleanup(func() {
+		if err := client.Close(); err != nil {
+			t.Fatalf("Close() error = %v", err)
+		}
+	})
+
+	response, err := client.Get(server.URL + "/healthz")
+	if err != nil {
+		t.Fatalf("Get() error = %v", err)
+	}
+	defer func() {
+		_ = response.Body.Close()
+	}()
+
+	if response.StatusCode != http.StatusOK {
+		t.Fatalf("status = %d, want %d", response.StatusCode, http.StatusOK)
+	}
+}