fix(store): track store's contiguous head

p2p/server_test.go

@@ -17,6 +17,9 @@ func TestExchangeServer_handleRequestTimeout(t *testing.T) {
 	peer := createMocknet(t, 1)
 	s, err := store.NewStore[*headertest.DummyHeader](datastore.NewMapDatastore())
 	require.NoError(t, err)
+	head := headertest.RandDummyHeader(t)
+	head.HeightI %= 1000 // make it a bit lower
+	s.Init(context.Background(), head)
 	server, err := NewExchangeServer[*headertest.DummyHeader](
 		peer[0],
 		s,

store/batch.go

@@ -6,7 +6,7 @@ import (
 	"github.com/celestiaorg/go-header"
 )
 
-// batch keeps an adjacent range of headers and loosely mimics the Store
+// batch keeps a range of headers and loosely mimics the Store
 // interface. NOTE: Can fully implement Store for a use case.
 //
 // It keeps a mapping 'height -> header' and 'hash -> height'
@@ -76,7 +76,11 @@ func (b *batch[H]) getByHeight(height uint64) H {
 		return zero
 	}
 
-	return b.headers[height-base-1]
+	h := b.headers[height-base-1]
+	if h.Height() == height {
+		return h
+	}
+	return zero
 }
 
 // Append appends new headers to the batch.

store/heightsub.go

@@ -5,126 +5,133 @@ import (
 	"errors"
 	"sync"
 	"sync/atomic"
-
-	"github.com/celestiaorg/go-header"
 )
 
 // errElapsedHeight is thrown when a requested height was already provided to heightSub.
 var errElapsedHeight = errors.New("elapsed height")
 
 // heightSub provides a minimalistic mechanism to wait till header for a height becomes available.
-type heightSub[H header.Header[H]] struct {
+type heightSub struct {
 	// height refers to the latest locally available header height
 	// that has been fully verified and inserted into the subjective chain
 	height       atomic.Uint64
-	heightReqsLk sync.Mutex
-	heightReqs   map[uint64]map[chan H]struct{}
+	heightSubsLk sync.Mutex
+	heightSubs   map[uint64]*sub
+}
+
+type sub struct {
+	signal chan struct{}
+	count  int
 }
 
 // newHeightSub instantiates new heightSub.
-func newHeightSub[H header.Header[H]]() *heightSub[H] {
-	return &heightSub[H]{
-		heightReqs: make(map[uint64]map[chan H]struct{}),
+func newHeightSub() *heightSub {
+	return &heightSub{
+		heightSubs: make(map[uint64]*sub),
+	}
+}
+
+// Init the heightSub with a given height.
+// Notifies all awaiting [Wait] calls lower than height.
+func (hs *heightSub) Init(height uint64) {
+	hs.height.Store(height)
+
+	hs.heightSubsLk.Lock()
+	defer hs.heightSubsLk.Unlock()
+
+	for h := range hs.heightSubs {
+		if h < height {
+			hs.notify(h, true)
+		}
 	}
 }
 
 // Height reports current height.
-func (hs *heightSub[H]) Height() uint64 {
+func (hs *heightSub) Height() uint64 {
 	return hs.height.Load()
 }
 
 // SetHeight sets the new head height for heightSub.
-func (hs *heightSub[H]) SetHeight(height uint64) {
-	hs.height.Store(height)
+// Notifies all awaiting [Wait] calls in range from [heightSub.Height] to height.
+func (hs *heightSub) SetHeight(height uint64) {
+	for {
+		curr := hs.height.Load()
+		if curr >= height {
+			return
+		}
+		if !hs.height.CompareAndSwap(curr, height) {
+			continue
+		}
+
+		hs.heightSubsLk.Lock()
+		defer hs.heightSubsLk.Unlock() //nolint:gocritic // we have a return below
+
+		for ; curr <= height; curr++ {
+			hs.notify(curr, true)
+		}
+		return
+	}
 }
 
-// Sub subscribes for a header of a given height.
-// It can return errElapsedHeight, which means a requested header was already provided
+// Wait for a given height to be published.
+// It can return errElapsedHeight, which means a requested height was already seen
 // and caller should get it elsewhere.
-func (hs *heightSub[H]) Sub(ctx context.Context, height uint64) (H, error) {
-	var zero H
+func (hs *heightSub) Wait(ctx context.Context, height uint64) error {
 	if hs.Height() >= height {
-		return zero, errElapsedHeight
+		return errElapsedHeight
 	}
 
-	hs.heightReqsLk.Lock()
+	hs.heightSubsLk.Lock()
 	if hs.Height() >= height {
 		// This is a rare case we have to account for.
 		// The lock above can park a goroutine long enough for hs.height to change for a requested height,
 		// leaving the request never fulfilled and the goroutine deadlocked.
-		hs.heightReqsLk.Unlock()
-		return zero, errElapsedHeight
+		hs.heightSubsLk.Unlock()
+		return errElapsedHeight
 	}
-	resp := make(chan H, 1)
-	reqs, ok := hs.heightReqs[height]
+
+	sac, ok := hs.heightSubs[height]
 	if !ok {
-		reqs = make(map[chan H]struct{})
-		hs.heightReqs[height] = reqs
+		sac = &sub{
+			signal: make(chan struct{}, 1),
+		}
+		hs.heightSubs[height] = sac
 	}
-	reqs[resp] = struct{}{}
-	hs.heightReqsLk.Unlock()
+	sac.count++
+	hs.heightSubsLk.Unlock()
 
 	select {
-	case resp := <-resp:
-		return resp, nil
+	case <-sac.signal:
+		return nil
 	case <-ctx.Done():
 		// no need to keep the request, if the op has canceled
-		hs.heightReqsLk.Lock()
-		delete(reqs, resp)
-		if len(reqs) == 0 {
-			delete(hs.heightReqs, height)
-		}
-		hs.heightReqsLk.Unlock()
-		return zero, ctx.Err()
+		hs.heightSubsLk.Lock()
+		hs.notify(height, false)
+		hs.heightSubsLk.Unlock()
+		return ctx.Err()
 	}
 }
 
-// Pub processes all the outstanding subscriptions matching the given headers.
-// Pub is only safe when called from one goroutine.
-// For Pub to work correctly, heightSub has to be initialized with SetHeight
-// so that given headers are contiguous to the height on heightSub.
-func (hs *heightSub[H]) Pub(headers ...H) {
-	ln := len(headers)
-	if ln == 0 {
-		return
-	}
+// Notify and release the waiters in [Wait].
+// Note: do not advance heightSub's height.
+func (hs *heightSub) Notify(heights ...uint64) {
+	hs.heightSubsLk.Lock()
+	defer hs.heightSubsLk.Unlock()
 
-	height := hs.Height()
-	from, to := headers[0].Height(), headers[ln-1].Height()
-	if height+1 != from && height != 0 { // height != 0 is needed to enable init from any height and not only 1
-		log.Fatalf("PLEASE FILE A BUG REPORT: headers given to the heightSub are in the wrong order: expected %d, got %d", height+1, from)
-		return
+	for _, h := range heights {
+		hs.notify(h, true)
 	}
-	hs.SetHeight(to)
-
-	hs.heightReqsLk.Lock()
-	defer hs.heightReqsLk.Unlock()
-
-	// there is a common case where we Pub only header
-	// in this case, we shouldn't loop over each heightReqs
-	// and instead read from the map directly
-	if ln == 1 {
-		reqs, ok := hs.heightReqs[from]
-		if ok {
-			for req := range reqs {
-				req <- headers[0] // reqs must always be buffered, so this won't block
-			}
-			delete(hs.heightReqs, from)
-		}
+}
+
+func (hs *heightSub) notify(height uint64, all bool) {
+	sac, ok := hs.heightSubs[height]
+	if !ok {
 		return
 	}
 
-	// instead of looping over each header in 'headers', we can loop over each request
-	// which will drastically decrease idle iterations, as there will be less requests than headers
-	for height, reqs := range hs.heightReqs {
-		// then we look if any of the requests match the given range of headers
-		if height >= from && height <= to {
-			// and if so, calculate its position and fulfill requests
-			h := headers[height-from]
-			for req := range reqs {
-				req <- h // reqs must always be buffered, so this won't block
-			}
-			delete(hs.heightReqs, height)
-		}
+	sac.count--
+	if all || sac.count == 0 {
+		close(sac.signal)
+		delete(hs.heightSubs, height)
 	}
 }