From 32ec784e1e21db2c9517d97cc8be1f56f7190449 Mon Sep 17 00:00:00 2001
From: Abdul Rabbani <abdulrabbani00@gmail.com>
Date: Thu, 23 Jun 2022 08:50:09 -0400
Subject: [PATCH] Close channels in the correct location

---
 pkg/beaconclient/capturehead.go   | 2 ++
 pkg/beaconclient/incomingsse.go   | 3 +--
 pkg/beaconclient/processevents.go | 2 --
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/pkg/beaconclient/capturehead.go b/pkg/beaconclient/capturehead.go
index 51cecae..5ab00e2 100644
--- a/pkg/beaconclient/capturehead.go
+++ b/pkg/beaconclient/capturehead.go
@@ -39,6 +39,8 @@ func (bc *BeaconClient) StopHeadTracking(ctx context.Context, skipSee bool) {
 			bc.HeadTracking.SseClient.Unsubscribe(bc.HeadTracking.MessagesCh)
 			bc.ReOrgTracking.SseClient.Unsubscribe(bc.ReOrgTracking.MessagesCh)
 			log.Info("Successfully unsubscribed to SSE client")
+			close(bc.ReOrgTracking.MessagesCh)
+			close(bc.HeadTracking.MessagesCh)
 		}
 		log.Info("Successfully stopped the head tracking service.")
 	default:
diff --git a/pkg/beaconclient/incomingsse.go b/pkg/beaconclient/incomingsse.go
index 5831407..0ebd1d5 100644
--- a/pkg/beaconclient/incomingsse.go
+++ b/pkg/beaconclient/incomingsse.go
@@ -47,8 +47,7 @@ func handleIncomingSseEvent[P ProcessedEvents](ctx context.Context, eventHandler
 	for {
 		select {
 		case <-ctx.Done():
-			close(eventHandler.MessagesCh)
-			close(eventHandler.ErrorCh)
+			close(eventHandler.ProcessCh)
 			return
 		case message := <-eventHandler.MessagesCh:
 			// Message can be nil if its a keep-alive message
diff --git a/pkg/beaconclient/processevents.go b/pkg/beaconclient/processevents.go
index ed8da5c..f95cb31 100644
--- a/pkg/beaconclient/processevents.go
+++ b/pkg/beaconclient/processevents.go
@@ -33,7 +33,6 @@ func (bc *BeaconClient) handleReorg(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
-			close(bc.ReOrgTracking.ProcessCh)
 			return
 		case reorg := <-bc.ReOrgTracking.ProcessCh:
 			log.WithFields(log.Fields{"reorg": reorg}).Debug("Received a new reorg message.")
@@ -55,7 +54,6 @@ func (bc *BeaconClient) handleHead(ctx context.Context, maxWorkers int) {
 	for {
 		select {
 		case <-ctx.Done():
-			close(bc.HeadTracking.ProcessCh)
 			close(workCh)
 			return
 		case head := <-bc.HeadTracking.ProcessCh: