fix(agent): bound concurrent beacon state and bad block fetches

pkg/agent/agent.go

@@ -55,7 +55,14 @@ type agent struct {
 	compressor *compression.Compressor
 }
 
-const namespace = "tracoor_agent"
+const (
+	namespace = "tracoor_agent"
+
+	logKeyPurpose = "purpose"
+	logKeySlot    = "slot"
+	labelAgent    = "agent"
+	labelQueue    = "queue"
+)
 
 func New(ctx context.Context, log logrus.FieldLogger, config *Config) (*agent, error) {
 	if config == nil {
@@ -127,9 +134,9 @@ func (s *agent) Start(ctx context.Context) error {
 			}
 
 			logCtx := s.log.WithFields(logrus.Fields{
-				"event_slot": event.Slot,
-				"event_root": fmt.Sprintf("%#x", event.Block),
-				"purpose":    "execution_block_trace",
+				"event_slot":  event.Slot,
+				"event_root":  fmt.Sprintf("%#x", event.Block),
+				logKeyPurpose: "execution_block_trace",
 			})
 
 			// Check if the block is too old to bother fetching.
@@ -195,9 +202,9 @@ func (s *agent) Start(ctx context.Context) error {
 		s.node.Beacon().Node().OnBlock(ctx, func(ctx context.Context, event *eth2v1.BlockEvent) error {
 			logCtx := s.log.WithFields(logrus.Fields{
 				"event_topic": "block",
-				"slot":        event.Slot,
+				logKeySlot:    event.Slot,
 				"root":        fmt.Sprintf("%#x", event.Block),
-				"purpose":     "beacon_state_and_block",
+				logKeyPurpose: "beacon_state_and_block",
 			})
 
 			if ignore, err := s.node.ShouldIgnoreEventFromSlot(event.Slot); err != nil {
@@ -224,7 +231,7 @@ func (s *agent) Start(ctx context.Context) error {
 					"event_old_head_block": rootAsString(chainReorg.OldHeadBlock),
 					"event_new_head_block": rootAsString(chainReorg.NewHeadBlock),
 					"event_depth":          chainReorg.Depth,
-					"purpose":              "chain_reorg",
+					logKeyPurpose:          "chain_reorg",
 					"event_slot":           chainReorg.Slot,
 				},
 			)

pkg/agent/consensus.go

@@ -77,6 +77,15 @@ func (s *agent) fetchAndIndexBeaconState(ctx context.Context, slot phase0.Slot)
 		stateID = fmt.Sprintf("%d", slot)
 	}
 
+	// Held until the upload finishes, not just the fetch, as the raw state and its
+	// compressed copy are both live until then.
+	releaseFetchSlot, err := s.acquireBeaconStateFetch(ctx)
+	if err != nil {
+		return errors.Wrap(err, "failed to acquire beacon state fetch slot")
+	}
+
+	defer releaseFetchSlot()
+
 	// Fetch the state
 	state, err := s.node.Beacon().Node().FetchRawBeaconState(ctx, stateID, string(mime.ContentTypeOctet))
 	if err != nil {
@@ -91,6 +100,10 @@ func (s *agent) fetchAndIndexBeaconState(ctx context.Context, slot phase0.Slot)
 		return errors.Wrap(err, "failed to compress beacon state")
 	}
 
+	// Drop the raw state before the upload so only the compressed copy is held
+	// for the duration of the store write.
+	state = nil
+
 	// Upload the state to the store
 	location, err = s.store.SaveBeaconState(ctx, &store.SaveParams{
 		Data:            &compressedState,
@@ -385,7 +398,7 @@ func (s *agent) fetchAndIndexBeaconBadBlocks(ctx context.Context, path string) e
 				})
 				if err != nil {
 					s.log.WithFields(logrus.Fields{
-						"slot":      slot,
+						logKeySlot:  slot,
 						"blockRoot": blockRoot,
 						"filePath":  filePath,
 					}).WithError(err).Error("Failed to save beacon bad block to store")
@@ -603,7 +616,7 @@ func (s *agent) fetchAndIndexBeaconBadBlobs(ctx context.Context, path string) er
 				})
 				if err != nil {
 					s.log.WithFields(logrus.Fields{
-						"slot":      slot,
+						logKeySlot:  slot,
 						"blockRoot": blockRoot,
 						"index":     index,
 						"filePath":  filePath,

pkg/agent/ethereum/config.go

@@ -27,6 +27,31 @@ type Config struct {
 	// which can cause the agent to get stuck as old states might not be available in
 	// the beacon node cache.
 	BeaconStateAgeThresholdEpochs uint64 `yaml:"beaconStateAgeThresholdEpochs" default:"1"`
+
+	// MaxConcurrentBeaconStateFetches bounds in-flight beacon state fetches across
+	// every agent in the process, capping peak memory independently of how many
+	// agents are configured. Applied process-wide by whichever agent starts first.
+	MaxConcurrentBeaconStateFetches int `yaml:"maxConcurrentBeaconStateFetches" default:"10"`
+
+	// MaxConcurrentExecutionBadBlockFetches bounds in-flight bad block fetches,
+	// as above.
+	MaxConcurrentExecutionBadBlockFetches int `yaml:"maxConcurrentExecutionBadBlockFetches" default:"10"`
+}
+
+func (c *Config) GetMaxConcurrentBeaconStateFetches() int {
+	if c.MaxConcurrentBeaconStateFetches <= 0 {
+		return 0
+	}
+
+	return c.MaxConcurrentBeaconStateFetches
+}
+
+func (c *Config) GetMaxConcurrentExecutionBadBlockFetches() int {
+	if c.MaxConcurrentExecutionBadBlockFetches <= 0 {
+		return 0
+	}
+
+	return c.MaxConcurrentExecutionBadBlockFetches
 }
 
 func (c *Config) Validate() error {

pkg/agent/execution.go

@@ -97,6 +97,15 @@ func (s *agent) fetchAndIndexExecutionBlockTrace(ctx context.Context, blockNumbe
 }
 
 func (s *agent) fetchAndIndexExecutionBadBlocks(ctx context.Context) error {
+	// Held across the indexing pass below, as the decoded slice is retained until
+	// it completes.
+	releaseFetchSlot, err := s.acquireExecutionBadBlockFetch(ctx)
+	if err != nil {
+		return errors.Wrap(err, "failed to acquire execution bad block fetch slot")
+	}
+
+	defer releaseFetchSlot()
+
 	// Fetch the bad blocks from the execution node.
 	blocks, err := s.node.Execution().GetBadBlocks(ctx)
 	if err != nil {

pkg/agent/limiter.go

@@ -0,0 +1,62 @@
+package agent
+
+import (
+	"context"
+	"sync"
+
+	"golang.org/x/sync/semaphore"
+)
+
+// Both paths read a whole response into memory, so peak usage is the limit
+// multiplied by the response size. States are roughly an order of magnitude
+// larger than bad block responses.
+const (
+	defaultMaxConcurrentBeaconStateFetches       = 10
+	defaultMaxConcurrentExecutionBadBlockFetches = 10
+)
+
+// fetchLimiter is a lazily sized process-wide concurrency budget. It is
+// package-level because `single` mode runs an agent per node in one process and
+// constructs each independently, so a per-agent budget would bound nothing. The
+// first caller fixes the size.
+type fetchLimiter struct {
+	once sync.Once
+	sem  *semaphore.Weighted
+}
+
+func (l *fetchLimiter) acquire(ctx context.Context, limit, def int) (func(), error) {
+	l.once.Do(func() {
+		if limit <= 0 {
+			limit = def
+		}
+
+		l.sem = semaphore.NewWeighted(int64(limit))
+	})
+
+	if err := l.sem.Acquire(ctx, 1); err != nil {
+		return nil, err
+	}
+
+	return func() { l.sem.Release(1) }, nil
+}
+
+var (
+	beaconStateFetchLimiter       fetchLimiter
+	executionBadBlockFetchLimiter fetchLimiter
+)
+
+func (s *agent) acquireBeaconStateFetch(ctx context.Context) (func(), error) {
+	return beaconStateFetchLimiter.acquire(
+		ctx,
+		s.Config.Ethereum.GetMaxConcurrentBeaconStateFetches(),
+		defaultMaxConcurrentBeaconStateFetches,
+	)
+}
+
+func (s *agent) acquireExecutionBadBlockFetch(ctx context.Context) (func(), error) {
+	return executionBadBlockFetchLimiter.acquire(
+		ctx,
+		s.Config.Ethereum.GetMaxConcurrentExecutionBadBlockFetches(),
+		defaultMaxConcurrentExecutionBadBlockFetches,
+	)
+}

pkg/agent/metrics.go

@@ -39,23 +39,23 @@ func GetMetricsInstance(namespace string) *Metrics {
 				Namespace: namespace,
 				Name:      "queue_size",
 				Help:      "The size of the queue",
-			}, []string{"queue", "agent"}),
+			}, []string{labelQueue, labelAgent}),
 			queueItemProcessingTime: prometheus.NewHistogramVec(prometheus.HistogramOpts{
 				Namespace: namespace,
 				Name:      "queue_item_processing_time_seconds",
 				Help:      "The time it takes to process an item from the queue",
 				Buckets:   prometheus.LinearBuckets(0, 3, 10),
-			}, []string{"queue", "agent"}),
+			}, []string{labelQueue, labelAgent}),
 			itemExported: prometheus.NewCounterVec(prometheus.CounterOpts{
 				Namespace: namespace,
 				Name:      "item_exported",
 				Help:      "The number of items exported",
-			}, []string{"queue", "agent"}),
+			}, []string{labelQueue, labelAgent}),
 			queueItemSkipped: prometheus.NewCounterVec(prometheus.CounterOpts{
 				Namespace: namespace,
 				Name:      "queue_item_skipped",
 				Help:      "The number of items skipped",
-			}, []string{"queue", "agent"}),
+			}, []string{labelQueue, labelAgent}),
 		}
 
 		prometheus.MustRegister(metricsInstance.queueSize)

