noncebalancer: use endpointsharding, ignore ready status

cmd/boulder-wfe2/main.go

@@ -19,6 +19,7 @@ import (
 	"github.com/letsencrypt/boulder/goodkey/sagoodkey"
 	bgrpc "github.com/letsencrypt/boulder/grpc"
 	"github.com/letsencrypt/boulder/grpc/noncebalancer"
+	noncebalancerv1 "github.com/letsencrypt/boulder/grpc/noncebalancerv1"
 	"github.com/letsencrypt/boulder/issuance"
 	"github.com/letsencrypt/boulder/nonce"
 	rapb "github.com/letsencrypt/boulder/ra/proto"
@@ -318,9 +319,11 @@ func main() {
 	cmd.FailOnError(err, "Failed to load credentials and create gRPC connection to get nonce service")
 	gnc := nonce.NewGetter(getNonceConn)
 
-	if c.WFE.RedeemNonceService.SRVResolver != noncebalancer.SRVResolverScheme {
+	if c.WFE.RedeemNonceService.SRVResolver != noncebalancer.SRVResolverScheme &&
+		c.WFE.RedeemNonceService.SRVResolver != noncebalancerv1.SRVResolverScheme {
 		cmd.Fail(fmt.Sprintf(
-			"'redeemNonceService.SRVResolver' must be set to %q", noncebalancer.SRVResolverScheme),
+			"'redeemNonceService.SRVResolver' must be set to %q or %q",
+			noncebalancer.SRVResolverScheme, noncebalancerv1.SRVResolverScheme),
 		)
 	}
 	redeemNonceConn, err := bgrpc.ClientSetup(c.WFE.RedeemNonceService, tlsConfig, stats, clk)

cmd/config.go

@@ -306,7 +306,7 @@ type GRPCClientConfig struct {
 	// implementation of the SRV resolver should be used. The default is 'srv'
 	// For more details, see the documentation in:
 	// grpc/internal/resolver/dns/dns_resolver.go.
-	SRVResolver string `validate:"excluded_with=ServerAddress,isdefault|oneof=srv nonce-srv"`
+	SRVResolver string `validate:"excluded_with=ServerAddress,isdefault|oneof=srv nonce-srv nonce-srv-v2"`
 
 	// ServerAddress is a single <hostname|IPv4|[IPv6]>:<port> or `:<port>` that
 	// the gRPC client will, if necessary, resolve via DNS and then connect to.

grpc/internal/resolver/dns/dns_resolver.go

@@ -40,6 +40,7 @@ import (
 	"github.com/letsencrypt/boulder/bdns"
 	"github.com/letsencrypt/boulder/grpc/internal/backoff"
 	"github.com/letsencrypt/boulder/grpc/noncebalancer"
+	noncebalancerv1 "github.com/letsencrypt/boulder/grpc/noncebalancerv1"
 )
 
 var logger = grpclog.Component("srv")
@@ -54,6 +55,7 @@ var (
 func init() {
 	resolver.Register(NewDefaultSRVBuilder())
 	resolver.Register(NewNonceSRVBuilder())
+	resolver.Register(NewNonceSRVBuilderV2())
 }
 
 const defaultDNSSvrPort = "53"
@@ -90,9 +92,15 @@ func NewDefaultSRVBuilder() resolver.Builder {
 	return &srvBuilder{scheme: "srv"}
 }
 
-// NewNonceSRVBuilder creates a srvBuilder which is used to factory SRV DNS
+// NewNonceSRVBuilder creates a srvBuilder which is used to build SRV DNS
 // resolvers with a custom grpc.Balancer used by nonce-service clients.
 func NewNonceSRVBuilder() resolver.Builder {
+	return &srvBuilder{scheme: noncebalancerv1.SRVResolverScheme, balancer: noncebalancerv1.Name}
+}
+
+// NewNonceSRVBuilderV2 creates a srvBuilder which is used to build SRV DNS
+// resolvers with a custom nonce balancer used by nonce-service clients.
+func NewNonceSRVBuilderV2() resolver.Builder {
 	return &srvBuilder{scheme: noncebalancer.SRVResolverScheme, balancer: noncebalancer.Name}
 }
 

grpc/noncebalancer/noncebalancer.go

@@ -4,24 +4,26 @@ import (
 	"errors"
 	"sync"
 
-	"github.com/letsencrypt/boulder/nonce"
-
 	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/status"
+
+	"github.com/letsencrypt/boulder/nonce"
 )
 
 const (
 	// Name is the name used to register the nonce balancer with the gRPC
 	// runtime.
-	Name = "nonce"
+	Name = "noncev2"
 
 	// SRVResolverScheme is the scheme used to invoke an instance of the SRV
 	// resolver which will use the noncebalancer to pick backends. It would be
 	// ideal to export this from the SRV resolver package but that package is
 	// internal.
-	SRVResolverScheme = "nonce-srv"
+	SRVResolverScheme = "nonce-srv-v2"
 )
 
 // ErrNoBackendsMatchPrefix indicates that no backends were found which match
@@ -40,38 +42,43 @@ var errMissingHMACKeyCtxKey = errors.New("nonce.HMACKeyCtxKey value required in
 var errInvalidPrefixCtxKeyType = errors.New("nonce.PrefixCtxKey value in RPC context must be a string")
 var errInvalidHMACKeyCtxKeyType = errors.New("nonce.HMACKeyCtxKey value in RPC context must be a byte slice")
 
-// pickerBuilder implements the base.PickerBuilder interface. It's used to
-// create new Picker instances. It should only be used by nonce-service clients.
-type pickerBuilder struct{}
-
-// Build implements the base.PickerBuilder interface. It is called by the gRPC
-// runtime when the balancer is first initialized and when the set of backend
-// (SubConn) addresses changes.
-func (b *pickerBuilder) Build(buildInfo base.PickerBuildInfo) balancer.Picker {
-	if len(buildInfo.ReadySCs) == 0 {
-		// The Picker must be rebuilt if there are no backends available.
-		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
-	}
-	return &picker{
-		backends: buildInfo.ReadySCs,
-	}
+// prefixBasedPicker implements the balancer.Picker interface. It delegates to a child Picker
+// based on the endpoint (IP address and port) that Picker represents.
+// The child picker is provided by endpointsharding's Balancer implementation
+// (https://pkg.go.dev/google.golang.org/grpc/balancer/endpointsharding), which
+// abstracts away the creation and management of SubConns for us.
+//
+// We happen to know the child Picker is created by the "pickfirst" balancer, but
+// since each child Picker only has a single Endpoint anyhow, it doesn't really matter.
+type prefixBasedPicker struct {
+	// This is the full list of (address -> Picker) pairs passed in by the nonceBalancer.
+	// In particular it is not filtered based on the state of any SubConn, since a given
+	// address' SubConn may be temporarily unavailable while reconnecting, and we still
+	// want to attempt sending traffic to that endpoint if we receive the corresponding
+	// prefix.
+	addrToPicker map[string]balancer.Picker
+
+	// A mapping from nonce prefix to the child picker for that backend. This is derived,
+	// on first Pick call, from the address of each backend plus the HMAC key passed in a
+	// context.Context. We don't derive it on construction because we don't have access to
+	// the HMAC key at that point.
+	prefixToPicker     map[string]balancer.Picker
+	prefixToPickerOnce sync.Once
 }
 
-// picker implements the balancer.Picker interface. It picks a backend (SubConn)
-// based on the nonce prefix contained in each request's Context.
-type picker struct {
-	backends            map[balancer.SubConn]base.SubConnInfo
-	prefixToBackend     map[string]balancer.SubConn
-	prefixToBackendOnce sync.Once
+// newPicker creates a prefixBasedPicker with the given map of addresses to child pickers.
+func newPicker(m map[string]balancer.Picker) *prefixBasedPicker {
+	return &prefixBasedPicker{
+		addrToPicker: m,
+	}
 }
 
 // Pick implements the balancer.Picker interface. It is called by the gRPC
 // runtime for each RPC message. It is responsible for picking a backend
 // (SubConn) based on the context of each RPC message.
-func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
-	if len(p.backends) == 0 {
-		// This should never happen, the Picker should only be built when there
-		// are backends available.
+func (p *prefixBasedPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	if len(p.addrToPicker) == 0 {
+		// Should never happen.
 		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
 	}
 
@@ -87,14 +94,14 @@ func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
 		return balancer.PickResult{}, errInvalidHMACKeyCtxKeyType
 	}
 
-	p.prefixToBackendOnce.Do(func() {
+	p.prefixToPickerOnce.Do(func() {
 		// First call to Pick with a new Picker.
-		prefixToBackend := make(map[string]balancer.SubConn)
-		for sc, scInfo := range p.backends {
-			scPrefix := nonce.DerivePrefix(scInfo.Address.Addr, hmacKey)
-			prefixToBackend[scPrefix] = sc
+		prefixToPicker := make(map[string]balancer.Picker)
+		for addr, picker := range p.addrToPicker {
+			prefix := nonce.DerivePrefix(addr, hmacKey)
+			prefixToPicker[prefix] = picker
 		}
-		p.prefixToBackend = prefixToBackend
+		p.prefixToPicker = prefixToPicker
 	})
 
 	// Get the destination prefix from the RPC context.
@@ -109,16 +116,66 @@ func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
 		return balancer.PickResult{}, errInvalidPrefixCtxKeyType
 	}
 
-	sc, ok := p.prefixToBackend[destPrefix]
+	childPicker, ok := p.prefixToPicker[destPrefix]
 	if !ok {
 		// No backend SubConn was found for the destination prefix.
 		return balancer.PickResult{}, ErrNoBackendsMatchPrefix.Err()
 	}
-	return balancer.PickResult{SubConn: sc}, nil
+	return childPicker.Pick(info)
+}
+
+// builder builds a nonceBalancer, which internally uses `endpointsharding.NewBalancer`.
+//
+// The embedded `endpointsharding` balancer manages a set of child pickers that all use
+// `pickfirst` on an endpoint that consists of a single IP address (because our `"nonce-srv-v2"`
+// resolver returns single-IP endpoints).
+type builder struct{}
+
+func (b builder) Name() string {
+	return Name
+}
+
+func (b builder) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+	childBalancerBuilder := balancer.Get(pickfirst.Name).Build
+	nb := &nonceBalancer{
+		ClientConn: cc,
+	}
+	nb.Balancer = endpointsharding.NewBalancer(nb, bOpts, childBalancerBuilder, endpointsharding.Options{})
+	return nb
+}
+
+// nonceBalancer sends nonce redemption requests to backends based on the nonce prefix,
+// which maps to a specific IP address and port pair.
+type nonceBalancer struct {
+	balancer.Balancer
+	balancer.ClientConn
+}
+
+// UpdateState creates a `prefixBasedPicker` that is aware of the IP address and port of all
+// the child pickers available, including ones that may not have an active connection.
+//
+// The child pickers are all `pickfirst` across a single IP address.
+func (b *nonceBalancer) UpdateState(state balancer.State) {
+	if state.ConnectivityState != connectivity.Ready {
+		b.ClientConn.UpdateState(state)
+		return
+	}
+
+	addrToPicker := make(map[string]balancer.Picker)
+	for _, childState := range endpointsharding.ChildStatesFromPicker(state.Picker) {
+		// We expect our Endpoints to always have single Addresses, but might as well
+		// be robust to the possibility there are more.
+		for _, addr := range childState.Endpoint.Addresses {
+			addrToPicker[addr.Addr] = childState.State.Picker
+		}
+	}
+	b.ClientConn.UpdateState(balancer.State{
+		ConnectivityState: state.ConnectivityState,
+		// Here's where we build our nonce-aware picker.
+		Picker: newPicker(addrToPicker),
+	})
 }
 
 func init() {
-	balancer.Register(
-		base.NewBalancerBuilder(Name, &pickerBuilder{}, base.Config{}),
-	)
+	balancer.Register(builder{})
 }