allocator: coordinate graceful exit of signaled members

Allow scenarios like unattended upgrades where many members are
signaled to exit at once. Members now mark themselves as exiting
rather than zeroing their item limit, and the allocator gradually
sheds their capacity from available excess slots, oldest first,
capped so that enough members remain to satisfy replication.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: williamhbaker

allocator/alloc_state.go

@@ -1,8 +1,10 @@
 package allocator
 
 import (
+	"cmp"
 	"hash/crc64"
 	"math"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -51,6 +53,13 @@ type State struct {
 	// These share cardinality with |Members|.
 	MemberTotalCount   []int
 	MemberPrimaryCount []int
+
+	// Number of item slots to shed from each member's ItemLimit, reducing its
+	// effective capacity to ItemLimit - ShedCapacity. Exiting members are granted
+	// ShedCapacity from available excess cluster capacity, ordered by age
+	// (CreateRevision), so that the oldest exiting members drain first.
+	// Shares cardinality with |Members|.
+	ShedCapacity []int
 }
 
 // NewObservedState returns a *State instance which extracts and updates itself
@@ -89,12 +98,15 @@ func (s *State) observe() {
 	s.NetworkHash = 0
 	s.MemberTotalCount = make([]int, len(s.Members))
 	s.MemberPrimaryCount = make([]int, len(s.Members))
+	s.ShedCapacity = make([]int, len(s.Members))
 
 	// Walk Members to:
 	//  * Group the set of ordered |Zones| across all Members.
 	//  * Initialize |ZoneSlots|.
 	//  * Initialize |MemberSlots|.
 	//  * Initialize |NetworkHash|.
+	//  * Collect indices of exiting members.
+	var exiting []int
 	for i := range s.Members {
 		var m = memberAt(s.Members, i)
 		var slots = m.ItemLimit()
@@ -114,6 +126,10 @@ func (s *State) observe() {
 
 		s.MemberSlots += slots
 		s.NetworkHash = foldCRC(s.NetworkHash, s.Members[i].Raw.Key, slots)
+
+		if m.IsExiting() {
+			exiting = append(exiting, i)
+		}
 	}
 
 	// Fetch |localMember| identified by |LocalKey|.
@@ -135,11 +151,14 @@ func (s *State) observe() {
 			return strings.Compare(itemAt(s.Items, l).ID, assignmentAt(s.Assignments, r).ItemID)
 		},
 	}
+
+	var maxReplicationFactor int
 	for cur, ok := it.Next(); ok; cur, ok = it.Next() {
 		var item = itemAt(s.Items, cur.Left)
 		var slots = item.DesiredReplication()
 
 		s.ItemSlots += slots
+		maxReplicationFactor = max(maxReplicationFactor, slots)
 		s.NetworkHash = foldCRC(s.NetworkHash, s.Items[cur.Left].Raw.Key, slots)
 
 		for r := cur.RightBegin; r != cur.RightEnd; r++ {
@@ -161,11 +180,39 @@ func (s *State) observe() {
 			}
 		}
 	}
+
+	// Compute ShedCapacity for exiting members. ShedCapacity is granted to
+	// exiting members, oldest first, up to each member's ItemLimit.
+	// Excess slots available to grant as ShedCapacity. When the cluster is
+	// overloaded (MemberSlots < ItemSlots), there is no excess to shed.
+	var excessSlots = max(0, s.MemberSlots-s.ItemSlots)
+	slices.SortFunc(exiting, func(a, b int) int {
+		return cmp.Compare(s.Members[a].Raw.CreateRevision, s.Members[b].Raw.CreateRevision)
+	})
+
+	// We must retain at least maxReplicationFactor members at full capacity
+	// to satisfy replication requirements. This is a numerical bound, not
+	// zone-aware: we shed the oldest members first, expecting that replacement
+	// members will restore zone diversity as they join.
+	var maxShedding = len(s.Members) - maxReplicationFactor
+	if maxShedding < 0 {
+		maxShedding = 0
+	}
+	for n, i := range exiting {
+		if n == maxShedding || excessSlots == 0 {
+			break
+		}
+		var shed = min(memberAt(s.Members, i).ItemLimit(), excessSlots)
+		s.ShedCapacity[i] = shed
+		s.MemberSlots -= shed
+		excessSlots -= shed
+		s.NetworkHash = foldCRC(s.NetworkHash, s.Members[i].Raw.Key, shed)
+	}
 }
 
 // shouldExit returns true iff the local Member is able to safely exit.
 func (s *State) shouldExit() bool {
-	return memberAt(s.Members, s.LocalMemberInd).ItemLimit() == 0 && len(s.LocalItems) == 0
+	return memberAt(s.Members, s.LocalMemberInd).IsExiting() && len(s.LocalItems) == 0
 }
 
 // isLeader returns true iff the local Member key is ordered first on
@@ -201,16 +248,22 @@ func (s *State) debugLog() {
 	}).Info("extracted State")
 }
 
