fix: avs subscriber error not being handled correctly (yetanotherco#626)

Author: taturosati

aggregator/internal/pkg/subscriber.go

@@ -1,53 +1,35 @@
 package pkg
 
-import (
-	"errors"
-	"fmt"
-	"time"
-
-	"github.com/ethereum/go-ethereum/accounts/abi/bind"
-)
-
-const (
-	MaxRetries    = 100
-	RetryInterval = 1 * time.Second
-)
-
 func (agg *Aggregator) SubscribeToNewTasks() error {
-	for retries := 0; retries < MaxRetries; retries++ {
-		err := agg.tryCreateTaskSubscriber()
-		if err == nil {
-			_ = agg.subscribeToNewTasks() // This will block until an error occurs
-		}
-
-		message := fmt.Sprintf("Failed to subscribe to new tasks. Retrying in %v", RetryInterval)
-		agg.AggregatorConfig.BaseConfig.Logger.Info(message)
-		time.Sleep(RetryInterval)
+	err := agg.subscribeToNewTasks()
+	if err != nil {
+		return err
 	}
 
-	return errors.New("failed to subscribe to new tasks after max retries")
-}
-
-func (agg *Aggregator) subscribeToNewTasks() error {
 	for {
 		select {
 		case err := <-agg.taskSubscriber.Err():
-			return err
+			agg.AggregatorConfig.BaseConfig.Logger.Info("Failed to subscribe to new tasks", "err", err)
+			agg.taskSubscriber.Unsubscribe()
+			err = agg.subscribeToNewTasks()
+			if err != nil {
+				return err
+			}
 		case newBatch := <-agg.NewBatchChan:
 			agg.AddNewTask(newBatch.BatchMerkleRoot, newBatch.TaskCreatedBlock)
 		}
 	}
+
 }
 
-func (agg *Aggregator) tryCreateTaskSubscriber() error {
+func (agg *Aggregator) subscribeToNewTasks() error {
 	var err error
 
-	agg.AggregatorConfig.BaseConfig.Logger.Info("Subscribing to Ethereum serviceManager task events")
-	agg.taskSubscriber, err = agg.avsSubscriber.AvsContractBindings.ServiceManager.WatchNewBatch(&bind.WatchOpts{},
-		agg.NewBatchChan, nil)
+	agg.taskSubscriber, err = agg.avsSubscriber.SubscribeToNewTasks(agg.NewBatchChan)
 
 	if err != nil {
 		agg.AggregatorConfig.BaseConfig.Logger.Info("Failed to create task subscriber", "err", err)
 	}
+
 	return err
 }

core/chainio/avs_subscriber.go

@@ -1,6 +1,9 @@
 package chainio
 
 import (
+	"fmt"
+	"time"
+
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/ethereum/go-ethereum/event"
 	servicemanager "github.com/yetanotherco/aligned_layer/contracts/bindings/AlignedLayerServiceManager"
@@ -9,6 +12,11 @@ import (
 	sdklogging "github.com/Layr-Labs/eigensdk-go/logging"
 )
 
+const (
+	MaxRetries    = 100
+	RetryInterval = 1 * time.Second
+)
+
 // NOTE(marian): Leaving this commented code here as it may be useful in the short term.
 // type AvsSubscriberer interface {
 // 	SubscribeToNewTasks(newTaskCreatedChan chan *cstaskmanager.ContractAlignedLayerTaskManagerNewTaskCreated) event.Subscription
@@ -42,15 +50,22 @@ func NewAvsSubscriberFromConfig(baseConfig *config.BaseConfig) (*AvsSubscriber,
 	}, nil
 }
 
-func (s *AvsSubscriber) SubscribeToNewTasks(newTaskCreatedChan chan *servicemanager.ContractAlignedLayerServiceManagerNewBatch) event.Subscription {
-	sub, err := s.AvsContractBindings.ServiceManager.WatchNewBatch(
-		&bind.WatchOpts{}, newTaskCreatedChan, nil,
-	)
-	if err != nil {
-		s.logger.Error("Failed to subscribe to new AlignedLayer tasks", "err", err)
+func (s *AvsSubscriber) SubscribeToNewTasks(newTaskCreatedChan chan *servicemanager.ContractAlignedLayerServiceManagerNewBatch) (event.Subscription, error) {
+	for i := 0; i < MaxRetries; i++ {
+		sub, err := s.AvsContractBindings.ServiceManager.WatchNewBatch(
+			&bind.WatchOpts{}, newTaskCreatedChan, nil,
+		)
+		if err != nil {
+			s.logger.Warn("Failed to subscribe to new AlignedLayer tasks", "err", err)
+			time.Sleep(RetryInterval)
+			continue
+		}
+
+		s.logger.Info("Subscribed to new AlignedLayer tasks")
+		return sub, nil
 	}
-	s.logger.Infof("Subscribed to new AlignedLayer tasks")
-	return sub
+
+	return nil, fmt.Errorf("Failed to subscribe to new AlignedLayer tasks after %d retries", MaxRetries)
 }
 
 // func (s *AvsSubscriber) SubscribeToTaskResponses(taskResponseChan chan *cstaskmanager.ContractAlignedLayerTaskManagerTaskResponded) event.Subscription {

operator/pkg/operator.go

@@ -6,12 +6,13 @@ import (
 	"crypto/ecdsa"
 	"encoding/binary"
 	"fmt"
-	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/yetanotherco/aligned_layer/operator/risc_zero"
 	"log"
 	"sync"
 	"time"
 
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/yetanotherco/aligned_layer/operator/risc_zero"
+
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/yetanotherco/aligned_layer/metrics"
 
@@ -118,13 +119,15 @@ func NewOperatorFromConfig(configuration config.OperatorConfig) (*Operator, erro
 	return operator, nil
 }
 
-func (o *Operator) SubscribeToNewTasks() event.Subscription {
-	sub := o.avsSubscriber.SubscribeToNewTasks(o.NewTaskCreatedChan)
-	return sub
+func (o *Operator) SubscribeToNewTasks() (event.Subscription, error) {
+	return o.avsSubscriber.SubscribeToNewTasks(o.NewTaskCreatedChan)
 }
 
 func (o *Operator) Start(ctx context.Context) error {
-	sub := o.SubscribeToNewTasks()
+	sub, err := o.SubscribeToNewTasks()
+	if err != nil {
+		log.Fatal("Could not subscribe to new tasks")
+	}
 
 	var metricsErrChan <-chan error
 	if o.Config.Operator.EnableMetrics {
@@ -143,7 +146,10 @@ func (o *Operator) Start(ctx context.Context) error {
 		case err := <-sub.Err():
 			o.Logger.Infof("Error in websocket subscription", "err", err)
 			sub.Unsubscribe()
-			sub = o.SubscribeToNewTasks()
+			sub, err = o.SubscribeToNewTasks()
+			if err != nil {
+				o.Logger.Fatal("Could not subscribe to new tasks")
+			}
 		case newBatchLog := <-o.NewTaskCreatedChan:
 			err := o.ProcessNewBatchLog(newBatchLog)
 			if err != nil {