Merge remote-tracking branch 'origin/master' into arize-dev/50458f1-from-0f3a43-in-upstream

Author: ddowker

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,16 @@ 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 }
+// TODO(whb): Zero'ing R is for backward compatibility; remove once deployment is complete.
+func (m *testMember) SetExiting()     { m.E = true; m.R = 0 }
+func (m testMember) IsExiting() bool  { return m.E }
 
 func (m *testMember) MarshalString() string {
 	if b, err := json.Marshal(m); err != nil {