fix(store): properly update store head

store/heightsub.go

@@ -34,8 +34,17 @@ func (hs *heightSub[H]) Height() uint64 {
 }
 
 // SetHeight sets the new head height for heightSub.
+// Only higher then current height will be set.
 func (hs *heightSub[H]) SetHeight(height uint64) {
-	hs.height.Store(height)
+	for {
+		curr := hs.height.Load()
+		if curr >= height {
+			return
+		}
+		if hs.height.CompareAndSwap(curr, height) {
+			return
+		}
+	}
 }
 
 // Sub subscribes for a header of a given height.
@@ -89,12 +98,7 @@ func (hs *heightSub[H]) Pub(headers ...H) {
 		return
 	}
 
-	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
-	}
 	hs.SetHeight(to)
 
 	hs.heightReqsLk.Lock()

store/heightsub_test.go

@@ -47,6 +47,37 @@ func TestHeightSub(t *testing.T) {
 	}
 }
 
+func TestHeightSubNonAdjacement(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	hs := newHeightSub[*headertest.DummyHeader]()
+
+	{
+		h := headertest.RandDummyHeader(t)
+		h.HeightI = 100
+		hs.SetHeight(99)
+		hs.Pub(h)
+	}
+
+	{
+		go func() {
+			// fixes flakiness on CI
+			time.Sleep(time.Millisecond)
+
+			h1 := headertest.RandDummyHeader(t)
+			h1.HeightI = 200
+			h2 := headertest.RandDummyHeader(t)
+			h2.HeightI = 300
+			hs.Pub(h1, h2)
+		}()
+
+		h, err := hs.Sub(ctx, 200)
+		assert.NoError(t, err)
+		assert.NotNil(t, h)
+	}
+}
+
 func TestHeightSubCancellation(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()

store/store.go

@@ -1,9 +1,11 @@
 package store
 
 import (
+	"cmp"
 	"context"
 	"errors"
 	"fmt"
+	"slices"
 	"sync/atomic"
 	"time"
 
@@ -51,6 +53,8 @@ type Store[H header.Header[H]] struct {
 	writesDn chan struct{}
 	// writeHead maintains the current write head
 	writeHead atomic.Pointer[H]
+
+	knownHeaders map[uint64]H
 	// pending keeps headers pending to be written in one batch
 	pending *batch[H]
 
@@ -108,6 +112,8 @@ func newStore[H header.Header[H]](ds datastore.Batching, opts ...Option) (*Store
 		writesDn:    make(chan struct{}),
 		pending:     newBatch[H](params.WriteBatchSize),
 		Params:      params,
+
+		knownHeaders: make(map[uint64]H),
 	}, nil
 }
 
@@ -164,27 +170,30 @@ func (s *Store[H]) Stop(ctx context.Context) error {
 }
 
 func (s *Store[H]) Height() uint64 {
-	return s.heightSub.Height()
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	head, err := s.Head(ctx)
+	if err != nil {
+		// TODO(cristaloleg): log? panic? retry?
+		return 0
+	}
+	return head.Height()
 }
 
+// Head returns the highest contiguous header written to the store.
 func (s *Store[H]) Head(ctx context.Context, _ ...header.HeadOption[H]) (H, error) {
-	head, err := s.GetByHeight(ctx, s.heightSub.Height())
-	if err == nil {
-		return head, nil
+	headPtr := s.writeHead.Load()
+	if headPtr != nil {
+		return *headPtr, nil
 	}
 
-	var zero H
-	head, err = s.readHead(ctx)
-	switch {
-	default:
+	head, err := s.readHead(ctx)
+	if err != nil {
+		var zero H
 		return zero, err
-	case errors.Is(err, datastore.ErrNotFound), errors.Is(err, header.ErrNotFound):
-		return zero, header.ErrNoHead
-	case err == nil:
-		s.heightSub.SetHeight(head.Height())
-		log.Infow("loaded head", "height", head.Height(), "hash", head.Hash())
-		return head, nil
 	}
+	return head, nil
 }
 
 func (s *Store[H]) Get(ctx context.Context, hash header.Hash) (H, error) {
@@ -300,6 +309,9 @@ func (s *Store[H]) HasAt(_ context.Context, height uint64) bool {
 	return height != uint64(0) && s.Height() >= height
 }
 
+// Append the given headers to the store. Real write to the disk happens
+// asynchronously and might fail without reporting error (just logging).
+// TODO(cristaloleg): add retries to the flush worker.
 func (s *Store[H]) Append(ctx context.Context, headers ...H) error {
 	lh := len(headers)
 	if lh == 0 {
@@ -319,10 +331,15 @@ func (s *Store[H]) Append(ctx context.Context, headers ...H) error {
 		head = *headPtr
 	}
 
+	continuousHead := head
+
+	slices.SortFunc(headers, func(a, b H) int {
+		return cmp.Compare(a.Height(), b.Height())
+	})
+
 	// collect valid headers
 	verified := make([]H, 0, lh)
 	for i, h := range headers {
-
 		err = head.Verify(h)
 		if err != nil {
 			var verErr *header.VerifyError
@@ -344,11 +361,18 @@ func (s *Store[H]) Append(ctx context.Context, headers ...H) error {
 		}
 		verified = append(verified, h)
 		head = h
+
+		if continuousHead.Height()+1 == head.Height() {
+			continuousHead = head
+		} else {
+			s.knownHeaders[head.Height()] = head
+		}
 	}
 
 	onWrite := func() {
-		newHead := verified[len(verified)-1]
+		newHead := s.tryAdvanceHead(continuousHead)
 		s.writeHead.Store(&newHead)
+
 		log.Infow("new head", "height", newHead.Height(), "hash", newHead.Hash())
 		s.metrics.newHead(newHead.Height())
 	}
@@ -493,6 +517,22 @@ func (s *Store[H]) get(ctx context.Context, hash header.Hash) ([]byte, error) {
 	return data, nil
 }
 
+// try advance heighest header if we saw a higher continuous before.
+func (s *Store[H]) tryAdvanceHead(highestHead H) H {
+	curr := highestHead.Height()
+
+	for len(s.knownHeaders) > 0 {
+		h, ok := s.knownHeaders[curr+1]
+		if !ok {
+			break
+		}
+		highestHead = h
+		delete(s.knownHeaders, curr+1)
+		curr++
+	}
+	return highestHead
+}
+
 // indexTo saves mapping between header Height and Hash to the given batch.
 func indexTo[H header.Header[H]](ctx context.Context, batch datastore.Batch, headers ...H) error {
 	for _, h := range headers {

store/store_test.go

@@ -141,6 +141,86 @@ func TestStore_Append_BadHeader(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestStore_Append_stableHeadWhenGaps(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	t.Cleanup(cancel)
+
+	suite := headertest.NewTestSuite(t)
+
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	store := NewTestStore(t, ctx, ds, suite.Head(), WithWriteBatchSize(4))
+
+	head, err := store.Head(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, head.Hash(), suite.Head().Hash())
+
+	firstChunk := suite.GenDummyHeaders(5)
+	for i := range firstChunk {
+		t.Log("firstChunk:", firstChunk[i].Height(), firstChunk[i].Hash())
+	}
+	missedChunk := suite.GenDummyHeaders(5)
+	for i := range missedChunk {
+		t.Log("missedChunk:", missedChunk[i].Height(), missedChunk[i].Hash())
+	}
+	lastChunk := suite.GenDummyHeaders(5)
+	for i := range lastChunk {
+		t.Log("lastChunk:", lastChunk[i].Height(), lastChunk[i].Hash())
+	}
+
+	wantHead := firstChunk[len(firstChunk)-1]
+	t.Log("wantHead", wantHead.Height(), wantHead.Hash())
+	latestHead := lastChunk[len(lastChunk)-1]
+	t.Log("latestHead", latestHead.Height(), latestHead.Hash())
+
+	{
+		err := store.Append(ctx, firstChunk...)
+		require.NoError(t, err)
+		// wait for batch to be written.
+		time.Sleep(100 * time.Millisecond)
+
+		// head is advanced to the last known header.
+		head, err := store.Head(ctx)
+		require.NoError(t, err)
+		assert.Equal(t, head.Hash(), wantHead.Hash())
+
+		// check that store height is aligned with the head.
+		height := store.Height()
+		assert.Equal(t, height, head.Height())
+	}
+	{
+		err := store.Append(ctx, lastChunk...)
+		require.NoError(t, err)
+		// wait for batch to be written.
+		time.Sleep(100 * time.Millisecond)
+
+		// head is not advanced due to a gap.
+		head, err := store.Head(ctx)
+		require.NoError(t, err)
+		assert.Equal(t, head.Height(), wantHead.Height())
+		t.Log("head", head.Height(), head.Hash())
+		assert.Equal(t, head.Hash(), wantHead.Hash())
+
+		// check that store height is aligned with the head.
+		height := store.Height()
+		assert.Equal(t, height, head.Height())
+	}
+	{
+		err := store.Append(ctx, missedChunk...)
+		require.NoError(t, err)
+		// wait for batch to be written.
+		time.Sleep(100 * time.Millisecond)
+
+		// after appending missing headers we're on the latest header.
+		head, err := store.Head(ctx)
+		require.NoError(t, err)
+		assert.Equal(t, head.Hash(), latestHead.Hash())
+
+		// check that store height is aligned with the head.
+		height := store.Height()
+		assert.Equal(t, height, head.Height())
+	}
+}
+
 // TestStore_GetRange tests possible combinations of requests and ensures that
 // the store can handle them adequately (even malformed requests)
 func TestStore_GetRange(t *testing.T) {

sync/sync_head_test.go

@@ -2,6 +2,7 @@ package sync
 
 import (
 	"context"
+	"errors"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -121,20 +122,17 @@ func (t *wrappedGetter) Head(ctx context.Context, options ...header.HeadOption[*
 }
 
 func (t *wrappedGetter) Get(ctx context.Context, hash header.Hash) (*headertest.DummyHeader, error) {
-	// TODO implement me
-	panic("implement me")
+	return nil, errors.New("implement me")
 }
 
 func (t *wrappedGetter) GetByHeight(ctx context.Context, u uint64) (*headertest.DummyHeader, error) {
-	// TODO implement me
-	panic("implement me")
+	return nil, errors.New("implement me")
 }
 
 func (t *wrappedGetter) GetRangeByHeight(
 	ctx context.Context,
 	from *headertest.DummyHeader,
 	to uint64,
 ) ([]*headertest.DummyHeader, error) {
-	// TODO implement me
-	panic("implement me")
+	return nil, errors.New("implement me")
 }