From 31ce7088970410306c4edf546bec2d09a7176af9 Mon Sep 17 00:00:00 2001
From: Yacov Manevich <yacov.manevich@avalabs.org>
Date: Thu, 29 Aug 2024 00:40:53 +0200
Subject: [PATCH] Abort query if not enough stake connected

This commit makes the snowman engine to abort early and avoid sending a query
if the node considers the connected stake as insufficient to achieve a successful poll.

The rational for this, is to prevent the node to try to agree on blocks when it is in a partition.
The intent here is twofold: (1) To prevent needless resource expenditure,
and (2) to prevent snowball from increasing its preference strength while the protocol has no way to make progress.
The higher preference strength in a partition, the longer it takes to recover and achieve consensus once the partition heals.

Signed-off-by: Yacov Manevich <yacov.manevich@avalabs.org>
---
 go.mod                             |  1 +
 snow/engine/snowman/engine.go      | 33 +++++++++++++++++++++++--
 snow/engine/snowman/engine_test.go | 39 ++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index 30476015824e..c54e99dc31a2 100644
--- a/go.mod
+++ b/go.mod
@@ -33,6 +33,7 @@ require (
 	github.com/mr-tron/base58 v1.2.0
 	github.com/nbutton23/zxcvbn-go v0.0.0-20180912185939-ae427f1e4c1d
 	github.com/onsi/ginkgo/v2 v2.13.1
+	github.com/onsi/gomega v1.29.0
 	github.com/pires/go-proxyproto v0.6.2
 	github.com/prometheus/client_golang v1.16.0
 	github.com/prometheus/client_model v0.3.0
diff --git a/snow/engine/snowman/engine.go b/snow/engine/snowman/engine.go
index 1b2cfb5c11d4..1125f86cd6c2 100644
--- a/snow/engine/snowman/engine.go
+++ b/snow/engine/snowman/engine.go
@@ -31,7 +31,10 @@ import (
 	"github.com/ava-labs/avalanchego/utils/units"
 )
 
-const nonVerifiedCacheSize = 64 * units.MiB
+const (
+	nonVerifiedCacheSize = 64 * units.MiB
+	errInsufficientStake = "insufficient connected stake"
+)
 
 var _ common.Engine = (*Engine)(nil)
 
@@ -190,12 +193,19 @@ func (e *Engine) Gossip(ctx context.Context) error {
 		return nil
 	}
 
+	blkID := e.Consensus.Preference()
+
+	if e.abortDueToInsufficientConnectedStake(blkID) {
+		return nil
+	}
+
 	e.requestID++
+
 	e.Sender.SendPullQuery(
 		ctx,
 		set.Of(vdrID),
 		e.requestID,
-		e.Consensus.Preference(),
+		blkID,
 		nextHeightToAccept,
 	)
 	return nil
@@ -875,6 +885,10 @@ func (e *Engine) sendQuery(
 		zap.Stringer("validators", e.Validators),
 	)
 
+	if e.abortDueToInsufficientConnectedStake(blkID) {
+		return
+	}
+
 	vdrIDs, err := e.Validators.Sample(e.Ctx.SubnetID, e.Params.K)
 	if err != nil {
 		e.Ctx.Log.Warn("dropped query for block",
@@ -916,6 +930,21 @@ func (e *Engine) sendQuery(
 	}
 }
 
+func (e *Engine) abortDueToInsufficientConnectedStake(blkID ids.ID) bool {
+	stakeConnectedRatio := e.Config.ConnectedValidators.ConnectedPercent()
+	minConnectedStakeToQuery := float64(e.Params.AlphaConfidence) / float64(e.Params.K)
+
+	if stakeConnectedRatio < minConnectedStakeToQuery {
+		e.Ctx.Log.Debug("dropped query for block",
+			zap.String("reason", errInsufficientStake),
+			zap.Stringer("blkID", blkID),
+			zap.Float64("ratio", stakeConnectedRatio),
+		)
+		return true
+	}
+	return false
+}
+
 // issue [blk] to consensus
 // If [push] is true, a push query will be used. Otherwise, a pull query will be
 // used.
diff --git a/snow/engine/snowman/engine_test.go b/snow/engine/snowman/engine_test.go
index 26eccf232ed1..73227b14250e 100644
--- a/snow/engine/snowman/engine_test.go
+++ b/snow/engine/snowman/engine_test.go
@@ -22,6 +22,7 @@ import (
 	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
 	"github.com/ava-labs/avalanchego/snow/consensus/snowman/snowmantest"
 	"github.com/ava-labs/avalanchego/snow/engine/common"
+	"github.com/ava-labs/avalanchego/snow/engine/common/tracker"
 	"github.com/ava-labs/avalanchego/snow/engine/enginetest"
 	"github.com/ava-labs/avalanchego/snow/engine/snowman/ancestor"
 	"github.com/ava-labs/avalanchego/snow/engine/snowman/block/blocktest"
@@ -29,6 +30,7 @@ import (
 	"github.com/ava-labs/avalanchego/snow/snowtest"
 	"github.com/ava-labs/avalanchego/snow/validators"
 	"github.com/ava-labs/avalanchego/utils"
+	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/utils/set"
 	"github.com/ava-labs/avalanchego/version"
 )
@@ -3182,3 +3184,40 @@ func TestShouldIssueBlock(t *testing.T) {
 		})
 	}
 }
+
+type mockConnVDR struct {
+	tracker.Peers
+	percent float64
+}
+
+func (m *mockConnVDR) ConnectedPercent() float64 {
+	return m.percent
+}
+
+type logBuffer struct {
+	bytes.Buffer
+}
+
+func (logBuffer) Close() error {
+	return nil
+}
+
+func TestEngineAbortQueryWhenInPartition(t *testing.T) {
+	require := require.New(t)
+
+	// Buffer to record the log entries
+	buff := logBuffer{}
+
+	conf := DefaultConfig(t)
+	// Overwrite the log to record what it says
+	conf.Ctx.Log = logging.NewLogger("", logging.NewWrappedCore(logging.Verbo, &buff, logging.Plain.ConsoleEncoder()))
+	conf.Params = snowball.DefaultParameters
+	conf.ConnectedValidators = &mockConnVDR{percent: 0.7, Peers: conf.ConnectedValidators}
+
+	_, _, _, _, engine := setup(t, conf)
+
+	// Gossip will cause a pull query if enough stake is connected
+	require.NoError(engine.Gossip(context.Background()))
+
+	require.Contains(buff.String(), errInsufficientStake)
+}