From 7b5b2826a755a66e9c2c4a41bd93749bebebe36f Mon Sep 17 00:00:00 2001
From: Oleg Kovalov <oleg@hey.com>
Date: Fri, 24 Jan 2025 11:15:36 +0100
Subject: [PATCH] review suggestions

---
 store/store.go      | 24 ++++++------------------
 store/store_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/store/store.go b/store/store.go
index d7962bc8..4dd0fe10 100644
--- a/store/store.go
+++ b/store/store.go
@@ -165,19 +165,7 @@ func (s *Store[H]) Stop(ctx context.Context) error {
 }
 
 func (s *Store[H]) Height() uint64 {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-
-	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
-		}
-		panic(err)
-	}
-	return head.Height()
+	return s.heightSub.Height()
 }
 
 func (s *Store[H]) Head(ctx context.Context, _ ...header.HeadOption[H]) (H, error) {
@@ -231,8 +219,8 @@ func (s *Store[H]) GetByHeight(ctx context.Context, height uint64) (H, error) {
 		return zero, errors.New("header/store: height must be bigger than zero")
 	}
 
-	if h, err := s.getByHeight(ctx, height); err == nil || ctx.Err() != nil {
-		return h, err
+	if h, err := s.getByHeight(ctx, height); err == nil {
+		return h, nil
 	}
 
 	// if the requested 'height' was not yet published
@@ -537,17 +525,17 @@ func (s *Store[H]) advanceContiguousHead(ctx context.Context, headers ...H) {
 	}
 
 	if currHeight > prevHeight {
-		s.updateContiguousHead(newHead, currHeight)
+		s.updateContiguousHead(ctx, newHead, currHeight)
 	}
 }
 
-func (s *Store[H]) updateContiguousHead(newHead H, newHeight uint64) {
+func (s *Store[H]) updateContiguousHead(ctx context.Context, newHead H, newHeight uint64) {
 	s.contiguousHead.Store(&newHead)
 	s.heightSub.SetHeight(newHeight)
 	log.Infow("new head", "height", newHead.Height(), "hash", newHead.Hash())
 	s.metrics.newHead(newHead.Height())
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()
 
 	b, err := newHead.Hash().MarshalJSON()
diff --git a/store/store_test.go b/store/store_test.go
index 804d4c25..a152a42b 100644
--- a/store/store_test.go
+++ b/store/store_test.go
@@ -375,6 +375,34 @@ func TestStoreGetByHeight_whenGaps(t *testing.T) {
 	}
 }
 
+func TestStoreGetByHeight_earlyAvailable(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	t.Cleanup(cancel)
+
+	suite := headertest.NewTestSuite(t)
+
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	store := NewTestStore(t, ctx, ds, suite.Head(), WithWriteBatchSize(10))
+
+	const skippedHeaders = 15
+	suite.GenDummyHeaders(skippedHeaders)
+	lastChunk := suite.GenDummyHeaders(1)
+
+	{
+		err := store.Append(ctx, lastChunk...)
+		require.NoError(t, err)
+
+		// wait for batch to be written.
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	{
+		h, err := store.GetByHeight(ctx, lastChunk[0].Height())
+		require.NoError(t, err)
+		require.Equal(t, h, lastChunk[0])
+	}
+}
+
 // 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) {