perf(store): almost readless parallel deletes (#320)

Closes #318

With these optimizations and (one more coming after this PR), we get to
~106s to prune the entire ~6.5mil chain:
* Avoids reading the full header entirely and only reads index height ->
hash index.
* Allows reading headers ad hoc with OnDelete only when a user needs,
rather than always.
* Implements parallel deletion that is triggered when there are more
than 10000 deletions.

  
While tested manually, some cases miss coverage:  
* error cases
* context deadline handling
* parallel codepath (unit tested locally by manually changing threshold)

Author: Wondertan

headertest/store.go

@@ -20,7 +20,7 @@ type Store[H header.Header[H]] struct {
 	TailHeight uint64
 
 	onDeleteMu sync.Mutex
-	onDelete   []func(context.Context, []H) error
+	onDelete   []func(context.Context, uint64) error
 }
 
 // NewDummyStore creates a store for DummyHeader.
@@ -80,26 +80,26 @@ func (m *Store[H]) GetByHeight(_ context.Context, height uint64) (H, error) {
 }
 
 func (m *Store[H]) DeleteTo(ctx context.Context, to uint64) error {
-	var deleted []H
 	for h := m.TailHeight; h < to; h++ {
-		hdr, ok := m.Headers[h]
-		if ok {
-			delete(m.Headers, h)
-			deleted = append(deleted, hdr)
+		_, ok := m.Headers[h]
+		if !ok {
+			continue
 		}
-	}
 
-	m.TailHeight = to
-	for _, deleteFn := range m.onDelete {
-		err := deleteFn(ctx, deleted)
-		if err != nil {
-			return err
+		for _, deleteFn := range m.onDelete {
+			err := deleteFn(ctx, h)
+			if err != nil {
+				return err
+			}
 		}
+		delete(m.Headers, h) // must be after deleteFn
 	}
+
+	m.TailHeight = to
 	return nil
 }
 
-func (m *Store[H]) OnDelete(fn func(context.Context, []H) error) {
+func (m *Store[H]) OnDelete(fn func(context.Context, uint64) error) {
 	m.onDeleteMu.Lock()
 	defer m.onDeleteMu.Unlock()
 

interface.go

@@ -88,8 +88,10 @@ 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)
+	// OnDelete registers given handler to be called whenever a header with the height is being removed.
+	// OnDelete guarantees that the header is accessible for the handler with GetByHeight and is removed
+	// only after the handler terminates with nil error.
+	OnDelete(handler func(ctx context.Context, height uint64) error)
 }
 
 // Getter contains the behavior necessary for a component to retrieve

p2p/server_test.go

@@ -197,4 +197,4 @@ func (timeoutStore[H]) DeleteTo(ctx context.Context, _ uint64) error {
 	return ctx.Err()
 }
 
-func (timeoutStore[H]) OnDelete(fn func(context.Context, []H) error) {}
+func (timeoutStore[H]) OnDelete(fn func(context.Context, uint64) error) {}

store/height_indexer.go

@@ -34,7 +34,11 @@ func newHeightIndexer[H header.Header[H]](
 }
 
 // HashByHeight loads a header hash corresponding to the given height.
-func (hi *heightIndexer[H]) HashByHeight(ctx context.Context, h uint64) (header.Hash, error) {
+func (hi *heightIndexer[H]) HashByHeight(
+	ctx context.Context,
+	h uint64,
+	cache bool,
+) (header.Hash, error) {
 	if v, ok := hi.cache.Get(h); ok {
 		return v, nil
 	}
@@ -44,6 +48,8 @@ func (hi *heightIndexer[H]) HashByHeight(ctx context.Context, h uint64) (header.
 		return nil, err
 	}
 
-	hi.cache.Add(h, header.Hash(val))
+	if cache {
+		hi.cache.Add(h, header.Hash(val))
+	}
 	return val, nil
 }

store/store.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"slices"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -61,7 +60,7 @@ type Store[H header.Header[H]] struct {
 	syncCh chan chan struct{}
 
 	onDeleteMu sync.Mutex
-	onDelete   []func(context.Context, []H) error
+	onDelete   []func(context.Context, uint64) error
 
 	Params Parameters
 }
@@ -271,7 +270,7 @@ func (s *Store[H]) getByHeight(ctx context.Context, height uint64) (H, error) {
 		return h, nil
 	}
 
-	hash, err := s.heightIndex.HashByHeight(ctx, height)
+	hash, err := s.heightIndex.HashByHeight(ctx, height, true)
 	if err != nil {
 		var zero H
 		if errors.Is(err, datastore.ErrNotFound) {
@@ -357,149 +356,7 @@ 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
-	err := s.Sync(ctx)
-	if err != nil {
-		return err
-	}
-
-	head, err := s.Head(ctx)
-	if err != nil {
-		return fmt.Errorf("header/store: reading head: %w", err)
-	}
-	if head.Height()+1 < to {
-		_, err := s.getByHeight(ctx, to)
-		if err != nil {
-			return fmt.Errorf(
-				"header/store: delete to %d beyond current head(%d)",
-				to,
-				head.Height(),
-			)
-		}
-
-		//  if `to` is bigger than the current head and is stored - allow delete, making `to` a new head
-	}
-
-	tail, err := s.Tail(ctx)
-	if err != nil {
-		return fmt.Errorf("header/store: reading tail: %w", err)
-	}
-	if tail.Height() >= to {
-		return fmt.Errorf("header/store: delete to %d below current tail(%d)", to, tail.Height())
-	}
-
-	if err := s.deleteRange(ctx, tail.Height(), to); err != nil {
-		return fmt.Errorf("header/store: delete to height %d: %w", to, err)
-	}
-
-	if head.Height()+1 == to {
-		// this is the case where we have deleted all the headers
-		// wipe the store
-		if err := s.wipe(ctx); err != nil {
-			return fmt.Errorf("header/store: wipe: %w", err)
-		}
-	}
-
-	return nil
-}
-
-var maxHeadersLoadedPerDelete uint64 = 512
-
-func (s *Store[H]) deleteRange(ctx context.Context, from, to uint64) error {
-	log.Infow("deleting headers", "from_height", from, "to_height", to)
-	for from < to {
-		amount := min(to-from, maxHeadersLoadedPerDelete)
-		toDelete := from + amount
-		headers := make([]H, 0, amount)
-
-		for height := from; height < toDelete; height++ {
-			// take headers individually instead of range
-			// as getRangeByHeight can't deal with potentially missing headers
-			h, err := s.getByHeight(ctx, height)
-			if errors.Is(err, header.ErrNotFound) {
-				log.Warnf("header/store: attempt to delete header that's not found", height)
-				continue
-			}
-			if err != nil {
-				return fmt.Errorf("getting header while deleting: %w", err)
-			}
-
-			headers = append(headers, h)
-		}
-
-		batch, err := s.ds.Batch(ctx)
-		if err != nil {
-			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)
-			}
-			if err := batch.Delete(ctx, heightKey(h.Height())); err != nil {
-				return fmt.Errorf("delete height key (%d): %w", h.Height(), err)
-			}
-		}
-
-		err = s.setTail(ctx, batch, toDelete)
-		if err != nil {
-			return fmt.Errorf("setting tail to %d: %w", toDelete, err)
-		}
-
-		if err := batch.Commit(ctx); err != nil {
-			return fmt.Errorf("committing delete batch [%d:%d): %w", from, toDelete, err)
-		}
-
-		// cleanup caches after disk is flushed
-		for _, h := range headers {
-			s.cache.Remove(h.Hash().String())
-			s.heightIndex.cache.Remove(h.Height())
-		}
-		s.pending.DeleteRange(from, toDelete)
-
-		log.Infow("deleted header range", "from_height", from, "to_height", toDelete)
-
-		// move iterator
-		from = toDelete
-	}
-
-	return nil
-}
-
-func (s *Store[H]) setTail(ctx context.Context, batch datastore.Batch, to uint64) error {
+func (s *Store[H]) setTail(ctx context.Context, write datastore.Write, to uint64) error {
 	newTail, err := s.getByHeight(ctx, to)
 	if errors.Is(err, header.ErrNotFound) {
 		return nil
@@ -510,17 +367,16 @@ func (s *Store[H]) setTail(ctx context.Context, batch datastore.Batch, to uint64
 
 	// set directly to `to`, avoiding iteration in recedeTail
 	s.tailHeader.Store(&newTail)
-	if err := writeHeaderHashTo(ctx, batch, newTail, tailKey); err != nil {
+	if err := writeHeaderHashTo(ctx, write, newTail, tailKey); err != nil {
 		return fmt.Errorf("writing tailKey in batch: %w", err)
 	}
+	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() {
-		if err := writeHeaderHashTo(ctx, batch, newTail, headKey); err != nil {
+	head, _ := s.Head(ctx)
+	if head.IsZero() || to > head.Height() {
+		if err := writeHeaderHashTo(ctx, write, newTail, headKey); err != nil {
 			return fmt.Errorf("writing headKey in batch: %w", err)
 		}
 		s.contiguousHead.Store(&newTail)
@@ -828,7 +684,7 @@ func (s *Store[H]) deinit() {
 
 func writeHeaderHashTo[H header.Header[H]](
 	ctx context.Context,
-	batch datastore.Batch,
+	write datastore.Write,
 	h H,
 	key datastore.Key,
 ) error {
@@ -837,7 +693,7 @@ func writeHeaderHashTo[H header.Header[H]](
 		return err
 	}
 
-	if err := batch.Put(ctx, key, hashBytes); err != nil {
+	if err := write.Put(ctx, key, hashBytes); err != nil {
 		return err
 	}