feat!(store): add DeleteTo and update tail accordingly (#275)

Fixes #257

Author: cristaloleg

headertest/store.go

@@ -16,6 +16,7 @@ type Generator[H header.Header[H]] interface {
 type Store[H header.Header[H]] struct {
 	Headers    map[uint64]H
 	HeadHeight uint64
+	TailHeight uint64
 }
 
 // NewDummyStore creates a store for DummyHeader.
@@ -52,16 +53,12 @@ func (m *Store[H]) Head(context.Context, ...header.HeadOption[H]) (H, error) {
 }
 
 func (m *Store[H]) Tail(context.Context) (H, error) {
-	err := header.ErrNotFound
-
-	var tail H
-	for _, h := range m.Headers {
-		if tail.IsZero() || h.Height() < tail.Height() {
-			tail = h
-		}
+	tail, ok := m.Headers[m.TailHeight]
+	if !ok {
+		var zero H
+		return zero, header.ErrNotFound
 	}
-
-	return tail, err
+	return tail, nil
 }
 
 func (m *Store[H]) Get(_ context.Context, hash header.Hash) (H, error) {
@@ -82,6 +79,15 @@ 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 {
+	for h := m.TailHeight; h < to; h++ {
+		delete(m.Headers, h)
+	}
+
+	m.TailHeight = to
+	return nil
+}
+
 func (m *Store[H]) GetRange(ctx context.Context, from, to uint64) ([]H, error) {
 	return m.getRangeByHeight(ctx, from, to)
 }
@@ -122,10 +128,13 @@ func (m *Store[H]) HasAt(_ context.Context, height uint64) bool {
 
 func (m *Store[H]) Append(_ context.Context, headers ...H) error {
 	for _, header := range headers {
-		m.Headers[header.Height()] = header
-		// set head
-		if header.Height() > m.HeadHeight {
-			m.HeadHeight = header.Height()
+		height := header.Height()
+		m.Headers[height] = header
+		if height > m.HeadHeight {
+			m.HeadHeight = height
+		}
+		if height < m.TailHeight {
+			m.TailHeight = height
 		}
 	}
 	return nil

interface.go

@@ -87,6 +87,9 @@ type Store[H Header[H]] interface {
 
 	// GetRange returns the range [from:to).
 	GetRange(context.Context, uint64, uint64) ([]H, error)
+
+	// DeleteTo deletes the range [Tail():to).
+	DeleteTo(ctx context.Context, to uint64) error
 }
 
 // Getter contains the behavior necessary for a component to retrieve

p2p/server_test.go

@@ -190,3 +190,8 @@ func (timeoutStore[H]) GetRange(ctx context.Context, _, _ uint64) ([]H, error) {
 	<-ctx.Done()
 	return nil, ctx.Err()
 }
+
+func (timeoutStore[H]) DeleteTo(ctx context.Context, _ uint64) error {
+	<-ctx.Done()
+	return ctx.Err()
+}

store/batch.go

@@ -1,6 +1,8 @@
 package store
 
 import (
+	"maps"
+	"slices"
 	"sync"
 
 	"github.com/celestiaorg/go-header"
@@ -101,6 +103,21 @@ func (b *batch[H]) Has(hash header.Hash) bool {
 	return ok
 }
 
+// DeleteRange of headers from the batch.
+func (b *batch[H]) DeleteRange(from, to uint64) {
+	b.lk.Lock()
+	defer b.lk.Unlock()
+
+	maps.DeleteFunc(b.heights, func(_ string, height uint64) bool {
+		return from <= height && height < to
+	})
+
+	b.headers = slices.DeleteFunc(b.headers, func(h H) bool {
+		height := h.Height()
+		return from <= height && height < to
+	})
+}
+
 // Reset cleans references to batched headers.
 func (b *batch[H]) Reset() {
 	b.lk.Lock()

store/height_indexer.go

@@ -2,6 +2,7 @@ package store
 
 import (
 	"context"
+	"fmt"
 
 	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/ipfs/go-datastore"
@@ -47,3 +48,16 @@ func (hi *heightIndexer[H]) HashByHeight(ctx context.Context, h uint64) (header.
 	hi.cache.Add(h, header.Hash(val))
 	return val, nil
 }
+
+// deleteRange of heights from the index.
+func (hi *heightIndexer[H]) deleteRange(
+	ctx context.Context, batch datastore.Batch, from, to uint64,
+) error {
+	for h := from; h < to; h++ {
+		if err := batch.Delete(ctx, heightKey(h)); err != nil {
+			return fmt.Errorf("delete height key(%d): %w", h, err)
+		}
+		hi.cache.Remove(h)
+	}
+	return nil
+}

store/keys.go

@@ -18,5 +18,9 @@ func heightKey(h uint64) datastore.Key {
 }
 
 func headerKey[H header.Header[H]](h H) datastore.Key {
-	return datastore.NewKey(h.Hash().String())
+	return hashKey(h.Hash())
+}
+
+func hashKey(h header.Hash) datastore.Key {
+	return datastore.NewKey(h.String())
 }

store/store.go

@@ -345,6 +345,106 @@ func (s *Store[H]) HasAt(_ context.Context, height uint64) bool {
 	return height != uint64(0) && s.Height() >= height
 }
 
+// DeleteTo implements [header.Store] interface.
+func (s *Store[H]) DeleteTo(ctx context.Context, to uint64) error {
+	var from uint64
+
+	if tailPtr := s.tailHeader.Load(); tailPtr != nil {
+		from = (*tailPtr).Height()
+	}
+	if from >= to {
+		return nil
+	}
+	if headPtr := s.contiguousHead.Load(); headPtr != nil {
+		if height := (*headPtr).Height(); to > height {
+			return fmt.Errorf("header/store: higher then head (%d vs %d)", to, height)
+		}
+	}
+
+	if from >= to {
+		log.Debugw("header/store: attempt to delete empty range(%d, %d)", from, to)
+		return nil
+	}
+
+	if err := s.deleteRange(ctx, from, to); err != nil {
+		return fmt.Errorf("header/store: delete to height %d: %w", to, err)
+	}
+	return nil
+}
+
+func (s *Store[H]) deleteRange(ctx context.Context, from, to uint64) error {
+	batch, err := s.ds.Batch(ctx)
+	if err != nil {
+		return fmt.Errorf("delete batch: %w", err)
+	}
+
+	if err := s.prepareDeleteRangeBatch(ctx, batch, from, to); err != nil {
+		return fmt.Errorf("prepare: %w", err)
+	}
+
+	if err := s.heightIndex.deleteRange(ctx, batch, from, to); err != nil {
+		return fmt.Errorf("height index: %w", err)
+	}
+
+	newTail, err := s.updateTail(ctx, batch, to)
+	if err != nil {
+		return fmt.Errorf("update tail: %w", err)
+	}
+
+	if err := batch.Commit(ctx); err != nil {
+		return fmt.Errorf("delete commit: %w", err)
+	}
+
+	s.tailHeader.Store(&newTail)
+	return nil
+}
+
+func (s *Store[H]) prepareDeleteRangeBatch(
+	ctx context.Context, batch datastore.Batch, from, to uint64,
+) error {
+	for h := from; h < to; h++ {
+		hash, err := s.heightIndex.HashByHeight(ctx, h)
+		if err != nil {
+			if errors.Is(err, datastore.ErrNotFound) {
+				log.Errorw("removing non-existent header", "height", h)
+				continue
+			}
+			return fmt.Errorf("hash by height(%d): %w", h, err)
+		}
+		s.cache.Remove(hash.String())
+
+		if err := batch.Delete(ctx, hashKey(hash)); err != nil {
+			return fmt.Errorf("delete hash key: %w", err)
+		}
+	}
+
+	s.pending.DeleteRange(from, to)
+	return nil
+}
+
+func (s *Store[H]) updateTail(
+	ctx context.Context, batch datastore.Batch, to uint64,
+) (H, error) {
+	var zero H
+
+	newTail, err := s.getByHeight(ctx, to)
+	if err != nil {
+		if !errors.Is(err, header.ErrNotFound) {
+			return zero, fmt.Errorf("cannot fetch next tail: %w", err)
+		}
+		return zero, err
+	}
+
+	b, err := newTail.Hash().MarshalJSON()
+	if err != nil {
+		return zero, err
+	}
+	if err := batch.Put(ctx, tailKey, b); err != nil {
+		return zero, err
+	}
+	return newTail, nil
+}
+
 func (s *Store[H]) Append(ctx context.Context, headers ...H) error {
 	lh := len(headers)
 	if lh == 0 {

store/store_test.go

@@ -29,8 +29,6 @@ func TestStore(t *testing.T) {
 
 	assert.Equal(t, *store.tailHeader.Load(), suite.Head())
 
-	assert.Equal(t, *store.tailHeader.Load(), suite.Head())
-
 	head, err := store.Head(ctx)
 	require.NoError(t, err)
 	assert.EqualValues(t, suite.Head().Hash(), head.Hash())
@@ -470,6 +468,94 @@ func TestStore_GetRange(t *testing.T) {
 	}
 }
 
+func TestStore_DeleteTo(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	t.Cleanup(cancel)
+
+	suite := headertest.NewTestSuite(t)
+
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	store := NewTestStore(t, ctx, ds, suite.Head(), WithWriteBatchSize(10))
+
+	const count = 100
+	in := suite.GenDummyHeaders(count)
+	err := store.Append(ctx, in...)
+	require.NoError(t, err)
+
+	hashes := make(map[uint64]header.Hash, count)
+	for _, h := range in {
+		hashes[h.Height()] = h.Hash()
+	}
+
+	// wait until headers are written
+	time.Sleep(100 * time.Millisecond)
+
+	tests := []struct {
+		name      string
+		to        uint64
+		wantTail  uint64
+		wantError bool
+	}{
+		{
+			name:      "initial delete request",
+			to:        14,
+			wantTail:  14,
+			wantError: false,
+		},
+		{
+			name:      "no-op delete request",
+			to:        5,
+			wantTail:  14,
+			wantError: false,
+		},
+		{
+			name:      "valid delete request",
+			to:        50,
+			wantTail:  50,
+			wantError: false,
+		},
+		{
+			name:      "higher than head",
+			to:        1055,
+			wantTail:  30,
+			wantError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			from := (*store.tailHeader.Load()).Height()
+
+			// manually add something to the pending for assert at the bottom
+			if idx := from - 2; idx < count {
+				store.pending.Append(in[idx])
+				defer store.pending.Reset()
+			}
+
+			ctx, cancel := context.WithTimeout(ctx, time.Second)
+			defer cancel()
+
+			err := store.DeleteTo(ctx, tt.to)
+			if tt.wantError {
+				assert.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+
+			// check that cache and pending doesn't contain old headers
+			for h := from; h < tt.to; h++ {
+				hash := hashes[h]
+				assert.False(t, store.cache.Contains(hash.String()))
+				assert.False(t, store.pending.Has(hash))
+			}
+
+			tail, err := store.Tail(ctx)
+			require.NoError(t, err)
+			require.EqualValues(t, tail.Height(), tt.wantTail)
+		})
+	}
+}
+
 func TestStorePendingCacheMiss(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	t.Cleanup(cancel)