use datastore namespace

Author: guillaumemichel

provider/datastore/keystore.go

@@ -8,7 +8,8 @@ import (
 	"strings"
 
 	ds "github.com/ipfs/go-datastore"
-	query "github.com/ipfs/go-datastore/query"
+	"github.com/ipfs/go-datastore/namespace"
+	"github.com/ipfs/go-datastore/query"
 	"github.com/libp2p/go-libp2p-kad-dht/provider/internal/keyspace"
 	mh "github.com/multiformats/go-multihash"
 
@@ -64,7 +65,6 @@ type operationResponse struct {
 // keyStore indexes multihashes by their kademlia identifier.
 type keyStore struct {
 	ds         ds.Batching
-	base       ds.Key
 	prefixBits int
 	batchSize  int
 
@@ -147,8 +147,7 @@ func NewKeyStore(d ds.Batching, opts ...KeyStoreOption) (KeyStore, error) {
 	}
 
 	ks := &keyStore{
-		ds:         d,
-		base:       ds.NewKey(cfg.base),
+		ds:         namespace.Wrap(d, ds.NewKey(cfg.base)),
 		prefixBits: cfg.prefixBits,
 		batchSize:  cfg.batchSize,
 		requests:   make(chan operation),
@@ -183,7 +182,7 @@ func NewKeyStore(d ds.Batching, opts ...KeyStoreOption) (KeyStore, error) {
 // If the prefix is longer than `prefixBits`, only the first `prefixBits` bits
 // are used, allowing the returned key to serve as a query prefix for the
 // datastore.
-func dsKey[K kad.Key[K]](k K, prefixBits int, base ds.Key) ds.Key {
+func dsKey[K kad.Key[K]](k K, prefixBits int) ds.Key {
 	b := strings.Builder{}
 	l := k.BitLen()
 	for i := range min(prefixBits, l) {
@@ -193,7 +192,7 @@ func dsKey[K kad.Key[K]](k K, prefixBits int, base ds.Key) ds.Key {
 	if l == keyspace.KeyLen {
 		b.WriteString(base64.URLEncoding.EncodeToString(keyspace.KeyToBytes(k)[prefixBits/8:]))
 	}
-	return base.ChildString(b.String())
+	return ds.NewKey(b.String())
 }
 
 // decodeKey reconstructs a 256-bit binary key from a hierarchical datastore key string.
@@ -207,7 +206,6 @@ func dsKey[K kad.Key[K]](k K, prefixBits int, base ds.Key) ds.Key {
 //
 // Returns the reconstructed 256-bit key or an error if base64URL decoding fails.
 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)
 	for i := range s.prefixBits {
@@ -256,7 +254,7 @@ func (s *keyStore) worker() {
 				op.response <- operationResponse{err: err}
 
 			case opEmpty:
-				err := empty(op.ctx, s.ds, s.base, s.batchSize)
+				err := empty(op.ctx, s.ds, s.batchSize)
 				op.response <- operationResponse{err: err}
 
 			case opSize:
@@ -282,7 +280,7 @@ func (s *keyStore) put(ctx context.Context, keys []mh.Multihash) ([]mh.Multihash
 			continue
 		}
 		seen[k] = struct{}{}
-		dsk := dsKey(k, s.prefixBits, s.base)
+		dsk := dsKey(k, s.prefixBits)
 		ok, err := s.ds.Has(ctx, dsk)
 		if err != nil {
 			return nil, err
@@ -320,7 +318,7 @@ func (s *keyStore) put(ctx context.Context, keys []mh.Multihash) ([]mh.Multihash
 // get returns all keys whose bit256 representation matches the provided
 // prefix.
 func (s *keyStore) get(ctx context.Context, prefix bitstr.Key) ([]mh.Multihash, error) {
-	dsk := dsKey(prefix, s.prefixBits, s.base).String()
+	dsk := dsKey(prefix, s.prefixBits).String()
 	res, err := s.ds.Query(ctx, query.Query{Prefix: dsk})
 	if err != nil {
 		return nil, err
@@ -353,7 +351,7 @@ func (s *keyStore) get(ctx context.Context, prefix bitstr.Key) ([]mh.Multihash,
 // multihash whose kademlia identifier (bit256.Key) starts with the provided
 // bit-prefix.
 func (s *keyStore) containsPrefix(ctx context.Context, prefix bitstr.Key) (bool, error) {
-	dsk := dsKey(prefix, s.prefixBits, s.base).String()
+	dsk := dsKey(prefix, s.prefixBits).String()
 	longPrefix := prefix.BitLen() > s.prefixBits
 	q := query.Query{Prefix: dsk, KeysOnly: true}
 	if !longPrefix {
@@ -386,8 +384,8 @@ func (s *keyStore) containsPrefix(ctx context.Context, prefix bitstr.Key) (bool,
 
 // empty deletes all entries under the datastore prefix, assuming s.lk is
 // already held.
-func empty(ctx context.Context, d ds.Batching, base ds.Key, batchSize int) error {
-	res, err := d.Query(ctx, query.Query{Prefix: base.String(), KeysOnly: true})
+func empty(ctx context.Context, d ds.Batching, batchSize int) error {
+	res, err := d.Query(ctx, query.Query{KeysOnly: true})
 	if err != nil {
 		return err
 	}
@@ -431,7 +429,7 @@ func (s *keyStore) delete(ctx context.Context, keys []mh.Multihash) error {
 		return err
 	}
 	for _, h := range keys {
-		dsk := dsKey(keyspace.MhToBit256(h), s.prefixBits, s.base)
+		dsk := dsKey(keyspace.MhToBit256(h), s.prefixBits)
 		err := b.Delete(ctx, dsk)
 		if err != nil {
 			return err
@@ -442,7 +440,7 @@ func (s *keyStore) delete(ctx context.Context, keys []mh.Multihash) error {
 
 // size returns the number of keys currently stored in the KeyStore.
 func (s *keyStore) size(ctx context.Context) (size int, err error) {
-	q := query.Query{Prefix: s.base.String(), KeysOnly: true}
+	q := query.Query{KeysOnly: true}
 	res, err := s.ds.Query(ctx, q)
 	if err != nil {
 		return

provider/datastore/keystore_test.go

@@ -259,13 +259,12 @@ func testKeyStoreSizeImpl(t *testing.T, store KeyStore) {
 
 func TestDsKey(t *testing.T) {
 	s := keyStore{
-		base:       ds.NewKey("/base/prefix"),
 		prefixBits: 8,
 	}
 
 	k := bit256.ZeroKey()
-	dsk := dsKey(k, s.prefixBits, s.base)
-	expectedPrefix := "/base/prefix/0/0/0/0/0/0/0/0/"
+	dsk := dsKey(k, s.prefixBits)
+	expectedPrefix := "/0/0/0/0/0/0/0/0/"
 	require.Equal(t, expectedPrefix, dsk.String()[:len(expectedPrefix)])
 
 	s.prefixBits = 16
@@ -276,8 +275,8 @@ func TestDsKey(t *testing.T) {
 		require.NoError(t, err)
 		k := bit256.NewKey(b[:])
 
-		sdk := dsKey(k, s.prefixBits, s.base)
-		require.Equal(t, strings.Count(s.base.String(), "/")+s.prefixBits+1, strings.Count(sdk.String(), "/"))
+		sdk := dsKey(k, s.prefixBits)
+		require.Equal(t, s.prefixBits+1, strings.Count(sdk.String(), "/"))
 		decoded, err := s.decodeKey(sdk.String())
 		require.NoError(t, err)
 		require.Equal(t, k, decoded)

provider/datastore/resettable_keystore.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/ipfs/go-cid"
 	ds "github.com/ipfs/go-datastore"
+	"github.com/ipfs/go-datastore/namespace"
 	"github.com/libp2p/go-libp2p-kad-dht/provider/internal/keyspace"
 	mh "github.com/multiformats/go-multihash"
 )
@@ -26,17 +27,35 @@ type resetOp struct {
 }
 
 // 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.
+// operations using a dual-datastore architecture. It maintains two separate
+// datastores (primary and alternate) where only one is active at any time,
+// enabling atomic replacement of all stored keys without interrupting
+// concurrent operations.
 //
-// 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.
+// Architecture:
+//   - Primary datastore: Currently active storage for all read/write operations
+//   - Alternate datastore: Standby storage used during reset operations
+//   - The datastores use "/0" and "/1" namespace suffixes and can be swapped
+//
+// Reset Operation Flow:
+//  1. New keys from reset are written to the alternate (inactive) datastore
+//  2. Concurrent Put operations are automatically duplicated to both datastores
+//     to maintain consistency during the transition
+//  3. Once all reset keys are written, the datastores are atomically swapped
+//  4. The old datastore (now alternate) is cleaned up
+//
+// Thread Safety:
+//   - All operations are processed sequentially by a single worker goroutine
+//   - Reset operations are non-blocking for concurrent reads and writes
+//   - Only one reset operation can be active at a time
+//
+// The reset operation allows complete replacement of stored multihashes
+// without data loss or service interruption, making it suitable for
+// scenarios requiring periodic full dataset updates.
 type ResettableKeyStore struct {
 	keyStore
 
-	altBase         ds.Key
+	altDs           ds.Batching
 	resetInProgress bool
 	resetSync       chan []mh.Multihash // passes keys from worker to reset go routine
 	resetOps        chan resetOp        // reset operations that must be run in main go routine
@@ -59,15 +78,14 @@ func NewResettableKeyStore(d ds.Batching, opts ...KeyStoreOption) (*ResettableKe
 
 	rks := &ResettableKeyStore{
 		keyStore: keyStore{
-			ds:         d,
-			base:       ds.NewKey(cfg.base).ChildString("0"),
+			ds:         namespace.Wrap(d, ds.NewKey(cfg.base+"/0")),
 			prefixBits: cfg.prefixBits,
 			batchSize:  cfg.batchSize,
 			requests:   make(chan operation),
 			close:      make(chan struct{}),
 			done:       make(chan struct{}),
 		},
-		altBase:   ds.NewKey(cfg.base).ChildString("1"),
+		altDs:     namespace.Wrap(d, ds.NewKey(cfg.base+"/1")),
 		resetOps:  make(chan resetOp),
 		resetSync: make(chan []mh.Multihash, 128), // buffered to avoid blocking
 	}
@@ -105,7 +123,7 @@ func (s *ResettableKeyStore) worker() {
 				op.response <- operationResponse{err: err}
 
 			case opEmpty:
-				err := empty(op.ctx, s.ds, s.base, s.batchSize)
+				err := empty(op.ctx, s.ds, s.batchSize)
 				op.response <- operationResponse{err: err}
 
 			case opSize:
@@ -127,20 +145,21 @@ func (s *ResettableKeyStore) worker() {
 // handling during reset operations.
 func (s *ResettableKeyStore) put(ctx context.Context, keys []mh.Multihash) ([]mh.Multihash, error) {
 	if s.resetInProgress {
-		// Reset is in progress, write to alternate base in addition to current base
+		// Reset is in progress, write to alternate datastore in addition to
+		// current datastore
 		s.resetSync <- keys
 	}
 	return s.keyStore.put(ctx, keys)
 }
 
-// altPut writes the given multihashes to the alternate base in the datastore.
+// altPut writes the given multihashes to the alternate datastore.
 func (s *ResettableKeyStore) altPut(ctx context.Context, keys []mh.Multihash) error {
-	b, err := s.ds.Batch(ctx)
+	b, err := s.altDs.Batch(ctx)
 	if err != nil {
 		return err
 	}
 	for _, h := range keys {
-		dsk := dsKey(keyspace.MhToBit256(h), s.prefixBits, s.altBase)
+		dsk := dsKey(keyspace.MhToBit256(h), s.prefixBits)
 		if err := b.Put(ctx, dsk, h); err != nil {
 			return err
 		}
@@ -155,7 +174,7 @@ func (s *ResettableKeyStore) handleResetOp(op resetOp) {
 			op.response <- ErrResetInProgress
 			return
 		}
-		if err := empty(context.Background(), s.ds, s.altBase, s.batchSize); err != nil {
+		if err := empty(context.Background(), s.altDs, s.batchSize); err != nil {
 			op.response <- err
 			return
 		}
@@ -166,10 +185,10 @@ func (s *ResettableKeyStore) handleResetOp(op resetOp) {
 
 	// Cleanup operation
 	if op.success {
-		// Swap the keystore prefix bases.
-		oldBase := s.base
-		s.base = s.altBase
-		s.altBase = oldBase
+		// Swap the active datastore.
+		oldDs := s.ds
+		s.ds = s.altDs
+		s.altDs = oldDs
 	}
 	// Drain resetSync
 drain:
@@ -180,9 +199,9 @@ drain:
 			break drain
 		}
 	}
-	// Empty the unused base prefix
+	// Empty the unused datastore.
 	s.resetInProgress = false
-	op.response <- empty(context.Background(), s.ds, s.altBase, s.batchSize)
+	op.response <- empty(context.Background(), s.altDs, s.batchSize)
 }
 
 // ResetCids atomically replaces all stored keys with the CIDs received from
@@ -231,7 +250,7 @@ func (s *ResettableKeyStore) ResetCids(ctx context.Context, keysChan <-chan cid.
 		case <-s.done:
 			// Safe not to go through the worker since we are done, and we need to
 			// cleanup
-			empty(context.Background(), s.ds, s.altBase, s.batchSize)
+			empty(context.Background(), s.altDs, s.batchSize)
 		}
 	}()
 
@@ -260,7 +279,7 @@ func (s *ResettableKeyStore) ResetCids(ctx context.Context, keysChan <-chan cid.
 		return nil
 	}
 
-	// Read all the keys from the channel and put them in batch at the alternate base
+	// Read all the keys from the channel and write them to the altDs
 loop:
 	for {
 		select {