diff --git a/balancer.go b/balancer.go index 97cb5a2..933ab0a 100644 --- a/balancer.go +++ b/balancer.go @@ -16,6 +16,7 @@ import ( "errors" "fmt" "hash/maphash" + "sync" "google.golang.org/grpc/balancer" "google.golang.org/grpc/balancer/base" @@ -112,7 +113,10 @@ var logger = grpclog.Component("consistenthashring") // balancer.Register(consistent.NewBuilder(xxhash.Sum64)) // ``` func NewBuilder(hashfn hashring.HashFunc) Builder { - return &builder{hashfn: hashfn} + return &builder{ + hashfn: hashfn, + slots: make(map[string]*ringSlot), + } } type subConnMember struct { @@ -128,20 +132,53 @@ var _ hashring.Member = (*subConnMember)(nil) type builder struct { hashfn hashring.HashFunc + + mu sync.Mutex + slots map[string]*ringSlot } -// Builder combines both of gRPC's `balancer.Builder` and -// `balancer.ConfigParser` interfaces. +// Builder combines gRPC's `balancer.Builder` and `balancer.ConfigParser` +// interfaces. It also exposes read-only views of the hashrings of the +// balancers it builds. type Builder interface { balancer.Builder balancer.ConfigParser + + // RingFor returns a read-only view of the live hashring that routes + // requests to the given target. The target is the string form of the + // dial target, for example "kubernetes:///spicedb.default:50053". + // + // The returned view is always non-nil, and callers can get it before any + // balancer for the target exists. Its methods return ErrNoRing (or no + // members) until the balancer produces a ring. ClientConns that share a + // target string share one view, and the last writer wins. + RingFor(target string) RingView } var _ Builder = (*builder)(nil) func (b *builder) Name() string { return BalancerName } -func (b *builder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer { +// slotFor returns the ringSlot for the given target. It creates the slot if +// none exists. +func (b *builder) slotFor(target string) *ringSlot { + b.mu.Lock() + defer b.mu.Unlock() + + slot, ok := b.slots[target] + if !ok { + slot = &ringSlot{} + b.slots[target] = slot + } + + return slot +} + +func (b *builder) RingFor(target string) RingView { + return b.slotFor(target) +} + +func (b *builder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer { bal := &ringBalancer{ cc: cc, subConns: resolver.NewAddressMapV2[any](), @@ -152,6 +189,7 @@ func (b *builder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balance state: connectivity.Connecting, hasher: b.hashfn, picker: base.NewErrPicker(balancer.ErrNoSubConnAvailable), + slot: b.slotFor(opts.Target.String()), } return bal @@ -192,6 +230,9 @@ type ringBalancer struct { config *BalancerConfig hashring *hashring.Ring hasher hashring.HashFunc + // slot is where the balancer publishes its live ring for RingView + // readers. + slot *ringSlot resolverErr error // the last error reported by the resolver; cleared on successful resolution connErr error // the last connection error; cleared upon leaving TransientFailure @@ -235,6 +276,7 @@ func (b *ringBalancer) UpdateClientConnState(s balancer.ClientConnState) error { svcConfig := s.BalancerConfig.(*BalancerConfig) if b.config == nil || svcConfig.ReplicationFactor != b.config.ReplicationFactor { b.hashring = hashring.MustNew(b.hasher, svcConfig.ReplicationFactor) + b.slot.ring.Store(b.hashring) // The new ring starts empty: put every SubConn that has not // failed back on it. b.ringMembers = make(map[balancer.SubConn]struct{}) @@ -442,7 +484,11 @@ func (b *ringBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.Su } func (b *ringBalancer) Close() { - // No internal state to clean up and no need to call RemoveSubConn. + // Withdraw the published ring so that RingView readers stop routing to a + // balancer that no longer exists. The CompareAndSwap clears only this + // balancer's own ring, so it cannot remove the ring of a second balancer + // that shares the target's slot. + b.slot.ring.CompareAndSwap(b.hashring, nil) } func (b *ringBalancer) ExitIdle() { diff --git a/ringview.go b/ringview.go new file mode 100644 index 0000000..a32c030 --- /dev/null +++ b/ringview.go @@ -0,0 +1,59 @@ +package consistent + +import ( + "errors" + "sync/atomic" + + "github.com/authzed/consistent/hashring" +) + +// ErrNoRing means that no hashring exists for the target. gRPC builds +// balancers lazily, so before the first connection attempt there is no ring +// to read. A closed balancer also withdraws its ring. +var ErrNoRing = errors.New("no hashring available for target") + +// RingView is a read-only view of the live hashring a balancer routes with. +// +// A view is a stable handle. It follows the balancer's ring across +// replacements (a ReplicationFactor change replaces the ring) and it shows +// membership changes as they happen. Queries through a view therefore agree +// with the picker's routing decisions. The one exception is a membership +// change that races the query: that can misplace a request's locality, never +// its delivery. +type RingView interface { + // FindN returns the first n unique members of the ring for the given key, + // exactly as the balancer's picker sees them. It returns ErrNoRing if no + // ring exists yet for the target. + FindN(key []byte, n uint8) ([]hashring.Member, error) + + // Members returns all members currently on the ring. It returns nil if + // no ring exists yet for the target. + Members() []hashring.Member +} + +// ringSlot is the shared cell that connects a balancer to its RingView +// readers: the balancer publishes its live ring here and the readers load +// it. The pointer is nil until the balancer creates its first ring. The +// balancer swaps the pointer when it replaces the ring and clears the +// pointer on Close. +type ringSlot struct { + ring atomic.Pointer[hashring.Ring] +} + +var _ RingView = (*ringSlot)(nil) + +func (s *ringSlot) FindN(key []byte, n uint8) ([]hashring.Member, error) { + r := s.ring.Load() + if r == nil { + return nil, ErrNoRing + } + return r.FindN(key, n) +} + +func (s *ringSlot) Members() []hashring.Member { + r := s.ring.Load() + if r == nil { + return nil + } + return r.Members() +} diff --git a/ringview_test.go b/ringview_test.go new file mode 100644 index 0000000..24ee2f1 --- /dev/null +++ b/ringview_test.go @@ -0,0 +1,143 @@ +package consistent + +import ( + "net/url" + "testing" + + "github.com/cespare/xxhash/v2" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/balancer" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/resolver" +) + +func targetNamed(t *testing.T, s string) resolver.Target { + t.Helper() + u, err := url.Parse(s) + require.NoError(t, err) + return resolver.Target{URL: *u} +} + +// readyBalancerForTarget builds a balancer from bld for the given target and +// moves every SubConn for addrs to READY. +func readyBalancerForTarget(t *testing.T, bld Builder, target string, addrs ...resolver.Address) *ringBalancer { + t.Helper() + cc := newFakeClientConn() + cc.stateCh = make(chan balancer.State, 64) + b := bld.Build(cc, balancer.BuildOptions{Target: targetNamed(t, target)}).(*ringBalancer) + require.NoError(t, b.UpdateClientConnState(balancer.ClientConnState{ + ResolverState: resolver.State{Addresses: addrs}, + BalancerConfig: &BalancerConfig{ReplicationFactor: 100, Spread: 1}, + })) + for _, sci := range b.subConns.Values() { + sc := sci.(balancer.SubConn) + b.UpdateSubConnState(sc, balancer.SubConnState{ConnectivityState: connectivity.Connecting}) + b.UpdateSubConnState(sc, balancer.SubConnState{ConnectivityState: connectivity.Ready}) + } + return b +} + +// RingFor returns a usable view before any balancer exists for the target +// (gRPC builds balancers lazily). The view reports ErrNoRing until a +// balancer publishes a ring. +func TestRingForBeforeBalancerExists(t *testing.T) { + bld := NewBuilder(xxhash.Sum64) + + view := bld.RingFor("test:///backends") + require.NotNil(t, view) + + _, err := view.FindN([]byte("key"), 1) + require.ErrorIs(t, err, ErrNoRing) + require.Empty(t, view.Members()) +} + +// A view taken before the balancer exists starts to work once the balancer +// publishes its ring. The view agrees with the ring the picker uses. +func TestRingForSeesTheLiveRing(t *testing.T) { + bld := NewBuilder(xxhash.Sum64) + view := bld.RingFor("test:///backends") + + b := readyBalancerForTarget(t, bld, "test:///backends", + resolver.Address{ServerName: "t", Addr: "1"}, + resolver.Address{ServerName: "t", Addr: "2"}, + ) + + require.ElementsMatch(t, []string{"t1", "t2"}, keys(view.Members())) + + for _, key := range []string{"a", "b", "c", "d"} { + got, err := view.FindN([]byte(key), 1) + require.NoError(t, err) + want, err := b.hashring.FindN([]byte(key), 1) + require.NoError(t, err) + require.Equal(t, want[0].Key(), got[0].Key(), + "view and picker must agree on the owner of %q", key) + } +} + +// A view that the caller already holds shows membership changes. +func TestRingForReflectsMembershipChanges(t *testing.T) { + bld := NewBuilder(xxhash.Sum64) + view := bld.RingFor("test:///backends") + + addrs := []resolver.Address{{ServerName: "t", Addr: "1"}, {ServerName: "t", Addr: "2"}} + b := readyBalancerForTarget(t, bld, "test:///backends", addrs...) + require.ElementsMatch(t, []string{"t1", "t2"}, keys(view.Members())) + + require.NoError(t, b.UpdateClientConnState(balancer.ClientConnState{ + ResolverState: resolver.State{Addresses: addrs[:1]}, + })) + require.ElementsMatch(t, []string{"t1"}, keys(view.Members()), + "a removed backend must disappear from held views") +} + +// A ReplicationFactor change replaces the ring object. Held views must +// follow the replacement rather than serve the stale ring. +func TestRingForFollowsRingReplacement(t *testing.T) { + bld := NewBuilder(xxhash.Sum64) + view := bld.RingFor("test:///backends") + + addrs := []resolver.Address{{ServerName: "t", Addr: "1"}, {ServerName: "t", Addr: "2"}} + b := readyBalancerForTarget(t, bld, "test:///backends", addrs...) + oldRing := b.hashring + + require.NoError(t, b.UpdateClientConnState(balancer.ClientConnState{ + ResolverState: resolver.State{Addresses: addrs}, + BalancerConfig: &BalancerConfig{ReplicationFactor: 7, Spread: 1}, + })) + require.NotSame(t, oldRing, b.hashring, "sanity: the ring was replaced") + + require.ElementsMatch(t, []string{"t1", "t2"}, keys(view.Members())) + got, err := view.FindN([]byte("key"), 1) + require.NoError(t, err) + want, err := b.hashring.FindN([]byte("key"), 1) + require.NoError(t, err) + require.Equal(t, want[0].Key(), got[0].Key(), "view must consult the replacement ring") +} + +// Closing the balancer empties its views. A view must not continue to serve +// a ring that no longer routes anything. +func TestRingForEmptyAfterClose(t *testing.T) { + bld := NewBuilder(xxhash.Sum64) + view := bld.RingFor("test:///backends") + + b := readyBalancerForTarget(t, bld, "test:///backends", + resolver.Address{ServerName: "t", Addr: "1"}) + require.NotEmpty(t, view.Members()) + + b.Close() + _, err := view.FindN([]byte("key"), 1) + require.ErrorIs(t, err, ErrNoRing) + require.Empty(t, view.Members()) +} + +// Views are per-target: two targets that the same builder serves get +// separate views. +func TestRingForIsPerTarget(t *testing.T) { + bld := NewBuilder(xxhash.Sum64) + + readyBalancerForTarget(t, bld, "test:///alpha", resolver.Address{ServerName: "a", Addr: "1"}) + readyBalancerForTarget(t, bld, "test:///beta", resolver.Address{ServerName: "b", Addr: "1"}) + + require.ElementsMatch(t, []string{"a1"}, keys(bld.RingFor("test:///alpha").Members())) + require.ElementsMatch(t, []string{"b1"}, keys(bld.RingFor("test:///beta").Members())) +}