+// memberEffectiveLimit returns the effective item limit for the member at
+// index |ind|, accounting for any ShedCapacity granted to exiting members.
+func (s *State) memberEffectiveLimit(ind int) int {
+	return memberAt(s.Members, ind).ItemLimit() - s.ShedCapacity[ind]
+}
+
 // memberLoadRatio maps an |assignment| to a Member "load ratio". Given all
 // |Members| and their corresponding |counts| (1:1 with |Members|),
 // memberLoadRatio maps |assignment| to a Member and, if found, returns the
-// ratio of the Member's index in |counts| to the Member's ItemLimit. If the
-// Member is not found, infinity is returned.
+// ratio of the Member's index in |counts| to the Member's effective item
+// limit. If the Member is not found, infinity is returned.
 func (s *State) memberLoadRatio(assignment keyspace.KeyValue, counts []int) float32 {
 	var a = assignment.Decoded.(Assignment)
 
 	if ind, found := s.Members.Search(MemberKey(s.KS, a.MemberZone, a.MemberSuffix)); found {
-		return float32(counts[ind]) / float32(memberAt(s.Members, ind).ItemLimit())
+		return float32(counts[ind]) / float32(s.memberEffectiveLimit(ind))
 	}
 	return math.MaxFloat32
 }

allocator/alloc_state_test.go

