Repository: libp2p/go-libp2p-kad-dht

Issue:
provider: only reprovide skipped regions after `keystore.Reset()`
https://github.com/libp2p/go-libp2p-kad-dht/blob/0d37d0bb82136cf5d56e7e82a2ae9d5356fa89b6/provider/provider.go#L1992-L2004

If using a `ResettableKeystore`, we should only reprovide the regions that were reprovided since the node started. The regions that have been reprovided in a previous run, but less than `ReprovideInterval` ago don't need to be reprovided.



Discussion:

Solution patch:
diff --git a/provider/keystore/keystore.go b/provider/keystore/keystore.go
index 1c9401a..f93a909 100644
--- a/provider/keystore/keystore.go
+++ b/provider/keystore/keystore.go
@@ -219,7 +219,7 @@ func (s *keystore) worker() {
 				}
 
 			case opEmpty:
-				err := empty(op.ctx, s.ds, s.batchSize)
+				err := s.empty(op.ctx, s.ds)
 				op.response <- operationResponse{err: err}
 				if err == nil {
 					s.size = 0
@@ -298,15 +298,18 @@ func (s *keystore) put(ctx context.Context, keys []mh.Multihash) ([]mh.Multihash
 		}
 		if !ok {
 			if err := b.Put(ctx, dsk, h); err != nil {
-				return nil, err
+				return nil, fmt.Errorf("cannot put to batch: %w", err)
 			}
 			newKeys = append(newKeys, h)
 		}
 	}
 	if err := b.Commit(ctx); err != nil {
-		return nil, err
+		return nil, fmt.Errorf("cannot commit keystore updates: %w", err)
 	}
 	s.size += len(newKeys)
+	if err = s.ds.Sync(ctx, ds.NewKey("")); err != nil {
+		s.logger.Warnf("keystore: cannot sync datastore after put: %v", err)
+	}
 	return newKeys, nil
 }
 
@@ -367,8 +370,8 @@ func (s *keystore) containsPrefix(ctx context.Context, prefix bitstr.Key) (bool,
 	return false, nil
 }
 
