ethclient/lightclient: requestMap.close()

cmd/bltest/main.go

@@ -173,7 +173,6 @@ loop:
 			break loop
 		}
 	}
-
-	client.Stop()
+	client.Close()
 	return nil
 }

ethclient/lightclient/chain.go

@@ -39,7 +39,7 @@ const recentCanonicalLength = 256
 type canonicalChainFields struct {
 	chainLock      sync.Mutex
 	head, finality *btypes.ExecutionHeader
-	recent         map[uint64]common.Hash // nil until initialized
+	recent         map[uint64]common.Hash // nil while head == nil
 	recentTail     uint64                 // if recent != nil then recent hashes are available from recentTail to head
 	tailFetchCh    chan struct{}
 	finalized      *lru.Cache[uint64, common.Hash]  // finalized but not recent hashes
@@ -53,6 +53,10 @@ func (c *Client) initCanonicalChain() {
 	go c.tailFetcher()
 }
 
+func (c *Client) closeCanonicalChain() {
+	c.requests.close()
+}
+
 // Process implements request.Module in order to get notified about new heads.
 func (c *Client) Process(requester request.Requester, events []request.Event) {
 	if finality, ok := c.headTracker.ValidatedFinality(); ok {
@@ -268,6 +272,11 @@ func (c *Client) initBlocksAndHeaders() {
 	c.blockRequests = newRequestMap[common.Hash, *types.Block](c.requestBlock)
 }
 
+func (c *Client) closeBlocksAndHeaders() {
+	c.headerRequests.close()
+	c.blockRequests.close()
+}
+
 func (c *Client) requestHeader(ctx context.Context, hash common.Hash) (*types.Header, error) {
 	var header *types.Header
 	log.Debug("Starting RPC request", "type", "eth_getBlockByHash", "hash", hash, "full", false)

ethclient/lightclient/lightclient.go

@@ -92,7 +92,11 @@ func (c *Client) Start() {
 	}
 }
 
-func (c *Client) Stop() {
+func (c *Client) Close() {
+	c.closeBlocksAndHeaders()
+	c.closeCanonicalChain()
+	c.closeTxAndReceipts()
+	c.closeLightState()
 	c.scheduler.Stop()
 }
 

ethclient/lightclient/request_map.go

@@ -18,15 +18,29 @@ package lightclient
 
 import (
 	"context"
+	"errors"
 	"sync"
 )
 
 type requestMap[K comparable, V any] struct {
 	lock      sync.Mutex
+	closed    bool
 	requestFn func(context.Context, K) (V, error)
 	requests  map[K]*mappedRequest[K, V]
 }
 
+type mappedRequest[K comparable, V any] struct {
+	lock              sync.Mutex
+	rm                *requestMap[K, V]
+	key               K
+	refCount          int
+	delivered, closed bool
+	deliveredCh       chan struct{}
+	cancelFn          func() // called when delivered || closed || refCount == 0 becomes true
+	result            V
+	err               error
+}
+
 func newRequestMap[K comparable, V any](requestFn func(context.Context, K) (V, error)) *requestMap[K, V] {
 	return &requestMap[K, V]{
 		requestFn: requestFn,
@@ -34,6 +48,22 @@ func newRequestMap[K comparable, V any](requestFn func(context.Context, K) (V, e
 	}
 }
 
+func (rm *requestMap[K, V]) close() {
+	rm.lock.Lock()
+	defer rm.lock.Unlock()
+
+	if rm.closed {
+		return
+	}
+	for _, req := range rm.requests {
+		if !req.delivered && req.refCount != 0 {
+			req.cancelFn()
+			req.closed = true
+		}
+	}
+	rm.closed = true
+}
+
 func (rm *requestMap[K, V]) request(key K) *mappedRequest[K, V] {
 	rm.lock.Lock()
 	defer rm.lock.Unlock()
@@ -53,6 +83,14 @@ func (rm *requestMap[K, V]) request(key K) *mappedRequest[K, V] {
 		cancelFn:    cancelFn,
 	}
 	rm.requests[key] = r
+	if rm.closed {
+		// return a closed dummy request for simplicity
+		r.closed = true
+		r.cancelFn()
+		var null V
+		r.deliver(null, errors.New("request map is closed"))
+		return r
+	}
 	if rm.requestFn != nil {
 		go func() {
 			result, err := rm.requestFn(ctx, key)
@@ -91,25 +129,13 @@ func (rm *requestMap[K, V]) tryDeliver(key K, result V) {
 	}
 }
 
-type mappedRequest[K comparable, V any] struct {
-	lock        sync.Mutex
-	rm          *requestMap[K, V]
-	key         K
-	refCount    int
-	delivered   bool
-	deliveredCh chan struct{}
-	cancelFn    func() // called when delivered || refCount == 0 becomes true
-	result      V
-	err         error
-}
-
 func (r *mappedRequest[K, V]) deliver(result V, err error) {
 	r.lock.Lock()
 	if !r.delivered {
 		r.result, r.err = result, err
 		r.delivered = true
 		close(r.deliveredCh)
-		if r.refCount != 0 {
+		if !r.closed && r.refCount != 0 {
 			r.cancelFn()
 		}
 	}
@@ -133,7 +159,7 @@ func (r *mappedRequest[K, V]) release() {
 	r.refCount--
 	if r.refCount == 0 {
 		delete(r.rm.requests, r.key)
-		if !r.delivered {
+		if !r.delivered && !r.closed {
 			r.cancelFn()
 		}
 	}

ethclient/lightclient/state.go

@@ -60,6 +60,11 @@ func (c *Client) initLightState() {
 	c.codeRequests = newRequestMap[codeRequest, []byte](c.requestCode)
 }
 
+func (c *Client) closeLightState() {
+	c.proofRequests.close()
+	c.codeRequests.close()
+}
+
 func (c *Client) fetchProof(ctx context.Context, req proofRequest) (*gethclient.AccountResult, error) {
 	if proof, ok := c.proofCache.Get(req); ok {
 		return proof, nil

ethclient/lightclient/transactions.go

@@ -70,6 +70,10 @@ func (c *Client) initTxAndReceipts() {
 	c.receiptsRequests = newRequestMap[common.Hash, types.Receipts](c.requestBlockReceipts)
 }
 
+func (c *Client) closeTxAndReceipts() {
+	c.receiptsRequests.close()
+}
+
 func (c *Client) getTxByHash(ctx context.Context, txHash common.Hash) (tx *types.Transaction, isPending bool, err error) {
 	if pos, ok := c.txPosCache.Get(txHash); ok {
 		if hash, ok := c.getCachedHash(pos.blockNumber); ok && hash == pos.blockHash {