@@ -109,7 +109,7 @@ func (s *AllocStateSuite) TestExitCondition(c *gc.C) {
 	buildAllocKeySpaceFixture(c, ctx, client)
 	defer etcdtest.Cleanup()
 
-	var _, err = client.Put(ctx, "/root/members/us-east#allowed-to-exit", `{"R": 0}`)
+	var _, err = client.Put(ctx, "/root/members/us-east#allowed-to-exit", `{"R": 0, "E": true}`)
 	c.Assert(err, gc.IsNil)
 
 	var ks = NewAllocatorKeySpace("/root", testAllocDecoder{})
@@ -126,7 +126,7 @@ func (s *AllocStateSuite) TestExitCondition(c *gc.C) {
 	c.Check(states[1].shouldExit(), gc.Equals, true)
 
 	// While we're at it, expect |NetworkHash| changed with the new member.
-	c.Check(states[0].NetworkHash, gc.Equals, uint64(0xfce0237931d8c200))
+	c.Check(states[0].NetworkHash, gc.Equals, uint64(0x554c9f4e9605a7a1))
 }
 
 func (s *AllocStateSuite) TestLoadRatio(c *gc.C) {
@@ -153,6 +153,204 @@ func (s *AllocStateSuite) TestLoadRatio(c *gc.C) {
 	}
 }
 
+func (s *AllocStateSuite) TestLoadRatioWithShedding(c *gc.C) {
+	var client, ctx = etcdtest.TestClient(), context.Background()
+	defer etcdtest.Cleanup()
+
+	// m1 is exiting with R:5. m2 is not exiting with R:5.
+	// 2 items at R:1. MemberSlots=10, ItemSlots=2, excessSlots=8.
+	// m1 sheds its full capacity (5), so its effective limit is 0.
+	for _, kv := range [][2]string{
+		{"/root/items/item-a", `{"R": 1}`},
+		{"/root/items/item-b", `{"R": 1}`},
+
+		{"/root/members/zone-a#m1", `{"R": 5, "E": true}`},
+		{"/root/members/zone-a#m2", `{"R": 5, "E": false}`},
+
+		{"/root/assign/item-a#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-b#zone-a#m2#0", `consistent`},
+	} {
+		var _, err = client.Put(ctx, kv[0], kv[1])
+		c.Assert(err, gc.IsNil)
+	}
+
+	var ks = NewAllocatorKeySpace("/root", testAllocDecoder{})
+	var state = NewObservedState(ks, MemberKey(ks, "zone-a", "m1"), isConsistent)
+	c.Check(ks.Load(ctx, client, 0), gc.IsNil)
+
+	c.Check(state.ShedCapacity[0], gc.Equals, 5) // m1 fully shed.
+	c.Check(state.ShedCapacity[1], gc.Equals, 0)
+
+	// m1 has 1 assignment but effective capacity 0: load ratio is +Inf.
+	c.Check(state.memberLoadRatio(state.Assignments[0], state.MemberTotalCount), gc.Equals, float32(math.Inf(1)))
+	// m2 has 1 assignment and effective capacity 5: load ratio is 1/5.
+	c.Check(state.memberLoadRatio(state.Assignments[1], state.MemberTotalCount), gc.Equals, float32(1.0/5.0))
+}
+
+func (s *AllocStateSuite) TestShedCapacityAllAtOnce(c *gc.C) {
+	var client, ctx = etcdtest.TestClient(), context.Background()
+	defer etcdtest.Cleanup()
+
+	// 5 members, capacity 5 each. MemberSlots = 25.
+	// 2 items at R:1. ItemSlots = 2. excessSlots = 23.
+	// maxReplicationFactor = 1, maxShedding = 4.
+	// Both exiting members shed their full capacity.
+	for _, kv := range [][2]string{
+		{"/root/items/item-a", `{"R": 1}`},
+		{"/root/items/item-b", `{"R": 1}`},
+
+		{"/root/members/zone-a#m1", `{"R": 5, "E": true}`},
+		{"/root/members/zone-a#m2", `{"R": 5, "E": true}`},
+		{"/root/members/zone-a#m3", `{"R": 5, "E": false}`},
+		{"/root/members/zone-b#m4", `{"R": 5, "E": false}`},
+		{"/root/members/zone-b#m5", `{"R": 5, "E": false}`},
+
+		{"/root/assign/item-a#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-b#zone-b#m4#0", `consistent`},
+	} {
+		var _, err = client.Put(ctx, kv[0], kv[1])
+		c.Assert(err, gc.IsNil)
+	}
+
+	var ks = NewAllocatorKeySpace("/root", testAllocDecoder{})
+	var state = NewObservedState(ks, MemberKey(ks, "zone-a", "m1"), isConsistent)
+	c.Check(ks.Load(ctx, client, 0), gc.IsNil)
+
+	c.Check(state.ShedCapacity[0], gc.Equals, 5) // m1 sheds fully.
+	c.Check(state.ShedCapacity[1], gc.Equals, 5) // m2 sheds fully.
+	c.Check(state.ShedCapacity[2], gc.Equals, 0)
+	c.Check(state.ShedCapacity[3], gc.Equals, 0)
+	c.Check(state.ShedCapacity[4], gc.Equals, 0)
+}
+
+func (s *AllocStateSuite) TestShedCapacityReplicationConstrained(c *gc.C) {
+	var client, ctx = etcdtest.TestClient(), context.Background()
+	defer etcdtest.Cleanup()
+
+	// 4 members, capacity 10 each. MemberSlots = 40.
+	// 3 items: R:3, R:2, R:1. ItemSlots = 6. excessSlots = 34.
+	// maxReplicationFactor = 3, maxShedding = 1.
+	// m1 and m2 are both exiting, but only the oldest (m1) sheds.
+	for _, kv := range [][2]string{
+		{"/root/items/item-a", `{"R": 3}`},
+		{"/root/items/item-b", `{"R": 2}`},
+		{"/root/items/item-c", `{"R": 1}`},
+
+		{"/root/members/zone-a#m1", `{"R": 10, "E": true}`},
+		{"/root/members/zone-a#m2", `{"R": 10, "E": true}`},
+		{"/root/members/zone-b#m3", `{"R": 10, "E": false}`},
+		{"/root/members/zone-b#m4", `{"R": 10, "E": false}`},
+
+		{"/root/assign/item-a#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-a#zone-a#m2#1", `consistent`},
+		{"/root/assign/item-a#zone-b#m3#2", `consistent`},
+		{"/root/assign/item-b#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-b#zone-b#m4#1", `consistent`},
+		{"/root/assign/item-c#zone-b#m3#0", `consistent`},
+	} {
+		var _, err = client.Put(ctx, kv[0], kv[1])
+		c.Assert(err, gc.IsNil)
+	}
+
+	var ks = NewAllocatorKeySpace("/root", testAllocDecoder{})
+	var state = NewObservedState(ks, MemberKey(ks, "zone-a", "m1"), isConsistent)
+	c.Check(ks.Load(ctx, client, 0), gc.IsNil)
+
+	c.Check(state.ShedCapacity[0], gc.Equals, 10) // m1 sheds fully (oldest).
+	c.Check(state.ShedCapacity[1], gc.Equals, 0)  // m2 blocked by maxShedding.
+	c.Check(state.ShedCapacity[2], gc.Equals, 0)
+	c.Check(state.ShedCapacity[3], gc.Equals, 0)
+}
+
+func (s *AllocStateSuite) TestShedCapacityExcessConstrained(c *gc.C) {
+	var client, ctx = etcdtest.TestClient(), context.Background()
+	defer etcdtest.Cleanup()
+
+	// 4 members, capacity 3 each. MemberSlots = 12.
+	// 5 items at R:1. ItemSlots = 5. excessSlots = 7.
+	// maxReplicationFactor = 1, maxShedding = 3.
+	// 3 exiting members need 9 total shed, but only 7 excess available.
+	// m1: sheds 3 (full), m2: sheds 3 (full), m3: sheds 1 (partial).
+	for _, kv := range [][2]string{
+		{"/root/items/item-a", `{"R": 1}`},
+		{"/root/items/item-b", `{"R": 1}`},
+		{"/root/items/item-c", `{"R": 1}`},
+		{"/root/items/item-d", `{"R": 1}`},
+		{"/root/items/item-e", `{"R": 1}`},
+
+		{"/root/members/zone-a#m1", `{"R": 3, "E": true}`},
+		{"/root/members/zone-a#m2", `{"R": 3, "E": true}`},
+		{"/root/members/zone-b#m3", `{"R": 3, "E": true}`},
+		{"/root/members/zone-b#m4", `{"R": 3, "E": false}`},
+
+		{"/root/assign/item-a#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-b#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-c#zone-a#m2#0", `consistent`},
+		{"/root/assign/item-d#zone-b#m3#0", `consistent`},
+		{"/root/assign/item-e#zone-b#m4#0", `consistent`},
+	} {
+		var _, err = client.Put(ctx, kv[0], kv[1])
+		c.Assert(err, gc.IsNil)
+	}
+
+	var ks = NewAllocatorKeySpace("/root", testAllocDecoder{})
+	var state = NewObservedState(ks, MemberKey(ks, "zone-a", "m1"), isConsistent)
+	c.Check(ks.Load(ctx, client, 0), gc.IsNil)
+
+	c.Check(state.ShedCapacity[0], gc.Equals, 3) // m1 sheds fully.
+	c.Check(state.ShedCapacity[1], gc.Equals, 3) // m2 sheds fully.
+	c.Check(state.ShedCapacity[2], gc.Equals, 1) // m3 gets remaining excess.
+	c.Check(state.ShedCapacity[3], gc.Equals, 0) // non-exiting.
+}
+
+func (s *AllocStateSuite) TestShedCapacityOverloadedCluster(c *gc.C) {
+	var client, ctx = etcdtest.TestClient(), context.Background()
+	defer etcdtest.Cleanup()
+
+	// 2 members, capacity 2 each. MemberSlots = 4.
+	// 5 items at R:1. ItemSlots = 5. excessSlots = -1 (clamped to 0).
+	// No shedding is possible because the cluster is overloaded.
+	for _, kv := range [][2]string{
+		{"/root/items/item-a", `{"R": 1}`},
+		{"/root/items/item-b", `{"R": 1}`},
+		{"/root/items/item-c", `{"R": 1}`},
+		{"/root/items/item-d", `{"R": 1}`},
+		{"/root/items/item-e", `{"R": 1}`},
+
+		{"/root/members/zone-a#m1", `{"R": 2, "E": true}`},
+		{"/root/members/zone-b#m2", `{"R": 2, "E": false}`},
+
+		{"/root/assign/item-a#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-b#zone-a#m1#0", `consistent`},
+		{"/root/assign/item-c#zone-b#m2#0", `consistent`},
+		{"/root/assign/item-d#zone-b#m2#0", `consistent`},
+	} {
+		var _, err = client.Put(ctx, kv[0], kv[1])
+		c.Assert(err, gc.IsNil)
+	}
+
+	var ks = NewAllocatorKeySpace("/root", testAllocDecoder{})
+	var state = NewObservedState(ks, MemberKey(ks, "zone-a", "m1"), isConsistent)
+	c.Check(ks.Load(ctx, client, 0), gc.IsNil)
+
+	c.Check(state.ShedCapacity[0], gc.Equals, 0) // m1: no excess to shed.
+	c.Check(state.ShedCapacity[1], gc.Equals, 0)
+}
+
+func (s *AllocStateSuite) TestShedCapacityNoExiting(c *gc.C) {
+	var client, ctx = etcdtest.TestClient(), context.Background()
+	defer etcdtest.Cleanup()
+	buildAllocKeySpaceFixture(c, ctx, client)
+
+	var ks = NewAllocatorKeySpace("/root", testAllocDecoder{})
+	var state = NewObservedState(ks, MemberKey(ks, "us-west", "baz"), isConsistent)
+	c.Check(ks.Load(ctx, client, 0), gc.IsNil)
+
+	for i := range state.ShedCapacity {
+		c.Check(state.ShedCapacity[i], gc.Equals, 0)
+	}
+}
+
 var _ = gc.Suite(&AllocStateSuite{})
 
 func TestMain(m *testing.M) { etcdtest.TestMainWithEtcd(m) }

allocator/allocator_key_space.go

@@ -34,6 +34,8 @@ const (
 type MemberValue interface {
 	// ItemLimit is the maximum number of Items this Member may be assigned.
 	ItemLimit() int
+	// IsExiting returns true if this Member has been signaled to exit.
+	IsExiting() bool
 }
 
 // ItemValue is a user-defined Item representation which also supports required

allocator/allocator_key_space_test.go

@@ -354,11 +354,15 @@ func isConsistent(_ Item, assignment keyspace.KeyValue, allAssignments keyspace.
 	return assignment.Decoded.(Assignment).AssignmentValue.(testAssignment).consistent
 }
 
-type testMember struct{ R int }
+type testMember struct {
+	R int  // ItemLimit
+	E bool // Exiting
+}
 
-func (m testMember) ItemLimit() int  { return m.R }
-func (m testMember) Validate() error { return nil }
-func (m *testMember) ZeroLimit()     { m.R = 0 }
+func (m testMember) ItemLimit() int   { return m.R }
+func (m testMember) Validate() error  { return nil }
+func (m *testMember) SetExiting()     { m.E = true }
+func (m testMember) IsExiting() bool  { return m.E }
 
 func (m *testMember) MarshalString() string {
 	if b, err := json.Marshal(m); err != nil {