-// empty deletes all entries under the datastore prefix
-func empty(ctx context.Context, d ds.Batching, batchSize int) error {
+// empty deletes all entries in the supplied datastore.
+func (s *keystore) empty(ctx context.Context, d ds.Batching) error {
 	batch, err := d.Batch(ctx)
 	if err != nil {
 		return err
@@ -379,7 +382,7 @@ func empty(ctx context.Context, d ds.Batching, batchSize int) error {
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
-		if writeCount >= batchSize {
+		if writeCount >= s.batchSize {
 			writeCount = 0
 			if err = batch.Commit(ctx); err != nil {
 				return fmt.Errorf("cannot commit keystore updates: %w", err)
@@ -404,7 +407,7 @@ func empty(ctx context.Context, d ds.Batching, batchSize int) error {
 		}
 	}
 	if err = d.Sync(ctx, ds.NewKey("")); err != nil {
-		return fmt.Errorf("cannot sync datastore: %w", err)
+		s.logger.Warnf("keystore: cannot sync datastore after put: %v", err)
 	}
 	return nil
 }
@@ -435,8 +438,14 @@ func (s *keystore) delete(ctx context.Context, keys []mh.Multihash) error {
 			removedCount++
 		}
 	}
+	if err := b.Commit(ctx); err != nil {
+		return fmt.Errorf("cannot commit keystore updates: %w", err)
+	}
 	s.size -= removedCount
-	return b.Commit(ctx)
+	if err = s.ds.Sync(ctx, ds.NewKey("")); err != nil {
+		s.logger.Warnf("keystore: cannot sync datastore after delete: %v", err)
+	}
+	return nil
 }
 
 // refreshSize iterates over all keys in the supplied datastore to count the
@@ -543,11 +552,15 @@ func (s *keystore) Close() error {
 	select {
 	case <-s.close:
 		// Already closed
-		return nil
 	default:
 		close(s.close)
-		<-s.done              // Wait for worker to exit
-		err = s.persistSize() // Safe to access s.size after worker exits
+		<-s.done // Wait for worker to exit
+		if err = s.persistSize(); err != nil {
+			return fmt.Errorf("error persisting size on close: %w", err)
+		}
+		if err = s.ds.Sync(context.Background(), sizeKey); err != nil {
+			return fmt.Errorf("error syncing size on close: %w", err)
+		}
 	}
 	return err
 }
diff --git a/provider/keystore/resettable_keystore.go b/provider/keystore/resettable_keystore.go
index 89f236c..643d9ef 100644
--- a/provider/keystore/resettable_keystore.go
+++ b/provider/keystore/resettable_keystore.go
@@ -3,6 +3,7 @@ package keystore
 import (
 	"context"
 	"errors"
+	"fmt"
 	"sync/atomic"
 
 	"github.com/ipfs/go-cid"
@@ -22,6 +23,9 @@ const (
 	opCleanup
 )
 
+// activeNamespaceKey is the key used to persist which namespace (0 or 1) is currently active
+var activeNamespaceKey = ds.NewKey("active")
+
 type resetOp struct {
 	op       opType
 	success  bool
@@ -38,13 +42,15 @@ type resetOp struct {
 //   - 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
+//   - Active namespace is persisted to survive restarts
 //
 // 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
+//  4. The active namespace marker is updated to persist the swap
+//  5. The old datastore (now alternate) is cleaned up
 //
 // Thread Safety:
 //   - All operations are processed sequentially by a single worker goroutine
@@ -57,9 +63,11 @@ type resetOp struct {
 type ResettableKeystore struct {
 	keystore
 
+	baseDs          ds.Batching // Unwrapped datastore for storing active namespace marker
 	altDs           ds.Batching
 	altSize         atomic.Int64
 	resetInProgress bool
+	activeNamespace byte
 	resetOps        chan resetOp // reset operations that must be run in main go routine
 }
 
@@ -69,15 +77,49 @@ var _ Keystore = (*ResettableKeystore)(nil)
 // provided datastore. It automatically adds "/0" and "/1" suffixes to the
 // configured datastore path to create two alternate storage locations for
 // atomic reset operations.
+//
+// On initialization, it checks for a persisted active namespace marker to
+// determine which namespace was active during the previous session. This
+// ensures keystore data persists correctly across restarts.
 func NewResettableKeystore(d ds.Batching, opts ...Option) (*ResettableKeystore, error) {
 	cfg, err := getOpts(opts)
 	if err != nil {
 		return nil, err
 	}
 
+	// Create namespaced datastores
+	baseDs := namespace.Wrap(d, ds.NewKey(cfg.path))
+	datastores := []ds.Batching{
+		namespace.Wrap(d, ds.NewKey(cfg.path+"/0")),
+		namespace.Wrap(d, ds.NewKey(cfg.path+"/1")),
+	}
+
+	// Check if there's a persisted active namespace marker
+	ctx := context.Background()
+	activeDs, err := baseDs.Get(ctx, activeNamespaceKey)
+	var primaryDs, altDs ds.Batching
+	var activeIdx byte
+	if err != nil {
+		if err != ds.ErrNotFound {
+			return nil, err
+		}
+		// No marker found, default to namespace 0 as primary
+		primaryDs = datastores[0]
+		altDs = datastores[1]
+		activeIdx = 0
+	} else {
+		// Marker found, use it to determine active namespace
+		if len(activeDs) != 1 {
+			return nil, fmt.Errorf("invalid active namespace marker length: %d", len(activeDs))
+		}
+		activeIdx = activeDs[0]
+		primaryDs = datastores[activeIdx]
+		altDs = datastores[1-activeIdx]
+	}
+
 	rks := &ResettableKeystore{
 		keystore: keystore{
-			ds:         namespace.Wrap(d, ds.NewKey(cfg.path+"/0")),
+			ds:         primaryDs,
 			prefixBits: cfg.prefixBits,
 			batchSize:  cfg.batchSize,
 			requests:   make(chan operation),
@@ -85,8 +127,10 @@ func NewResettableKeystore(d ds.Batching, opts ...Option) (*ResettableKeystore,
 			done:       make(chan struct{}),
 			logger:     log.Logger(cfg.loggerName),
 		},
-		altDs:    namespace.Wrap(d, ds.NewKey(cfg.path+"/1")),
-		resetOps: make(chan resetOp),
+		baseDs:          baseDs,
+		altDs:           altDs,
+		activeNamespace: activeIdx,
+		resetOps:        make(chan resetOp),
 	}
 
 	// start worker goroutine
@@ -137,7 +181,7 @@ func (s *ResettableKeystore) worker() {
 				}
 
 			case opEmpty:
-				err := empty(op.ctx, s.ds, s.batchSize)
+				err := s.empty(op.ctx, s.ds)
 				op.response <- operationResponse{err: err}
 				if err == nil {
 					s.size = 0
@@ -204,18 +248,25 @@ func (s *ResettableKeystore) altPut(ctx context.Context, keys []mh.Multihash) er
 		}
 	}
 	s.altSize.Add(added)
-	return b.Commit(ctx)
+	if err := b.Commit(ctx); err != nil {
+		return fmt.Errorf("cannot commit keystore updates: %w", err)
+	}
+	if err = s.ds.Sync(ctx, ds.NewKey("")); err != nil {
+		s.logger.Warnf("keystore: failed to sync datastore after alt put: %v", err)
+	}
+	return nil
 }
 
 // handleResetOp processes reset operations that need to happen synchronously.
 func (s *ResettableKeystore) handleResetOp(op resetOp) {
+	ctx := context.Background()
 	if op.op == opStart {
 		if s.resetInProgress {
 			op.response <- ErrResetInProgress
 			return
 		}
 		s.altSize.Store(0)
-		if err := empty(context.Background(), s.altDs, s.batchSize); err != nil {
+		if err := s.empty(ctx, s.altDs); err != nil {
 			op.response <- err
 			return
 		}
@@ -231,10 +282,24 @@ func (s *ResettableKeystore) handleResetOp(op resetOp) {
 		s.ds = s.altDs
 		s.altDs = oldDs
 		s.size = int(s.altSize.Load())
+
+		// Toggle the active namespace index
+		s.activeNamespace = 1 - s.activeNamespace
+		// Persist the new active namespace
+		activeValue := []byte{s.activeNamespace}
+
+		// Write the active namespace marker
+		if err := s.baseDs.Put(ctx, activeNamespaceKey, activeValue); err != nil {
+			s.logger.Errorf("keystore: failed to persist active namespace marker: %v", err)
+		}
+		// Sync to ensure marker is persisted
+		if err := s.baseDs.Sync(ctx, activeNamespaceKey); err != nil {
+			s.logger.Warnf("keystore: failed to sync active namespace marker: %v", err)
+		}
 	}
 	// Empty the unused datastore.
 	s.resetInProgress = false
-	op.response <- empty(context.Background(), s.altDs, s.batchSize)
+	op.response <- s.empty(ctx, s.altDs)
 }
 
 // ResetCids atomically replaces all stored keys with the CIDs received from
@@ -281,9 +346,10 @@ func (s *ResettableKeystore) ResetCids(ctx context.Context, keysChan <-chan cid.
 		case s.resetOps <- resetOp{op: opCleanup, success: success, response: opsChan}:
 			<-opsChan
 		case <-s.done:
-			// Safe not to go through the worker since we are done, and we need to
-			// cleanup
-			empty(context.Background(), s.altDs, s.batchSize)
+			// Worker is done, which means the keystore is closing or has closed.
+			// The underlying datastore may already be closed, so attempting to
+			// empty it could cause a panic. Skip cleanup since the datastore
+			// will be cleaned up during normal shutdown.
 		}
 	}()
 
diff --git a/provider/provider.go b/provider/provider.go
index fad9d35..e92e649 100644
--- a/provider/provider.go
+++ b/provider/provider.go
@@ -867,6 +867,9 @@ func (s *SweepingProvider) closestPeersToPrefix(prefix bitstr.Key) ([]peer.ID, e
 	i, noFreshPeersFoundCount := 0, 0
 	// Go down the trie to fully cover prefix.
 	for ; i < maxExplorationPrefixSearches; i++ {
+		if s.closed() {
+			return nil, ErrClosed
+		}
 		if !s.connectivity.IsOnline() {
 			return nil, errors.New("node is offline")
 		}
@@ -991,6 +994,9 @@ func (s *SweepingProvider) sendProviderRecords(keysAllocations map[peer.ID][][]m
 			pmes := genProvideMessage(addrInfo)
 			defer wg.Done()
 			for job := range jobChan {
+				if s.closed() {
+					return
+				}
 				err := s.provideKeysToPeer(job.pid, job.batches, pmes)
 				if err != nil {
 					errCount.Add(1)
@@ -1012,8 +1018,13 @@ func (s *SweepingProvider) sendProviderRecords(keysAllocations map[peer.ID][][]m
 		}()
 	}
 
+loop:
 	for p, batches := range keysAllocations {
-		jobChan <- provideJob{p, batches}
+		select {
+		case jobChan <- provideJob{p, batches}:
+		case <-s.done:
+			break loop
+		}
 	}
 	close(jobChan)
 	wg.Wait()
@@ -1381,6 +1392,11 @@ func (s *SweepingProvider) enqueueExpiredRegionsNoLock(recentlyReprovided *trie.
 // persistSuccessfulReprovide logs a successful reprovide to the datastore for
 // resumption after restarts.
 func (s *SweepingProvider) persistSuccessfulReprovide(prefix bitstr.Key) {
+	// Don't persist if the provider is shutting down. The datastore may already
+	// be closed, which would cause a panic.
+	if s.closed() {
+		return
+	}
 	now := time.Now()
 	k := datastore.NewKey(path.Join(reprovideHistoryKeyPrefix, formatTimestampHex(now), string(prefix)))
 	if err := s.datastore.Put(s.ctx, k, []byte{}); err != nil {
@@ -1392,6 +1408,11 @@ func (s *SweepingProvider) persistSuccessfulReprovide(prefix bitstr.Key) {
 // loadRecentlyReprovidedRegions loads all regions that have been successfully
 // reprovided within the last reprovide interval from the datastore.
 func (s *SweepingProvider) loadRecentlyReprovidedRegions(now time.Time) (*trie.Trie[bitstr.Key, struct{}], error) {
+	// Don't load if the provider is shutting down. The datastore may already
+	// be closed, which would cause a panic.
+	if s.closed() {
+		return trie.New[bitstr.Key, struct{}](), nil
+	}
 	s.gcReprovideHistoryIfNeeded(now)
 
 	q := query.Query{
@@ -1424,6 +1445,11 @@ func (s *SweepingProvider) loadRecentlyReprovidedRegions(now time.Time) (*trie.T
 // gcReprovideHistoryIfNeeded removes reprovide log entries older than the
 // reprovide interval. GC runs at most once per reprovide interval.
 func (s *SweepingProvider) gcReprovideHistoryIfNeeded(now time.Time) {
+	// Don't run GC if the provider is shutting down. The datastore may already
+	// be closed, which would cause a panic.
+	if s.closed() {
+		return
+	}
 	lastGC := time.Unix(s.lastReprovideHistoryGC.Load(), 0)
 	if now.Sub(lastGC) < s.reprovideInterval {
 		// Only run GC once per reprovide interval.
@@ -1629,6 +1655,11 @@ func (s *SweepingProvider) batchProvide(prefix bitstr.Key, keys []mh.Multihash)
 }
 
 func (s *SweepingProvider) batchReprovide(prefix bitstr.Key) {
+	// Check if provider is shutting down before accessing keystore/datastore.
+	// The underlying datastores may be closed externally during shutdown.
+	if s.closed() {
+		return
+	}
 	addrInfo, ok := s.selfAddrInfo()
 	if !ok {
 		// Don't provide if the node doesn't have a valid address to include in the
@@ -2019,7 +2050,7 @@ func (s *SweepingProvider) RefreshSchedule() error {
 	var reprovideAll bool
 	// Insert prefixes into the schedule
 	s.scheduleLk.Lock()
-	if len(gaps) > 1 || gaps[0] != bitstr.Key("") {
+	if s.bootstrapped.Load() {
 		_, resettableKeystore := s.keystore.(*keystore.ResettableKeystore)
 		if resettableKeystore && len(toInsert) > s.schedule.Size() {
 			// Schedule size is about to double



Test patch:
diff --git a/provider/keystore/resettable_keystore_test.go b/provider/keystore/resettable_keystore_test.go
index a3b5c2d..c93bf03 100644
--- a/provider/keystore/resettable_keystore_test.go
+++ b/provider/keystore/resettable_keystore_test.go
@@ -279,3 +279,221 @@ func TestKeystoreResetSizePersistence(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, numKeys, size2, "size should persist correctly after restart (without reset)")
 }
+
+func TestKeystoreActiveNamespacePersistenceX(t *testing.T) {
+	ds := ds.NewMapDatastore()
+	defer ds.Close()
+
+	ctx := context.Background()
+
+	// Create initial keystore
+	store, err := NewResettableKeystore(ds)
+	require.NoError(t, err)
+
+	// Add initial keys
+	const initialKeys = 50
+	initialMhs := random.Multihashes(initialKeys)
+	_, err = store.Put(ctx, initialMhs...)
+	require.NoError(t, err)
+
+	// Verify initial size
+	size, err := store.Size(ctx)
+	require.NoError(t, err)
+	require.Equal(t, initialKeys, size, "initial size should be %d", initialKeys)
+
+	// Perform reset with different keys
+	const resetKeys = 75
+	resetMhs := random.Multihashes(resetKeys)
+	resetChan := make(chan cid.Cid, resetKeys)
+	for _, h := range resetMhs {
+		resetChan <- cid.NewCidV1(cid.Raw, h)
+	}
+	close(resetChan)
+
+	err = store.ResetCids(ctx, resetChan)
+	require.NoError(t, err)
+
+	// Verify size after reset
+	size, err = store.Size(ctx)
+	require.NoError(t, err)
+	require.Equal(t, resetKeys, size, "size after reset should be %d", resetKeys)
+
+	// Close the keystore
+	err = store.Close()
+	require.NoError(t, err)
+
+	// Reopen keystore - it should restore from the active namespace (post-reset)
+	store2, err := NewResettableKeystore(ds)
+	require.NoError(t, err)
+	defer store2.Close()
+
+	// Verify size is correctly restored from the active datastore
+	size2, err := store2.Size(ctx)
+	require.NoError(t, err)
+	require.Equal(t, resetKeys, size2, "size should persist correctly after restart (with reset)")
+
+	// Verify that the reset keys are still present
+	const prefixBits = 6
+	for _, h := range resetMhs {
+		prefix := bitstr.Key(key.BitString(keyspace.MhToBit256(h))[:prefixBits])
+		got, err := store2.Get(ctx, prefix)
+		require.NoError(t, err)
+		found := false
+		for _, m := range got {
+			if string(m) == string(h) {
+				found = true
+				break
+			}
+		}
+		require.True(t, found, "reset key %v should be present after restart", h)
+	}
+
+	// Verify that the initial keys (before reset) are NOT present
+	for _, h := range initialMhs {
+		prefix := bitstr.Key(key.BitString(keyspace.MhToBit256(h))[:prefixBits])
+		got, err := store2.Get(ctx, prefix)
+		require.NoError(t, err)
+		for _, m := range got {
+			require.NotEqual(t, string(m), string(h), "initial key %v should not be present after restart", h)
+		}
+	}
+}
+
+func TestKeystoreActiveNamespacePersistenceMultipleResets(t *testing.T) {
+	ds := ds.NewMapDatastore()
+	defer ds.Close()
+
+	ctx := context.Background()
+
+	// Create initial keystore
+	store, err := NewResettableKeystore(ds)
+	require.NoError(t, err)
+
+	// Perform first reset
+	const firstResetKeys = 100
+	firstResetMhs := random.Multihashes(firstResetKeys)
+	resetChan1 := make(chan cid.Cid, firstResetKeys)
+	for _, h := range firstResetMhs {
+		resetChan1 <- cid.NewCidV1(cid.Raw, h)
+	}
+	close(resetChan1)
+
+	err = store.ResetCids(ctx, resetChan1)
+	require.NoError(t, err)
+
+	// Close and reopen
+	err = store.Close()
+	require.NoError(t, err)
+
+	store2, err := NewResettableKeystore(ds)
+	require.NoError(t, err)
+
+	// Verify first reset keys are present
+	size, err := store2.Size(ctx)
+	require.NoError(t, err)
+	require.Equal(t, firstResetKeys, size, "size after first reset and restart should be %d", firstResetKeys)
+
+	// Perform second reset
+	const secondResetKeys = 50
+	secondResetMhs := random.Multihashes(secondResetKeys)
+	resetChan2 := make(chan cid.Cid, secondResetKeys)
+	for _, h := range secondResetMhs {
+		resetChan2 <- cid.NewCidV1(cid.Raw, h)
+	}
+	close(resetChan2)
+
+	err = store2.ResetCids(ctx, resetChan2)
+	require.NoError(t, err)
+
+	// Close and reopen again
+	err = store2.Close()
+	require.NoError(t, err)
+
+	store3, err := NewResettableKeystore(ds)
+	require.NoError(t, err)
+	defer store3.Close()
+
+	// Verify second reset keys are present and first reset keys are gone
+	size, err = store3.Size(ctx)
+	require.NoError(t, err)
+	require.Equal(t, secondResetKeys, size, "size after second reset and restart should be %d", secondResetKeys)
+
+	// Verify that the second reset keys are present
+	const prefixBits = 6
+	for _, h := range secondResetMhs {
+		prefix := bitstr.Key(key.BitString(keyspace.MhToBit256(h))[:prefixBits])
+		got, err := store3.Get(ctx, prefix)
+		require.NoError(t, err)
+		found := false
+		for _, m := range got {
+			if string(m) == string(h) {
+				found = true
+				break
+			}
+		}
+		require.True(t, found, "second reset key %v should be present after restart", h)
+	}
+
+	// Verify that the first reset keys are NOT present
+	for _, h := range firstResetMhs {
+		prefix := bitstr.Key(key.BitString(keyspace.MhToBit256(h))[:prefixBits])
+		got, err := store3.Get(ctx, prefix)
+		require.NoError(t, err)
+		for _, m := range got {
+			require.NotEqual(t, string(m), string(h), "first reset key %v should not be present after second reset and restart", h)
+		}
+	}
+}
+
+// TestKeystoreCloseDuringReset tests that closing the keystore during a
+// ResetCids operation does not cause a panic. This reproduces the race
+// condition where the underlying datastore is closed while ResetCids is
+// still running.
+func TestKeystoreCloseDuringReset(t *testing.T) {
+	ds := ds.NewMapDatastore()
+	defer ds.Close()
+
+	store, err := NewResettableKeystore(ds)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+
+	// Create a slow channel that will feed keys gradually
+	const resetKeys = 1000
+	resetChan := make(chan cid.Cid)
+
+	// Start reset in background
+	resetDone := make(chan error, 1)
+	go func() {
+		resetDone <- store.ResetCids(ctx, resetChan)
+	}()
+
+	// Feed some keys, then close the keystore
+	go func() {
+		resetMhs := random.Multihashes(resetKeys)
+		for i, h := range resetMhs {
+			resetChan <- cid.NewCidV1(cid.Raw, h)
+			// Close the keystore partway through
+			if i == 100 {
+				// Give it a moment to process some keys
+				go func() {
+					err := store.Close()
+					// Close might return an error if operations are still running,
+					// but it shouldn't panic
+					if err != nil {
+						t.Logf("Close returned error (expected): %v", err)
+					}
+				}()
+			}
+		}
+		close(resetChan)
+	}()
+
+	// Wait for reset to complete or fail
+	err = <-resetDone
+	// ResetCids should return ErrClosed when the keystore is closed
+	if err != nil && err != ErrClosed {
+		t.Logf("ResetCids returned error (may be expected during close): %v", err)
+	}
+	// The important thing is we didn't panic
+}