From 618ef67630ca549b7f6e5a23bc133a121da08652 Mon Sep 17 00:00:00 2001
From: Stephen Buttolph <stephen@avalabs.org>
Date: Tue, 28 Nov 2023 16:52:36 -0500
Subject: [PATCH] Use uniform pull gossip

---
 node/overridden_manager.go        |  4 ++++
 snow/engine/snowman/transitive.go |  2 +-
 snow/validators/manager.go        | 19 +++++++++++++++++++
 snow/validators/set.go            | 23 +++++++++++++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/node/overridden_manager.go b/node/overridden_manager.go
index 91d8c198a4c3..80295f8636ea 100644
--- a/node/overridden_manager.go
+++ b/node/overridden_manager.go
@@ -68,6 +68,10 @@ func (o *overriddenManager) Sample(_ ids.ID, size int) ([]ids.NodeID, error) {
 	return o.manager.Sample(o.subnetID, size)
 }
 
+func (o *overriddenManager) UniformSample(_ ids.ID, size int) ([]ids.NodeID, error) {
+	return o.manager.UniformSample(o.subnetID, size)
+}
+
 func (o *overriddenManager) GetMap(ids.ID) map[ids.NodeID]*validators.GetValidatorOutput {
 	return o.manager.GetMap(o.subnetID)
 }
diff --git a/snow/engine/snowman/transitive.go b/snow/engine/snowman/transitive.go
index 1e4f891e5604..bd0a0aeedb01 100644
--- a/snow/engine/snowman/transitive.go
+++ b/snow/engine/snowman/transitive.go
@@ -163,7 +163,7 @@ func (t *Transitive) Gossip(ctx context.Context) error {
 			zap.Stringer("validators", t.Validators),
 		)
 
-		vdrIDs, err := t.Validators.Sample(t.Ctx.SubnetID, 1)
+		vdrIDs, err := t.Validators.UniformSample(t.Ctx.SubnetID, 1)
 		if err != nil {
 			t.Ctx.Log.Error("skipping block gossip",
 				zap.String("reason", "no validators"),
diff --git a/snow/validators/manager.go b/snow/validators/manager.go
index c42ea779d96b..8cf634f29bd7 100644
--- a/snow/validators/manager.go
+++ b/snow/validators/manager.go
@@ -85,6 +85,10 @@ 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)
 
+	// UniformSample returns a collection of validatorIDs in the subnet.
+	// If sampling the requested size isn't possible, an error will be returned.
+	UniformSample(subnetID ids.ID, size int) ([]ids.NodeID, error)
+
 	// Map of the validators in this subnet
 	GetMap(subnetID ids.ID) map[ids.NodeID]*GetValidatorOutput
 
@@ -253,6 +257,21 @@ func (m *manager) Sample(subnetID ids.ID, size int) ([]ids.NodeID, error) {
 	return set.Sample(size)
 }
 
+func (m *manager) UniformSample(subnetID ids.ID, size int) ([]ids.NodeID, error) {
+	if size == 0 {
+		return nil, nil
+	}
+
+	m.lock.RLock()
+	set, exists := m.subnetToVdrs[subnetID]
+	m.lock.RUnlock()
+	if !exists {
+		return nil, ErrMissingValidators
+	}
+
+	return set.UniformSample(size)
+}
+
 func (m *manager) GetMap(subnetID ids.ID) map[ids.NodeID]*GetValidatorOutput {
 	m.lock.RLock()
 	set, exists := m.subnetToVdrs[subnetID]
diff --git a/snow/validators/set.go b/snow/validators/set.go
index dfa294a70bbe..564cd107153a 100644
--- a/snow/validators/set.go
+++ b/snow/validators/set.go
@@ -243,6 +243,13 @@ func (s *vdrSet) Sample(size int) ([]ids.NodeID, error) {
 	return s.sample(size)
 }
 
+func (s *vdrSet) UniformSample(size int) ([]ids.NodeID, error) {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+
+	return s.uniformSample(size)
+}
+
 func (s *vdrSet) sample(size int) ([]ids.NodeID, error) {
 	if !s.samplerInitialized {
 		if err := s.sampler.Initialize(s.weights); err != nil {
@@ -263,6 +270,22 @@ func (s *vdrSet) sample(size int) ([]ids.NodeID, error) {
 	return list, nil
 }
 
+func (s *vdrSet) uniformSample(size int) ([]ids.NodeID, error) {
+	uniform := sampler.NewUniform()
+	uniform.Initialize(uint64(len(s.vdrSlice)))
+
+	indices, err := uniform.Sample(size)
+	if err != nil {
+		return nil, err
+	}
+
+	list := make([]ids.NodeID, size)
+	for i, index := range indices {
+		list[i] = s.vdrSlice[index].NodeID
+	}
+	return list, nil
+}
+
 func (s *vdrSet) TotalWeight() (uint64, error) {
 	s.lock.RLock()
 	defer s.lock.RUnlock()