From 09d798e59c4863067b3558e17cd0646cca16e947 Mon Sep 17 00:00:00 2001
From: Oleg Kovalov <oleg@hey.com>
Date: Tue, 15 Oct 2024 12:18:59 +0200
Subject: [PATCH] fix review suggestions

---
 store/heightsub.go      | 7 ++++++-
 store/heightsub_test.go | 1 +
 store/store.go          | 4 +++-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/store/heightsub.go b/store/heightsub.go
index 772933a4..a27c3b30 100644
--- a/store/heightsub.go
+++ b/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()
diff --git a/store/heightsub_test.go b/store/heightsub_test.go
index 88f4434b..92eecec4 100644
--- a/store/heightsub_test.go
+++ b/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()
diff --git a/store/store.go b/store/store.go
index d7e700c1..0a0acaae 100644
--- a/store/store.go
+++ b/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.