provider: ResettableKeyStore (#1146)

Author: guillaumemichel

dual/provider/options.go

@@ -19,7 +19,7 @@ const (
 )
 
 type config struct {
-	keyStore *datastore.KeyStore
+	keyStore datastore.KeyStore
 
 	reprovideInterval [2]time.Duration // [0] = LAN, [1] = WAN
 	maxReprovideDelay [2]time.Duration
@@ -87,7 +87,7 @@ var DefaultConfig = func(cfg *config) error {
 	return nil
 }
 
-func WithKeyStore(keyStore *datastore.KeyStore) Option {
+func WithKeyStore(keyStore datastore.KeyStore) Option {
 	return func(cfg *config) error {
 		if keyStore == nil {
 			return errors.New("provider config: keyStore cannot be nil")

dual/provider/provider.go

@@ -19,7 +19,7 @@ type SweepingProvider struct {
 	dht      *dual.DHT
 	LAN      *provider.SweepingProvider
 	WAN      *provider.SweepingProvider
-	keyStore *datastore.KeyStore
+	keyStore datastore.KeyStore
 }
 
 // New creates a new SweepingProvider that manages provides and reprovides for

provider/datastore/keystore.go

@@ -3,9 +3,11 @@ package datastore
 import (
 	"context"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/ipfs/go-cid"
 	ds "github.com/ipfs/go-datastore"
@@ -18,8 +20,19 @@ import (
 	"github.com/probe-lab/go-libdht/kad/key/bitstr"
 )
 
-// KeyStore indexes multihashes by their kademlia identifier.
-type KeyStore struct {
+// KeyStore provides thread-safe storage and retrieval of multihashes, indexed
+// by their kademlia 256-bit identifier.
+type KeyStore interface {
+	Put(context.Context, ...mh.Multihash) ([]mh.Multihash, error)
+	Get(context.Context, bitstr.Key) ([]mh.Multihash, error)
+	ContainsPrefix(context.Context, bitstr.Key) (bool, error)
+	Delete(context.Context, ...mh.Multihash) error
+	Empty(context.Context) error
+	Size(context.Context) (int, error)
+}
+
+// keyStore indexes multihashes by their kademlia identifier.
+type keyStore struct {
 	lk sync.Mutex
 
 	ds         ds.Batching
@@ -91,50 +104,22 @@ func WithBatchSize(size int) KeyStoreOption {
 }
 
 // NewKeyStore creates a new KeyStore backed by the provided datastore.
-func NewKeyStore(d ds.Batching, opts ...KeyStoreOption) (*KeyStore, error) {
+func NewKeyStore(d ds.Batching, opts ...KeyStoreOption) (KeyStore, error) {
 	var cfg keyStoreCfg
 	opts = append([]KeyStoreOption{KeyStoreDefaultCfg}, opts...)
 	for i, o := range opts {
 		if err := o(&cfg); err != nil {
 			return nil, fmt.Errorf("KeyStore option %d failed: %w", i, err)
 		}
 	}
-	return &KeyStore{
+	return &keyStore{
 		ds:         d,
 		base:       ds.NewKey(cfg.base),
 		prefixBits: cfg.prefixBits,
 		batchSize:  cfg.batchSize,
 	}, nil
 }
 
-// ResetCids purges the KeyStore and repopulates it with the provided cids'
-// multihashes.
-func (s *KeyStore) ResetCids(ctx context.Context, keysChan <-chan cid.Cid) error {
-	s.lk.Lock()
-	defer s.lk.Unlock()
-
-	err := s.emptyLocked(ctx)
-	if err != nil {
-		return fmt.Errorf("KeyStore empty failed during reset: %w", err)
-	}
-	keys := make([]mh.Multihash, 0, s.batchSize)
-	for c := range keysChan {
-		keys = append(keys, c.Hash())
-		if len(keys) == cap(keys) {
-			_, err = s.putLocked(ctx, keys...)
-			if err != nil {
-				return fmt.Errorf("KeyStore put failed during reset: %w", err)
-			}
-			keys = keys[:0]
-		}
-	}
-	_, err = s.putLocked(ctx, keys...)
-	if err != nil {
-		return fmt.Errorf("KeyStore put failed during reset: %w", err)
-	}
-	return nil
-}
-
 // dsKey returns the datastore key for the provided binary key.
 //
 // The function creates a hierarchical datastore key by expanding bits into
@@ -179,7 +164,7 @@ func dsKey[K kad.Key[K]](k K, prefixBits int, base ds.Key) ds.Key {
 // "base/bit0/bit1/.../bitN/base64url_suffix"
 //
 // Returns the reconstructed 256-bit key or an error if base64URL decoding fails.
-func (s *KeyStore) decodeKey(dsk string) (bit256.Key, error) {
+func (s *keyStore) decodeKey(dsk string) (bit256.Key, error) {
 	dsk = dsk[len(s.base.String()):] // remove leading prefix
 	bs := make([]byte, 32)
 	// Extract individual bits from odd positions (skip '/' separators)
@@ -204,7 +189,7 @@ type pair struct {
 
 // putLocked stores the provided keys while assuming s.lk is already held, and
 // returns the keys that weren't present already in the keystore.
-func (s *KeyStore) putLocked(ctx context.Context, keys ...mh.Multihash) ([]mh.Multihash, error) {
+func (s *keyStore) putLocked(ctx context.Context, keys ...mh.Multihash) ([]mh.Multihash, error) {
 	seen := make(map[bit256.Key]struct{}, len(keys))
 	toPut := make([]pair, 0, len(keys))
 
@@ -252,7 +237,7 @@ func (s *KeyStore) putLocked(ctx context.Context, keys ...mh.Multihash) ([]mh.Mu
 // Put stores the provided keys in the underlying datastore, grouping them by
 // the first prefixLen bits. It returns only the keys that were not previously
 // persisted in the datastore (i.e., newly added keys).
-func (s *KeyStore) Put(ctx context.Context, keys ...mh.Multihash) ([]mh.Multihash, error) {
+func (s *keyStore) Put(ctx context.Context, keys ...mh.Multihash) ([]mh.Multihash, error) {
 	if len(keys) == 0 {
 		return nil, nil
 	}
@@ -264,7 +249,7 @@ func (s *KeyStore) Put(ctx context.Context, keys ...mh.Multihash) ([]mh.Multihas
 
 // Get returns all keys whose bit256 representation matches the provided
 // prefix.
-func (s *KeyStore) Get(ctx context.Context, prefix bitstr.Key) ([]mh.Multihash, error) {
+func (s *keyStore) Get(ctx context.Context, prefix bitstr.Key) ([]mh.Multihash, error) {
 	s.lk.Lock()
 	defer s.lk.Unlock()
 
@@ -300,7 +285,7 @@ func (s *KeyStore) Get(ctx context.Context, prefix bitstr.Key) ([]mh.Multihash,
 // ContainsPrefix reports whether the KeyStore currently holds at least one
 // multihash whose kademlia identifier (bit256.Key) starts with the provided
 // bit-prefix.
-func (s *KeyStore) ContainsPrefix(ctx context.Context, prefix bitstr.Key) (bool, error) {
+func (s *keyStore) ContainsPrefix(ctx context.Context, prefix bitstr.Key) (bool, error) {
 	s.lk.Lock()
 	defer s.lk.Unlock()
 
@@ -337,8 +322,8 @@ func (s *KeyStore) ContainsPrefix(ctx context.Context, prefix bitstr.Key) (bool,
 
 // emptyLocked deletes all entries under the datastore prefix, assuming s.lk is
 // already held.
-func (s *KeyStore) emptyLocked(ctx context.Context) error {
-	res, err := s.ds.Query(ctx, query.Query{Prefix: s.base.String(), KeysOnly: true})
+func emptyLocked(ctx context.Context, d ds.Batching, base ds.Key, batchSize int) error {
+	res, err := d.Query(ctx, query.Query{Prefix: base.String(), KeysOnly: true})
 	if err != nil {
 		return err
 	}
@@ -348,7 +333,7 @@ func (s *KeyStore) emptyLocked(ctx context.Context) error {
 		if len(keys) == 0 {
 			return nil
 		}
-		b, err := s.ds.Batch(ctx)
+		b, err := d.Batch(ctx)
 		if err != nil {
 			return nil
 		}
@@ -359,7 +344,7 @@ func (s *KeyStore) emptyLocked(ctx context.Context) error {
 		}
 		return b.Commit(ctx)
 	}
-	keys := make([]ds.Key, 0, s.batchSize)
+	keys := make([]ds.Key, 0, batchSize)
 	for r := range res.Next() {
 		if r.Error != nil {
 			return r.Error
@@ -376,15 +361,15 @@ func (s *KeyStore) emptyLocked(ctx context.Context) error {
 }
 
 // Empty deletes all entries under the datastore prefix.
-func (s *KeyStore) Empty(ctx context.Context) error {
+func (s *keyStore) Empty(ctx context.Context) error {
 	s.lk.Lock()
 	defer s.lk.Unlock()
 
-	return s.emptyLocked(ctx)
+	return emptyLocked(ctx, s.ds, s.base, s.batchSize)
 }
 
 // Delete removes the given keys from datastore.
-func (s *KeyStore) Delete(ctx context.Context, keys ...mh.Multihash) error {
+func (s *keyStore) Delete(ctx context.Context, keys ...mh.Multihash) error {
 	if len(keys) == 0 {
 		return nil
 	}
@@ -409,7 +394,7 @@ func (s *KeyStore) Delete(ctx context.Context, keys ...mh.Multihash) error {
 //
 // The size is obtained by iterating over all keys in the underlying
 // datastore, so it may be expensive for large stores.
-func (s *KeyStore) Size(ctx context.Context) (size int, err error) {
+func (s *keyStore) Size(ctx context.Context) (size int, err error) {
 	s.lk.Lock()
 	defer s.lk.Unlock()
 
@@ -429,3 +414,158 @@ func (s *KeyStore) Size(ctx context.Context) (size int, err error) {
 	}
 	return
 }
+
+var ErrResetInProgress = errors.New("reset already in progress")
+
+const (
+	ready uint32 = iota
+	inProgress
+	cleanup
+)
+
+// ResettableKeyStore is a KeyStore implementation that supports atomic reset
+// operations. It maintains two alternate bases in the underlying datastore and
+// can swap between them to provide atomic replacement of all stored keys.
+//
+// The reset operation allows replacing all stored multihashes with a new set
+// without interrupting concurrent read/write operations. During a reset, new
+// writes are duplicated to both the current and alternate storage bases to
+// maintain consistency.
+type ResettableKeyStore struct {
+	keyStore
+
+	altBase ds.Key
+	// resetInProgress tracks the state of the reset operation:
+	//   - 0: no reset in progress, ready for new reset
+	//   - 1: reset in progress, writing to alternate base
+	//   - 2: post reset cleanup
+	resetState atomic.Uint32
+}
+
+var _ KeyStore = (*ResettableKeyStore)(nil)
+
+// NewResettableKeyStore creates a new ResettableKeyStore backed by the
+// provided datastore. It automatically adds "/0" and "/1" suffixes to the
+// configured base prefix to create two alternate storage locations for atomic
+// reset operations.
+func NewResettableKeyStore(d ds.Batching, opts ...KeyStoreOption) (*ResettableKeyStore, error) {
+	var cfg keyStoreCfg
+	opts = append([]KeyStoreOption{KeyStoreDefaultCfg}, opts...)
+	for i, o := range opts {
+		if err := o(&cfg); err != nil {
+			return nil, fmt.Errorf("KeyStore option %d failed: %w", i, err)
+		}
+	}
+	return &ResettableKeyStore{
+		keyStore: keyStore{
+			ds:         d,
+			base:       ds.NewKey(cfg.base).ChildString("0"),
+			prefixBits: cfg.prefixBits,
+			batchSize:  cfg.batchSize,
+		},
+		altBase: ds.NewKey(cfg.base).ChildString("1"),
+	}, nil
+}
+
+func batchPut(ctx context.Context, d ds.Batching, base ds.Key, prefixBits int, keys []mh.Multihash) error {
+	b, err := d.Batch(ctx)
+	if err != nil {
+		return err
+	}
+	for _, h := range keys {
+		dsk := dsKey(keyspace.MhToBit256(h), prefixBits, base)
+		if err := b.Put(ctx, dsk, h); err != nil {
+			return err
+		}
+	}
+	return b.Commit(ctx)
+}
+
+// ResetCids atomically replaces all stored keys with the CIDs received from
+// keysChan. The operation is thread-safe and non-blocking for concurrent reads
+// and writes.
+//
+// During the reset:
+//   - New keys from keysChan are written to an alternate storage location
+//   - Concurrent Put operations are duplicated to both current and alternate
+//     locations
+//   - Once all keys are processed, storage locations are atomically swapped
+//   - The old storage location is cleaned up
+//
+// Returns ErrResetInProgress if another reset operation is already running.
+// The operation can be cancelled via context, which will clean up partial
+// state.
+func (s *ResettableKeyStore) ResetCids(ctx context.Context, keysChan <-chan cid.Cid) error {
+	if !s.resetState.CompareAndSwap(ready, inProgress) {
+		return ErrResetInProgress
+	}
+	defer s.resetState.Store(ready)
+
+	// Make sure the alternate base is empty
+	if err := emptyLocked(ctx, s.ds, s.altBase, s.batchSize); err != nil {
+		return fmt.Errorf("ResettableKeyStore empty failed during reset: %w", err)
+	}
+
+	// Read all the keys from the channel and put them in batch at the alternate
+	// base
+	keys := make([]mh.Multihash, 0, s.batchSize)
+loop:
+	for {
+		select {
+		case <-ctx.Done():
+			// Context cancelled, abort reset and clean up alternate base
+			s.resetState.Store(cleanup)
+			if err := emptyLocked(ctx, s.ds, s.altBase, s.batchSize); err != nil {
+				return fmt.Errorf("ResettableKeyStore empty failed during reset: %w", err)
+			}
+			return nil
+		case c, ok := <-keysChan:
+			if !ok {
+				break loop
+			}
+			keys = append(keys, c.Hash())
+			if len(keys) == cap(keys) {
+				if err := batchPut(ctx, s.ds, s.altBase, s.prefixBits, keys); err != nil {
+					return fmt.Errorf("ResettableKeyStore put failed during reset: %w", err)
+				}
+				keys = keys[:0]
+			}
+		}
+	}
+	// Put final batch
+	if err := batchPut(ctx, s.ds, s.altBase, s.prefixBits, keys); err != nil {
+		return fmt.Errorf("ResettableKeyStore put failed during reset: %w", err)
+	}
+
+	// Swap the bases. The base in use is now the one to which we just wrote.
+	s.lk.Lock()
+	oldBase := s.base
+	s.base = s.altBase
+	s.resetState.Store(cleanup)
+	s.lk.Unlock()
+	s.altBase = oldBase
+
+	// Clean up the old base
+	if err := emptyLocked(ctx, s.ds, oldBase, s.batchSize); err != nil {
+		return fmt.Errorf("ResettableKeyStore empty failed during reset: %w", err)
+	}
+
+	return nil
+}
+
+// Put stores the provided keys in the underlying datastore, with special
+// handling during reset operations. When a reset is in progress, keys are
+// written to both the current storage location and the alternate location
+// being prepared for the reset.
+//
+// This ensures that keys written during a reset operation remain available
+// after the reset completes. Returns only the keys that were not previously
+// persisted.
+func (s *ResettableKeyStore) Put(ctx context.Context, keys ...mh.Multihash) ([]mh.Multihash, error) {
+	if s.resetState.Load() == inProgress {
+		// Reset is in progress, write to alternate base in addition to current
+		// base.
+		batchPut(ctx, s.ds, s.altBase, s.prefixBits, keys)
+	}
+	return s.keyStore.Put(ctx, keys...)
+}

provider/datastore/keystore_test.go

@@ -151,7 +151,7 @@ func TestKeyStoreReset(t *testing.T) {
 	ds := ds.NewMapDatastore()
 	defer ds.Close()
 
-	store, err := NewKeyStore(ds)
+	store, err := NewResettableKeyStore(ds)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -289,7 +289,7 @@ func TestKeyStoreSize(t *testing.T) {
 }
 
 func TestDsKey(t *testing.T) {
-	s := KeyStore{
+	s := keyStore{
 		base:       ds.NewKey("/base/prefix"),
 		prefixBits: 8,
 	}