Use uniform pull gossip

ava-labs · Nov 28, 2023 · 618ef67 · 618ef67
1 parent 6b79747
commit 618ef67
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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()