Add a ShardedClient class for partitioned setups

client/client.go

@@ -3,8 +3,10 @@ package client
 import (
 	"context"
 	"crypto/tls"
+	"fmt"
 	"net/http"
 	"net/url"
+	"sync"
 	"time"
 
 	"git.autistici.org/ai3/go-common/clientutil"
@@ -12,7 +14,7 @@ import (
 	"git.autistici.org/id/keystore"
 )
 
-// Client for the keystore API.
+// Client for the keystore API (for a specific backend).
 type Client struct {
 	*http.Client
 	backendURL string
@@ -24,7 +26,8 @@ type Config struct {
 	TLSConfig  *clientutil.TLSClientConfig `yaml:"tls_config"`
 }
 
-// New returns a new Client with the given Config.
+// New returns a new Client with the given Config. Use this when the
+// keystore service runs on a single global instance.
 func New(config *Config) (*Client, error) {
 	u, err := url.Parse(config.BackendURL)
 	if err != nil {
@@ -76,3 +79,60 @@ func (c *Client) Close(ctx context.Context, username string) error {
 	var resp keystore.CloseResponse
 	return clientutil.DoJSONHTTPRequest(ctx, c.Client, c.backendURL+"/api/close", &req, &resp)
 }
+
+// ShardedClient for the keystore API (sharded service).
+type ShardedClient struct {
+	baseURL   *url.URL
+	tlsConfig *tls.Config
+	mx        sync.Mutex
+	shards    map[string]*Client
+}
+
+// NewSharded creates a ShardedClient for the keystore service. Use it
+// when the service is partitioned (sharded) across multiple backends.
+func NewSharded(config *Config) (*ShardedClient, error) {
+	u, err := url.Parse(config.BackendURL)
+	if err != nil {
+		return nil, err
+	}
+
+	var tlsConfig *tls.Config
+	if config.TLSConfig != nil {
+		tlsConfig, err = config.TLSConfig.TLSConfig()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return &ShardedClient{
+		baseURL:   u,
+		tlsConfig: tlsConfig,
+		shards:    make(map[string]*Client),
+	}, nil
+}
+
+func (c *ShardedClient) getShardURL(shard string) *url.URL {
+	u := *c.baseURL
+	u.Host = fmt.Sprintf("%s.%s", shard, u.Host)
+	return &u
+}
+
+// Shard returns the Client for a specific service shard.
+func (c *ShardedClient) Shard(shard string) *Client {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
+	client, ok := c.shards[shard]
+	if !ok {
+		u := c.getShardURL(shard)
+		client = &Client{
+			Client: &http.Client{
+				Transport: clientutil.NewTransport([]string{u.Host}, c.tlsConfig, nil),
+				Timeout:   20 * time.Second,
+			},
+			backendURL: u.String(),
+		}
+		c.shards[shard] = client
+	}
+	return client
+}