From 5f219aa5cf6007c6a02861f6bc077542b8cf5e43 Mon Sep 17 00:00:00 2001
From: simlecode <69969590+simlecode@users.noreply.github.com>
Date: Thu, 17 Aug 2023 10:35:58 +0800
Subject: [PATCH] fix: Unsubscribe required on exit

---
 pkg/chain/store.go | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/pkg/chain/store.go b/pkg/chain/store.go
index d17a951692..5d161b7999 100644
--- a/pkg/chain/store.go
+++ b/pkg/chain/store.go
@@ -695,8 +695,17 @@ func (store *Store) SubHeadChanges(ctx context.Context) chan []*types.HeadChange
 	}}
 
 	go func() {
-		defer close(out)
-		var unsubOnce sync.Once
+		defer func() {
+			// Tell the caller we're done first, the following may block for a bit.
+			close(out)
+
+			// Unsubscribe.
+			store.headEvents.Unsub(subCh)
+
+			// Drain the channel.
+			for range subCh {
+			}
+		}()
 
 		for {
 			select {
@@ -716,9 +725,8 @@ func (store *Store) SubHeadChanges(ctx context.Context) chan []*types.HeadChange
 					log.Warnf("head change sub is slow, has %d buffered entries", len(out))
 				}
 			case <-ctx.Done():
-				unsubOnce.Do(func() {
-					go store.headEvents.Unsub(subCh)
-				})
+				log.Infof("exit sub head change: %v", ctx.Err())
+				return
 			}
 		}
 	}()