From b9cfbb2cdae0b1e6e4d713fc12ea5305fca94561 Mon Sep 17 00:00:00 2001
From: battlmonstr <battlmonstr@users.noreply.github.com>
Date: Wed, 15 May 2024 10:19:57 +0200
Subject: [PATCH] polygon/sync: refactor Service.Run with errgroup

---
 polygon/sync/service.go | 61 ++++++++---------------------------------
 1 file changed, 11 insertions(+), 50 deletions(-)

diff --git a/polygon/sync/service.go b/polygon/sync/service.go
index 995a1eedf05..64aaf1fb566 100644
--- a/polygon/sync/service.go
+++ b/polygon/sync/service.go
@@ -5,6 +5,7 @@ import (
 
 	lru "github.com/hashicorp/golang-lru/arc/v2"
 	"github.com/ledgerwatch/log/v3"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/ledgerwatch/erigon-lib/chain"
 	"github.com/ledgerwatch/erigon-lib/common"
@@ -108,57 +109,17 @@ func NewService(
 	}
 }
 
-func (s *service) Run(ctx context.Context) error {
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
+func (s *service) Run(parentCtx context.Context) error {
+	group, ctx := errgroup.WithContext(parentCtx)
 
-	var serviceErr error
-
-	go func() {
-		s.p2pService.Run(ctx)
-	}()
-
-	go func() {
-		err := s.store.Run(ctx)
-		if (err != nil) && (ctx.Err() == nil) {
-			serviceErr = err
-			cancel()
-		}
-	}()
-
-	go func() {
-		err := s.events.Run(ctx)
-		if (err != nil) && (ctx.Err() == nil) {
-			serviceErr = err
-			cancel()
-		}
-	}()
-
-	go func() {
-		// TODO: remove when heimdall.NewService is functional
-		if s.heimdallService == nil {
-			return
-		}
-		err := s.heimdallService.Run(ctx)
-		if (err != nil) && (ctx.Err() == nil) {
-			serviceErr = err
-			cancel()
-		}
-	}()
-
-	go func() {
-		err := s.sync.Run(ctx)
-		if (err != nil) && (ctx.Err() == nil) {
-			serviceErr = err
-			cancel()
-		}
-	}()
-
-	<-ctx.Done()
-
-	if serviceErr != nil {
-		return serviceErr
+	group.Go(func() error { s.p2pService.Run(ctx); return nil })
+	group.Go(func() error { return s.store.Run(ctx) })
+	group.Go(func() error { return s.events.Run(ctx) })
+	// TODO: remove the check when heimdall.NewService is functional
+	if s.heimdallService != nil {
+		group.Go(func() error { return s.heimdallService.Run(ctx) })
 	}
+	group.Go(func() error { return s.sync.Run(ctx) })
 
-	return ctx.Err()
+	return group.Wait()
 }