introduct getByHeight

store/store.go

@@ -243,12 +243,16 @@ func (s *Store[H]) GetByHeight(ctx context.Context, height uint64) (H, error) {
 		return h, nil
 	}
 
+	return s.getByHeight(ctx, height)
+}
+
+func (s *Store[H]) getByHeight(ctx context.Context, height uint64) (H, error) {
+	var zero H
 	hash, err := s.heightIndex.HashByHeight(ctx, height)
 	if err != nil {
 		if errors.Is(err, datastore.ErrNotFound) {
 			return zero, header.ErrNotFound
 		}
-
 		return zero, err
 	}
 
@@ -519,7 +523,8 @@ func (s *Store[H]) tryAdvanceHead(ctx context.Context, headers ...H) {
 		if headers[i].Height() != currHeight+1 {
 			break
 		}
-		s.writeHead.Store(&headers[i])
+		newHead := headers[i]
+		s.writeHead.Store(&newHead)
 		currHeight++
 	}
 
@@ -528,10 +533,11 @@ func (s *Store[H]) tryAdvanceHead(ctx context.Context, headers ...H) {
 
 	// advance based on already written headers.
 	for {
-		newHead, err := s.GetByHeight(ctx, currHeight+1)
+		h, err := s.getByHeight(ctx, currHeight+1)
 		if err != nil {
 			break
 		}
+		newHead := h
 		s.writeHead.Store(&newHead)
 		currHeight++
 	}

store/store_test.go

@@ -190,6 +190,7 @@ func TestStore_Append(t *testing.T) {
 
 	head, err = store.Head(ctx)
 	assert.NoError(t, err)
+	assert.Equal(t, head.Height(), headers[len(headers)-1].Height())
 	assert.Equal(t, head.Hash(), headers[len(headers)-1].Hash())
 }
 
@@ -222,6 +223,7 @@ func TestStore_Append_stableHeadWhenGaps(t *testing.T) {
 		// head is advanced to the last known header.
 		head, err := store.Head(ctx)
 		require.NoError(t, err)
+		assert.Equal(t, head.Height(), wantHead.Height())
 		assert.Equal(t, head.Hash(), wantHead.Hash())
 
 		// check that store height is aligned with the head.
@@ -253,6 +255,7 @@ func TestStore_Append_stableHeadWhenGaps(t *testing.T) {
 		// after appending missing headers we're on the latest header.
 		head, err := store.Head(ctx)
 		require.NoError(t, err)
+		assert.Equal(t, head.Height(), latestHead.Height())
 		assert.Equal(t, head.Hash(), latestHead.Hash())
 
 		// check that store height is aligned with the head.