feat(store)!: OnDelete method (#304)

Closes #301 
Depends on #302

Author: Wondertan

headertest/store.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/celestiaorg/go-header"
@@ -17,6 +18,9 @@ type Store[H header.Header[H]] struct {
 	Headers    map[uint64]H
 	HeadHeight uint64
 	TailHeight uint64
+
+	onDeleteMu sync.Mutex
+	onDelete   []func(context.Context, []H) error
 }
 
 // NewDummyStore creates a store for DummyHeader.
@@ -75,15 +79,33 @@ func (m *Store[H]) GetByHeight(_ context.Context, height uint64) (H, error) {
 	return zero, header.ErrNotFound
 }
 
-func (m *Store[H]) DeleteTo(_ context.Context, to uint64) error {
+func (m *Store[H]) DeleteTo(ctx context.Context, to uint64) error {
+	var deleted []H
 	for h := m.TailHeight; h < to; h++ {
-		delete(m.Headers, h)
+		hdr, ok := m.Headers[h]
+		if ok {
+			delete(m.Headers, h)
+			deleted = append(deleted, hdr)
+		}
 	}
 
 	m.TailHeight = to
+	for _, deleteFn := range m.onDelete {
+		err := deleteFn(ctx, deleted)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
+func (m *Store[H]) OnDelete(fn func(context.Context, []H) error) {
+	m.onDeleteMu.Lock()
+	defer m.onDeleteMu.Unlock()
+
+	m.onDelete = append(m.onDelete, fn)
+}
+
 func (m *Store[H]) GetRange(ctx context.Context, from, to uint64) ([]H, error) {
 	return m.getRangeByHeight(ctx, from, to)
 }

interface.go

@@ -87,6 +87,9 @@ type Store[H Header[H]] interface {
 
 	// DeleteTo deletes the range [Tail():to).
 	DeleteTo(ctx context.Context, to uint64) error
+
+	// OnDelete registers given handler to be called whenever headers are removed from the Store.
+	OnDelete(func(context.Context, []H) error)
 }
 
 // Getter contains the behavior necessary for a component to retrieve

p2p/server_test.go

@@ -196,3 +196,5 @@ func (timeoutStore[H]) DeleteTo(ctx context.Context, _ uint64) error {
 	<-ctx.Done()
 	return ctx.Err()
 }
+
+func (timeoutStore[H]) OnDelete(fn func(context.Context, []H) error) {}

store/store.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"slices"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -58,6 +60,9 @@ type Store[H header.Header[H]] struct {
 	// syncCh is a channel used to synchronize writes
 	syncCh chan chan struct{}
 
+	onDeleteMu sync.Mutex
+	onDelete   []func(context.Context, []H) error
+
 	Params Parameters
 }
 
@@ -252,6 +257,16 @@ func (s *Store[H]) GetByHeight(ctx context.Context, height uint64) (H, error) {
 }
 
 func (s *Store[H]) getByHeight(ctx context.Context, height uint64) (H, error) {
+	head, _ := s.Head(ctx)
+	if !head.IsZero() && head.Height() == height {
+		return head, nil
+	}
+
+	tail, _ := s.Tail(ctx)
+	if !tail.IsZero() && tail.Height() == height {
+		return tail, nil
+	}
+
 	if h := s.pending.GetByHeight(height); !h.IsZero() {
 		return h, nil
 	}
@@ -342,6 +357,21 @@ func (s *Store[H]) HasAt(ctx context.Context, height uint64) bool {
 	return head.Height() >= height && height >= tail.Height()
 }
 
+func (s *Store[H]) OnDelete(fn func(context.Context, []H) error) {
+	s.onDeleteMu.Lock()
+	defer s.onDeleteMu.Unlock()
+
+	s.onDelete = append(s.onDelete, func(ctx context.Context, h []H) (rerr error) {
+		defer func() {
+			err := recover()
+			if err != nil {
+				rerr = fmt.Errorf("header/store: user provided onDelete panicked with: %s", err)
+			}
+		}()
+		return fn(ctx, h)
+	})
+}
+
 // DeleteTo implements [header.Store] interface.
 func (s *Store[H]) DeleteTo(ctx context.Context, to uint64) error {
 	// ensure all the pending headers are synchronized
@@ -419,6 +449,22 @@ func (s *Store[H]) deleteRange(ctx context.Context, from, to uint64) error {
 			return fmt.Errorf("new batch: %w", err)
 		}
 
+		s.onDeleteMu.Lock()
+		onDelete := slices.Clone(s.onDelete)
+		s.onDeleteMu.Unlock()
+		for _, deleteFn := range onDelete {
+			if err := deleteFn(ctx, headers); err != nil {
+				// abort deletion if onDelete handler fails
+				// to ensure atomicity between stored headers and user specific data
+				// TODO(@Wondertan): Batch is not actually atomic and could write some data at this point
+				//  but its fine for now: https://github.com/celestiaorg/go-header/issues/307
+				// TODO2(@Wondertan): Once we move to txn, find a way to pass txn through context,
+				//  so that users can use it in their onDelete handlers
+				//  to ensure atomicity between deleted headers and user specific data
+				return fmt.Errorf("on delete handler: %w", err)
+			}
+		}
+
 		for _, h := range headers {
 			if err := batch.Delete(ctx, hashKey(h.Hash())); err != nil {
 				return fmt.Errorf("delete hash key (%X): %w", h.Hash(), err)