feat(store): store recovery

Author: Wondertan

store/store.go

@@ -555,12 +555,9 @@ func (s *Store[H]) setTail(ctx context.Context, batch datastore.Batch, to uint64
 	log.Infow("new tail", "height", newTail.Height(), "hash", newTail.Hash())
 	s.metrics.newTail(newTail.Height())
 
-	// update head as well, if delete went over it
-	head, err := s.Head(ctx)
-	if err != nil {
-		return err
-	}
-	if to > head.Height() {
+	// update head as well, if delete went over it or head doesn't exist
+	head, _ := s.Head(ctx)
+	if head.IsZero() || to > head.Height() {
 		if err := writeHeaderHashTo(ctx, batch, newTail, headKey); err != nil {
 			return fmt.Errorf("writing headKey in batch: %w", err)
 		}
@@ -797,14 +794,16 @@ func (s *Store[H]) nextHead(ctx context.Context) (head H, changed bool) {
 		return head, false
 	}
 
-	for {
+	for ctx.Err() == nil {
 		h, err := s.getByHeight(ctx, head.Height()+1)
 		if err != nil {
 			return head, changed
 		}
 		head = h
 		changed = true
 	}
+
+	return head, changed
 }
 
 // nextTail finds the new contiguous Tail by iterating the current Tail down until the older height Tail is found.
@@ -816,14 +815,16 @@ func (s *Store[H]) nextTail(ctx context.Context) (tail H, changed bool) {
 		return tail, false
 	}
 
-	for {
+	for ctx.Err() == nil {
 		h, err := s.getByHeight(ctx, tail.Height()-1)
 		if err != nil {
 			return tail, changed
 		}
 		tail = h
 		changed = true
 	}
+
+	return tail, changed
 }
 
 func (s *Store[H]) loadHeadAndTail(ctx context.Context) error {

store/store_recover.go

@@ -0,0 +1,68 @@
+package store
+
+import (
+	"context"
+
+	"github.com/celestiaorg/go-header"
+)
+
+// ResetTail resets the tail of the store to be at the given height.
+// The new tail must be present in the store.
+// WARNING: Only use this function if you know what you are doing.
+func ResetTail[H header.Header[H]](ctx context.Context, store *Store[H], height uint64) error {
+	batch, err := store.ds.Batch(ctx)
+	if err != nil {
+		return err
+	}
+
+	err = store.setTail(ctx, batch, height)
+	if err != nil {
+		return err
+	}
+
+	err = batch.Commit(ctx)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// ResetHead resets the head of the store to be at the given height.
+// The new head must be present in the store.
+// WARNING: Only use this function if you know what you are doing.
+func ResetHead[H header.Header[H]](ctx context.Context, store *Store[H], height uint64) error {
+	batch, err := store.ds.Batch(ctx)
+	if err != nil {
+		return err
+	}
+
+	newHead, err := store.getByHeight(ctx, height)
+	if err != nil {
+		return err
+	}
+
+	if err := writeHeaderHashTo(ctx, batch, newHead, headKey); err != nil {
+		return err
+	}
+	store.contiguousHead.Store(&newHead)
+
+	err = batch.Commit(ctx)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// FindHeader forward iterates over the store starting from the given height until it finds any stored header
+// or the context is cancelled.
+func FindHeader[H header.Header[H]](ctx context.Context, store *Store[H], startFrom uint64) (H, error) {
+	for height := startFrom; ctx.Err() == nil; height++ {
+		header, err := store.getByHeight(ctx, height)
+		if err == nil {
+			return header, nil
+		}
+	}
+
+	var zero H
+	return zero, ctx.Err()
+}