fix review suggestions

store/heightsub.go

@@ -3,6 +3,7 @@ package store
 import (
 	"context"
 	"errors"
+	"fmt"
 	"sync"
 	"sync/atomic"
 
@@ -33,7 +34,7 @@ func (hs *heightSub[H]) isInited() bool {
 }
 
 // setHeight sets the new head height for heightSub.
-// Only higher then current height will be set.
+// Only higher than current height can be set.
 func (hs *heightSub[H]) setHeight(height uint64) {
 	for {
 		curr := hs.height.Load()
@@ -98,6 +99,10 @@ func (hs *heightSub[H]) Pub(headers ...H) {
 	}
 
 	from, to := headers[0].Height(), headers[ln-1].Height()
+	if from >= to {
+		panic(fmt.Sprintf("from must be lower than to, have: %d and %d", from, to))
+	}
+
 	hs.setHeight(to)
 
 	hs.heightReqsLk.Lock()

store/heightsub_test.go

@@ -47,6 +47,7 @@ func TestHeightSub(t *testing.T) {
 	}
 }
 
+// Test heightSub can accept non-adj headers without a problem.
 func TestHeightSubNonAdjacement(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()

store/store.go

@@ -173,6 +173,7 @@ func (s *Store[H]) Height() uint64 {
 	head, err := s.Head(ctx)
 	if err != nil {
 		if errors.Is(err, context.Canceled) ||
+			errors.Is(err, context.DeadlineExceeded) ||
 			errors.Is(err, datastore.ErrNotFound) {
 			return 0
 		}
@@ -528,7 +529,8 @@ func (s *Store[H]) tryAdvanceHead(ctx context.Context, headers ...H) {
 		currHeight++
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, time.Second)
+	// TODO(cristaloleg): benchmark this timeout or make it dynamic.
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	// advance based on already written headers.