pkg/compression/compressor_test.go

@@ -10,6 +10,11 @@ import (
 	"github.com/ethpandaops/tracoor/pkg/compression"
 )
 
+const (
+	testFilename   = "test"
+	testFilenameGz = "test.gz"
+)
+
 func TestNewCompressor(t *testing.T) {
 	c := compression.NewCompressor()
 	assert.NotNil(t, c)
@@ -70,13 +75,13 @@ func TestCompressor_Decompress(t *testing.T) {
 		{
 			name:     "Decompress Gzip",
 			data:     compressed,
-			filename: "test.gz",
+			filename: testFilenameGz,
 			wantErr:  false,
 		},
 		{
 			name:     "Decompress with nil data",
 			data:     nil,
-			filename: "test.gz",
+			filename: testFilenameGz,
 			wantErr:  true,
 		},
 		{
@@ -111,13 +116,13 @@ func TestAddExtension(t *testing.T) {
 	}{
 		{
 			name:      "Add Gzip extension",
-			filename:  "test",
+			filename:  testFilename,
 			algorithm: compression.Gzip,
 			want:      "test.gz",
 		},
 		{
 			name:      "Extension already present",
-			filename:  "test.gz",
+			filename:  testFilenameGz,
 			algorithm: compression.Gzip,
 			want:      "test.gz",
 		},
@@ -142,13 +147,13 @@ func TestRemoveExtension(t *testing.T) {
 	}{
 		{
 			name:      "Remove Gzip extension",
-			filename:  "test.gz",
+			filename:  testFilenameGz,
 			algorithm: compression.Gzip,
 			want:      "test",
 		},
 		{
 			name:      "No extension to remove",
-			filename:  "test",
+			filename:  testFilename,
 			algorithm: compression.Gzip,
 			want:      "test",
 		},
@@ -179,13 +184,13 @@ func TestHasCompressionExtension(t *testing.T) {
 	}{
 		{
 			name:      "Has Gzip extension",
-			filename:  "test.gz",
+			filename:  testFilenameGz,
 			algorithm: compression.Gzip,
 			want:      true,
 		},
 		{
 			name:      "No Gzip extension",
-			filename:  "test",
+			filename:  testFilename,
 			algorithm: compression.Gzip,
 			want:      false,
 		},
@@ -210,7 +215,7 @@ func TestGetCompressionAlgorithm(t *testing.T) {
 	}{
 		{
 			name:     "Get Gzip algorithm",
-			filename: "test.gz",
+			filename: testFilenameGz,
 			want:     compression.Gzip,
 			wantErr:  false,
 		},

pkg/server/persistence/beacon_bad_blob_test.go

@@ -98,7 +98,7 @@ func TestListBeaconBadBlob(t *testing.T) {
 	}
 	page := &PaginationCursor{}
 
-	mock.ExpectQuery("SELECT \\* FROM").WithArgs(filter.ID).WillReturnRows(sqlmock.NewRows([]string{"id", "node"}).AddRow("test-id", "test-node"))
+	mock.ExpectQuery("SELECT \\* FROM").WithArgs(filter.ID).WillReturnRows(sqlmock.NewRows([]string{columnID, columnNode}).AddRow("test-id", "test-node"))
 
 	blobs, err := indexer.ListBeaconBadBlob(ctx, filter, page)
 	assert.NoError(t, err)
@@ -301,15 +301,15 @@ func TestBeaconBadBlobIndividualFilters(t *testing.T) {
 			name   string
 			filter BeaconBadBlobFilter
 		}{
-			{"ID", BeaconBadBlobFilter{ID: &beaconBlob.ID}},
-			{"Node", BeaconBadBlobFilter{Node: &beaconBlob.Node}},
-			{"Slot", BeaconBadBlobFilter{Slot: &slot}},
-			{"Epoch", BeaconBadBlobFilter{Epoch: &epoch}},
-			{"BlockRoot", BeaconBadBlobFilter{BlockRoot: &beaconBlob.BlockRoot}},
-			{"NodeVersion", BeaconBadBlobFilter{NodeVersion: &beaconBlob.NodeVersion}},
-			{"Location", BeaconBadBlobFilter{Location: &beaconBlob.Location}},
-			{"Network", BeaconBadBlobFilter{Network: &beaconBlob.Network}},
-			{"BeaconImplementation", BeaconBadBlobFilter{BeaconImplementation: &beaconBlob.BeaconImplementation}},
+			{fieldID, BeaconBadBlobFilter{ID: &beaconBlob.ID}},
+			{fieldNode, BeaconBadBlobFilter{Node: &beaconBlob.Node}},
+			{fieldSlot, BeaconBadBlobFilter{Slot: &slot}},
+			{fieldEpoch, BeaconBadBlobFilter{Epoch: &epoch}},
+			{fieldBlockRoot, BeaconBadBlobFilter{BlockRoot: &beaconBlob.BlockRoot}},
+			{fieldNodeVersion, BeaconBadBlobFilter{NodeVersion: &beaconBlob.NodeVersion}},
+			{fieldLocation, BeaconBadBlobFilter{Location: &beaconBlob.Location}},
+			{fieldNetwork, BeaconBadBlobFilter{Network: &beaconBlob.Network}},
+			{fieldBeaconImplementation, BeaconBadBlobFilter{BeaconImplementation: &beaconBlob.BeaconImplementation}},
 			{"Index", BeaconBadBlobFilter{Index: &index}},
 		}
 

pkg/server/persistence/beacon_bad_block_test.go

@@ -96,7 +96,7 @@ func TestListBeaconBadBlock(t *testing.T) {
 	}
 	page := &PaginationCursor{}
 
-	mock.ExpectQuery("SELECT \\* FROM").WithArgs(filter.ID).WillReturnRows(sqlmock.NewRows([]string{"id", "node"}).AddRow("test-id", "test-node"))
+	mock.ExpectQuery("SELECT \\* FROM").WithArgs(filter.ID).WillReturnRows(sqlmock.NewRows([]string{columnID, columnNode}).AddRow("test-id", "test-node"))
 
 	blocks, err := indexer.ListBeaconBadBlock(ctx, filter, page)
 	assert.NoError(t, err)
@@ -288,15 +288,15 @@ func TestBeaconBadBlockIndividualFilters(t *testing.T) {
 			name   string
 			filter BeaconBadBlockFilter
 		}{
-			{"ID", BeaconBadBlockFilter{ID: &beaconBlock.ID}},
-			{"Node", BeaconBadBlockFilter{Node: &beaconBlock.Node}},
-			{"Slot", BeaconBadBlockFilter{Slot: &slot}},
-			{"Epoch", BeaconBadBlockFilter{Epoch: &epoch}},
-			{"BlockRoot", BeaconBadBlockFilter{BlockRoot: &beaconBlock.BlockRoot}},
-			{"NodeVersion", BeaconBadBlockFilter{NodeVersion: &beaconBlock.NodeVersion}},
-			{"Location", BeaconBadBlockFilter{Location: &beaconBlock.Location}},
-			{"Network", BeaconBadBlockFilter{Network: &beaconBlock.Network}},
-			{"BeaconImplementation", BeaconBadBlockFilter{BeaconImplementation: &beaconBlock.BeaconImplementation}},
+			{fieldID, BeaconBadBlockFilter{ID: &beaconBlock.ID}},
+			{fieldNode, BeaconBadBlockFilter{Node: &beaconBlock.Node}},
+			{fieldSlot, BeaconBadBlockFilter{Slot: &slot}},
+			{fieldEpoch, BeaconBadBlockFilter{Epoch: &epoch}},
+			{fieldBlockRoot, BeaconBadBlockFilter{BlockRoot: &beaconBlock.BlockRoot}},
+			{fieldNodeVersion, BeaconBadBlockFilter{NodeVersion: &beaconBlock.NodeVersion}},
+			{fieldLocation, BeaconBadBlockFilter{Location: &beaconBlock.Location}},
+			{fieldNetwork, BeaconBadBlockFilter{Network: &beaconBlock.Network}},
+			{fieldBeaconImplementation, BeaconBadBlockFilter{BeaconImplementation: &beaconBlock.BeaconImplementation}},
 		}
 
 		for _, tc := range testCases {