Support all validator set diffs by height

node/overridden_manager.go

@@ -75,6 +75,10 @@ func (o *overriddenManager) Sample(_ ids.ID, size int) ([]ids.NodeID, error) {
 	return o.manager.Sample(o.subnetID, size)
 }
 
+func (o *overriddenManager) GetAllMaps() map[ids.ID]map[ids.NodeID]*validators.GetValidatorOutput {
+	return o.manager.GetAllMaps()
+}
+
 func (o *overriddenManager) GetMap(ids.ID) map[ids.NodeID]*validators.GetValidatorOutput {
 	return o.manager.GetMap(o.subnetID)
 }

snow/validators/manager.go

@@ -94,6 +94,9 @@ type Manager interface {
 	// If sampling the requested size isn't possible, an error will be returned.
 	Sample(subnetID ids.ID, size int) ([]ids.NodeID, error)
 
+	// GetAllMaps returns a copy of all validators of all subnets
+	GetAllMaps() map[ids.ID]map[ids.NodeID]*GetValidatorOutput
+
 	// Map of the validators in this subnet
 	GetMap(subnetID ids.ID) map[ids.NodeID]*GetValidatorOutput
 
@@ -257,6 +260,17 @@ func (m *manager) Sample(subnetID ids.ID, size int) ([]ids.NodeID, error) {
 	return set.Sample(size)
 }
 
+func (m *manager) GetAllMaps() map[ids.ID]map[ids.NodeID]*GetValidatorOutput {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+
+	set := make(map[ids.ID]map[ids.NodeID]*GetValidatorOutput, len(m.subnetToVdrs))
+	for subnetID, vdrs := range m.subnetToVdrs {
+		set[subnetID] = vdrs.Map()
+	}
+	return set
+}
+
 func (m *manager) GetMap(subnetID ids.ID) map[ids.NodeID]*GetValidatorOutput {
 	m.lock.RLock()
 	set, exists := m.subnetToVdrs[subnetID]

snow/validators/manager_test.go

@@ -375,6 +375,100 @@ func TestGetMap(t *testing.T) {
 	require.Empty(m.GetMap(subnetID))
 }
 
+func TestGetAllMaps(t *testing.T) {
+	require := require.New(t)
+
+	m := NewManager()
+	subnetID0 := ids.GenerateTestID()
+	subnetID1 := ids.GenerateTestID()
+
+	maps := m.GetAllMaps()
+	require.Empty(maps)
+
+	sk, err := localsigner.New()
+	require.NoError(err)
+
+	pk := sk.PublicKey()
+	nodeID0 := ids.GenerateTestNodeID()
+	require.NoError(m.AddStaker(subnetID0, nodeID0, pk, ids.Empty, 2))
+
+	maps = m.GetAllMaps()
+	require.Len(maps, 1)
+	map0, ok := maps[subnetID0]
+	require.True(ok)
+	require.Len(map0, 1)
+	require.Contains(map0, nodeID0)
+
+	node0 := map0[nodeID0]
+	require.Equal(nodeID0, node0.NodeID)
+	require.Equal(pk, node0.PublicKey)
+	require.Equal(uint64(2), node0.Weight)
+
+	nodeID1 := ids.GenerateTestNodeID()
+	require.NoError(m.AddStaker(subnetID1, nodeID1, nil, ids.Empty, 1))
+
+	maps = m.GetAllMaps()
+	require.Len(maps, 2)
+	map0, ok = maps[subnetID0]
+	require.True(ok)
+	require.Contains(map0, nodeID0)
+	map1, ok := maps[subnetID1]
+	require.True(ok)
+	require.Contains(map1, nodeID1)
+
+	node0 = map0[nodeID0]
+	require.Equal(nodeID0, node0.NodeID)
+	require.Equal(pk, node0.PublicKey)
+	require.Equal(uint64(2), node0.Weight)
+
+	node1 := map1[nodeID1]
+	require.Equal(nodeID1, node1.NodeID)
+	require.Nil(node1.PublicKey)
+	require.Equal(uint64(1), node1.Weight)
+
+	require.NoError(m.RemoveWeight(subnetID0, nodeID0, 1))
+	require.Equal(nodeID0, node0.NodeID)
+	require.Equal(pk, node0.PublicKey)
+	require.Equal(uint64(2), node0.Weight)
+
+	maps = m.GetAllMaps()
+	require.Len(maps, 2)
+	map0, ok = maps[subnetID0]
+	require.True(ok)
+	map1, ok = maps[subnetID1]
+	require.True(ok)
+	require.Contains(map0, nodeID0)
+	require.Contains(map1, nodeID1)
+
+	node0 = map0[nodeID0]
+	require.Equal(nodeID0, node0.NodeID)
+	require.Equal(pk, node0.PublicKey)
+	require.Equal(uint64(1), node0.Weight)
+
+	node1 = map1[nodeID1]
+	require.Equal(nodeID1, node1.NodeID)
+	require.Nil(node1.PublicKey)
+	require.Equal(uint64(1), node1.Weight)
+
+	require.NoError(m.RemoveWeight(subnetID0, nodeID0, 1))
+
+	maps = m.GetAllMaps()
+	require.Len(maps, 1)
+	require.NotContains(maps, subnetID0)
+	map1, ok = maps[subnetID1]
+	require.True(ok)
+	require.Contains(map1, nodeID1)
+
+	node1 = map1[nodeID1]
+	require.Equal(nodeID1, node1.NodeID)
+	require.Nil(node1.PublicKey)
+	require.Equal(uint64(1), node1.Weight)
+
+	require.NoError(m.RemoveWeight(subnetID1, nodeID1, 1))
+
+	require.Empty(m.GetAllMaps())
+}
+
 func TestWeight(t *testing.T) {
 	require := require.New(t)
 

vms/platformvm/state/disk_staker_diff_iterator.go

@@ -28,26 +28,44 @@ var (
 	errUnexpectedWeightValueLength = fmt.Errorf("expected weight value length %d", weightValueLength)
 )
 
-// marshalStartDiffKey is used to determine the starting key when iterating.
+// marshalStartDiffKeyBySubnetID is used to determine the starting key when iterating.
 //
-// Invariant: the result is a prefix of [marshalDiffKey] when called with the
+// Invariant: the result is a prefix of [marshalDiffKeyBySubnetID] when called with the
 // same arguments.
-func marshalStartDiffKey(subnetID ids.ID, height uint64) []byte {
+func marshalStartDiffKeyBySubnetID(subnetID ids.ID, height uint64) []byte {
 	key := make([]byte, startDiffKeyLength)
 	copy(key, subnetID[:])
 	packIterableHeight(key[ids.IDLen:], height)
 	return key
 }
 
-func marshalDiffKey(subnetID ids.ID, height uint64, nodeID ids.NodeID) []byte {
+// marshalStartDiffKeyByHeight is used to determine the starting key when iterating.
+//
+// Invariant: the result is a prefix of [marshalDiffKeyByHeight] when called with the
+// same arguments.
+func marshalStartDiffKeyByHeight(height uint64) []byte {
+	key := make([]byte, database.Uint64Size)
+	packIterableHeight(key, height)
+	return key
+}
+
+func marshalDiffKeyBySubnetID(subnetID ids.ID, height uint64, nodeID ids.NodeID) []byte {
 	key := make([]byte, diffKeyLength)
 	copy(key, subnetID[:])
 	packIterableHeight(key[ids.IDLen:], height)
 	copy(key[diffKeyNodeIDOffset:], nodeID.Bytes())
 	return key
 }
 
-func unmarshalDiffKey(key []byte) (ids.ID, uint64, ids.NodeID, error) {
+func marshalDiffKeyByHeight(height uint64, subnetID ids.ID, nodeID ids.NodeID) []byte {
+	key := make([]byte, diffKeyLength)
+	packIterableHeight(key, height)
+	copy(key[database.Uint64Size:], subnetID[:])
+	copy(key[diffKeyNodeIDOffset:], nodeID.Bytes())
+	return key
+}
+
+func unmarshalDiffKeyBySubnetID(key []byte) (ids.ID, uint64, ids.NodeID, error) {
 	if len(key) != diffKeyLength {
 		return ids.Empty, 0, ids.EmptyNodeID, errUnexpectedDiffKeyLength
 	}
@@ -61,6 +79,20 @@ func unmarshalDiffKey(key []byte) (ids.ID, uint64, ids.NodeID, error) {
 	return subnetID, height, nodeID, nil
 }
 
+func unmarshalDiffKeyByHeight(key []byte) (uint64, ids.ID, ids.NodeID, error) {
+	if len(key) != diffKeyLength {
+		return 0, ids.Empty, ids.EmptyNodeID, errUnexpectedDiffKeyLength
+	}
+	var (
+		subnetID ids.ID
+		nodeID   ids.NodeID
+	)
+	height := unpackIterableHeight(key)
+	copy(subnetID[:], key[database.Uint64Size:])
+	copy(nodeID[:], key[diffKeyNodeIDOffset:])
+	return height, subnetID, nodeID, nil
+}
+
 func marshalWeightDiff(diff *ValidatorWeightDiff) []byte {
 	value := make([]byte, weightValueLength)
 	if diff.Decrease {

vms/platformvm/state/disk_staker_diff_iterator_test.go

@@ -13,7 +13,7 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 )
 
-func FuzzMarshalDiffKey(f *testing.F) {
+func FuzzMarshalDiffKeyBySubnetID(f *testing.F) {
 	f.Fuzz(func(t *testing.T, data []byte) {
 		require := require.New(t)
 
@@ -25,31 +25,67 @@ func FuzzMarshalDiffKey(f *testing.F) {
 		fz := fuzzer.NewFuzzer(data)
 		fz.Fill(&subnetID, &height, &nodeID)
 
-		key := marshalDiffKey(subnetID, height, nodeID)
-		parsedSubnetID, parsedHeight, parsedNodeID, err := unmarshalDiffKey(key)
+		key := marshalDiffKeyBySubnetID(subnetID, height, nodeID)
+		parsedSubnetID, parsedHeight, parsedNodeID, err := unmarshalDiffKeyBySubnetID(key)
 		require.NoError(err)
 		require.Equal(subnetID, parsedSubnetID)
 		require.Equal(height, parsedHeight)
 		require.Equal(nodeID, parsedNodeID)
 	})
 }
 
-func FuzzUnmarshalDiffKey(f *testing.F) {
+func FuzzMarshalDiffKeyByHeight(f *testing.F) {
+	f.Fuzz(func(t *testing.T, data []byte) {
+		require := require.New(t)
+
+		var (
+			subnetID ids.ID
+			height   uint64
+			nodeID   ids.NodeID
+		)
+		fz := fuzzer.NewFuzzer(data)
+		fz.Fill(&height, &subnetID, &nodeID)
+
+		key := marshalDiffKeyByHeight(height, subnetID, nodeID)
+		parsedHeight, parsedSubnetID, parsedNodeID, err := unmarshalDiffKeyByHeight(key)
+		require.NoError(err)
+		require.Equal(subnetID, parsedSubnetID)
+		require.Equal(height, parsedHeight)
+		require.Equal(nodeID, parsedNodeID)
+	})
+}
+
+func FuzzUnmarshalDiffKeyBySubnetID(f *testing.F) {
 	f.Fuzz(func(t *testing.T, key []byte) {
 		require := require.New(t)
 
-		subnetID, height, nodeID, err := unmarshalDiffKey(key)
+		subnetID, height, nodeID, err := unmarshalDiffKeyBySubnetID(key)
 		if err != nil {
 			require.ErrorIs(err, errUnexpectedDiffKeyLength)
 			return
 		}
 
-		formattedKey := marshalDiffKey(subnetID, height, nodeID)
+		formattedKey := marshalDiffKeyBySubnetID(subnetID, height, nodeID)
 		require.Equal(key, formattedKey)
 	})
 }
 
-func TestDiffIteration(t *testing.T) {
+func FuzzUnmarshalDiffKeyByHeight(f *testing.F) {
+	f.Fuzz(func(t *testing.T, key []byte) {
+		require := require.New(t)
+
+		subnetID, height, nodeID, err := unmarshalDiffKeyByHeight(key)
+		if err != nil {
+			require.ErrorIs(err, errUnexpectedDiffKeyLength)
+			return
+		}
+
+		formattedKey := marshalDiffKeyByHeight(subnetID, height, nodeID)
+		require.Equal(key, formattedKey)
+	})
+}
+
+func TestDiffIterationBySubnetID(t *testing.T) {
 	require := require.New(t)
 
 	db := memdb.New()
@@ -60,13 +96,13 @@ func TestDiffIteration(t *testing.T) {
 	nodeID0 := ids.BuildTestNodeID([]byte{0x00})
 	nodeID1 := ids.BuildTestNodeID([]byte{0x01})
 
-	subnetID0Height0NodeID0 := marshalDiffKey(subnetID0, 0, nodeID0)
-	subnetID0Height1NodeID0 := marshalDiffKey(subnetID0, 1, nodeID0)
-	subnetID0Height1NodeID1 := marshalDiffKey(subnetID0, 1, nodeID1)
+	subnetID0Height0NodeID0 := marshalDiffKeyBySubnetID(subnetID0, 0, nodeID0)
+	subnetID0Height1NodeID0 := marshalDiffKeyBySubnetID(subnetID0, 1, nodeID0)
+	subnetID0Height1NodeID1 := marshalDiffKeyBySubnetID(subnetID0, 1, nodeID1)
 
-	subnetID1Height0NodeID0 := marshalDiffKey(subnetID1, 0, nodeID0)
-	subnetID1Height1NodeID0 := marshalDiffKey(subnetID1, 1, nodeID0)
-	subnetID1Height1NodeID1 := marshalDiffKey(subnetID1, 1, nodeID1)
+	subnetID1Height0NodeID0 := marshalDiffKeyBySubnetID(subnetID1, 0, nodeID0)
+	subnetID1Height1NodeID0 := marshalDiffKeyBySubnetID(subnetID1, 1, nodeID0)
+	subnetID1Height1NodeID1 := marshalDiffKeyBySubnetID(subnetID1, 1, nodeID1)
 
 	require.NoError(db.Put(subnetID0Height0NodeID0, nil))
 	require.NoError(db.Put(subnetID0Height1NodeID0, nil))
@@ -76,7 +112,7 @@ func TestDiffIteration(t *testing.T) {
 	require.NoError(db.Put(subnetID1Height1NodeID1, nil))
 
 	{
-		it := db.NewIteratorWithStartAndPrefix(marshalStartDiffKey(subnetID0, 0), subnetID0[:])
+		it := db.NewIteratorWithStartAndPrefix(marshalStartDiffKeyBySubnetID(subnetID0, 0), subnetID0[:])
 		defer it.Release()
 
 		expectedKeys := [][]byte{
@@ -91,7 +127,7 @@ func TestDiffIteration(t *testing.T) {
 	}
 
 	{
-		it := db.NewIteratorWithStartAndPrefix(marshalStartDiffKey(subnetID0, 1), subnetID0[:])
+		it := db.NewIteratorWithStartAndPrefix(marshalStartDiffKeyBySubnetID(subnetID0, 1), subnetID0[:])
 		defer it.Release()
 
 		expectedKeys := [][]byte{
@@ -107,3 +143,66 @@ func TestDiffIteration(t *testing.T) {
 		require.NoError(it.Error())
 	}
 }
+
+func TestDiffIterationByHeight(t *testing.T) {
+	require := require.New(t)
+
+	db := memdb.New()
+
+	subnetID0 := ids.ID{0x00}
+	subnetID1 := ids.ID{0x01}
+
+	nodeID0 := ids.BuildTestNodeID([]byte{0x00})
+	nodeID1 := ids.BuildTestNodeID([]byte{0x01})
+
+	height0SubnetID0NodeID0 := marshalDiffKeyByHeight(0, subnetID0, nodeID0)
+	height1SubnetID0NodeID0 := marshalDiffKeyByHeight(1, subnetID0, nodeID0)
+	height1SubnetID0NodeID1 := marshalDiffKeyByHeight(1, subnetID0, nodeID1)
+
+	height0SubnetID1NodeID0 := marshalDiffKeyByHeight(0, subnetID1, nodeID0)
+	height1SubnetID1NodeID0 := marshalDiffKeyByHeight(1, subnetID1, nodeID0)
+	height1SubnetID1NodeID1 := marshalDiffKeyByHeight(1, subnetID1, nodeID1)
+
+	require.NoError(db.Put(height0SubnetID0NodeID0, nil))
+	require.NoError(db.Put(height1SubnetID0NodeID0, nil))
+	require.NoError(db.Put(height1SubnetID0NodeID1, nil))
+	require.NoError(db.Put(height0SubnetID1NodeID0, nil))
+	require.NoError(db.Put(height1SubnetID1NodeID0, nil))
+	require.NoError(db.Put(height1SubnetID1NodeID1, nil))
+
+	{
+		it := db.NewIteratorWithStart(marshalStartDiffKeyByHeight(0))
+		defer it.Release()
+
+		expectedKeys := [][]byte{
+			height0SubnetID0NodeID0,
+			height0SubnetID1NodeID0,
+		}
+		for _, expectedKey := range expectedKeys {
+			require.True(it.Next())
+			require.Equal(expectedKey, it.Key())
+		}
+		require.False(it.Next())
+		require.NoError(it.Error())
+	}
+
+	{
+		it := db.NewIteratorWithStart(marshalStartDiffKeyByHeight(1))
+		defer it.Release()
+
+		expectedKeys := [][]byte{
+			height1SubnetID0NodeID0,
+			height1SubnetID0NodeID1,
+			height1SubnetID1NodeID0,
+			height1SubnetID1NodeID1,
+			height0SubnetID0NodeID0,
+			height0SubnetID1NodeID0,
+		}
+		for _, expectedKey := range expectedKeys {
+			require.True(it.Next())
+			require.Equal(expectedKey, it.Key())
+		}
+		require.False(it.Next())
+		require.NoError(it.Error())
+	}
+}