From b944e45db37390728e9d119c31ef13091dacfb91 Mon Sep 17 00:00:00 2001 From: ecordell Date: Fri, 25 Sep 2026 13:15:39 -0400 Subject: [PATCH] feat: expose read-only views of live hashrings via Builder.RingFor Applications that route work by hashring key need to know which member owns a key before they dial. One example is a batch: to group the items by owner, the caller must query the ring. The ring lives inside the balancer, and gRPC constructs balancers lazily and internally, so callers had no way to reach it. The builder is the rendezvous point. RingFor(target) returns a stable RingView handle with these properties: - Callers can get a view before any balancer exists for the target. - A view returns ErrNoRing until its balancer publishes a ring. - A view follows ring replacements (ReplicationFactor changes). - A view shows membership changes as they happen. - A view empties again when its balancer closes. A view reads the same ring object that the picker routes with, so callers and the picker agree on ownership. The one exception is a membership change that races the query. --- balancer.go | 56 +++++++++++++++++-- ringview.go | 59 +++++++++++++++++++ ringview_test.go | 143 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 253 insertions(+), 5 deletions(-) create mode 100644 ringview.go create mode 100644 ringview_test.go 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())) +}