keyring: update handle to state inside replication loop

.changelog/15227.txt

@@ -0,0 +1,3 @@
+```release-note:bug
+keyring: Fixed a bug where replication would stop after snapshot restores
+```

nomad/encrypter.go

@@ -18,10 +18,11 @@ import (
 	log "github.com/hashicorp/go-hclog"
 	kms "github.com/hashicorp/go-kms-wrapping/v2"
 	"github.com/hashicorp/go-kms-wrapping/v2/aead"
-	"golang.org/x/time/rate"
+	memdb "github.com/hashicorp/go-memdb"
 
 	"github.com/hashicorp/nomad/helper"
 	"github.com/hashicorp/nomad/helper/crypto"
+	"github.com/hashicorp/nomad/nomad/state"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -434,35 +435,42 @@ func (krr *KeyringReplicator) stop() {
 const keyringReplicationRate = 10
 
 func (krr *KeyringReplicator) run(ctx context.Context) {
-	limiter := rate.NewLimiter(keyringReplicationRate, keyringReplicationRate)
 	krr.logger.Debug("starting encryption key replication")
 	defer krr.logger.Debug("exiting key replication")
 
-	retryErrTimer, stop := helper.NewSafeTimer(time.Second * 1)
-	defer stop()
-
-START:
 	store := krr.srv.fsm.State()
+	abandonCh := store.AbandonCh()
+	minIndex := uint64(0)
 
 	for {
 		select {
 		case <-krr.srv.shutdownCtx.Done():
 			return
 		case <-ctx.Done():
 			return
+		case <-abandonCh:
+			// reset the blocking query to make sure we get all the keys from a
+			// restored snapshot
+			store = krr.srv.fsm.State()
+			abandonCh = store.AbandonCh()
+			minIndex = uint64(0)
 		default:
-			// Rate limit how often we attempt replication
-			err := limiter.Wait(ctx)
+			// ensure that we abandon the blocking query periodically to give
+			// other contexts a chance to fire
+			queryCtx, queryCancel := context.WithTimeout(ctx, time.Second)
+			defer queryCancel()
+
+			var rawIter any
+			var err error
+			rawIter, minIndex, err = store.BlockingQuery(getRootKeyMetas, minIndex, queryCtx)
 			if err != nil {
-				goto ERR_WAIT // rate limit exceeded
-			}
-
-			ws := store.NewWatchSet()
-			iter, err := store.RootKeyMetas(ws)
-			if err != nil {
-				krr.logger.Error("failed to fetch keyring", "error", err)
-				goto ERR_WAIT
+				if queryCtx.Err() == nil {
+					// we get errors for closed context but don't want to log on that
+					krr.logger.Error("failed to fetch keyring", "error", err)
+				}
+				continue
 			}
+			iter := rawIter.(memdb.ResultIterator)
 
 			for {
 				raw := iter.Next()
@@ -488,16 +496,20 @@ START:
 		}
 	}
 
-ERR_WAIT:
-	retryErrTimer.Reset(1 * time.Second)
+}
 
-	select {
-	case <-retryErrTimer.C:
-		goto START
-	case <-ctx.Done():
-		return
+// getRootKeyMetas implements state.QueryFn and is run as a blocking query to
+// detect new key metadata
+func getRootKeyMetas(ws memdb.WatchSet, store *state.StateStore) (interface{}, uint64, error) {
+	iter, err := store.RootKeyMetas(ws)
+	if err != nil {
+		return nil, 0, err
 	}
-
+	index, err := store.Index(state.TableRootKeyMeta)
+	if err != nil {
+		return nil, 0, err
+	}
+	return iter, index, nil
 }
 
 // replicateKey replicates a single key from peer servers that was present in

nomad/encrypter_test.go

@@ -1,13 +1,15 @@
 package nomad
 
 import (
+	"bytes"
 	"context"
 	"os"
 	"path/filepath"
 	"testing"
 	"time"
 
 	msgpackrpc "github.com/hashicorp/net-rpc-msgpackrpc"
+	"github.com/shoenig/test/must"
 	"github.com/stretchr/testify/require"
 
 	"github.com/hashicorp/nomad/ci"
@@ -300,6 +302,32 @@ func TestEncrypter_KeyringReplication(t *testing.T) {
 		time.Second*5, time.Second,
 		"expected new servers to get replicated keys")
 
+	// Scenario: reload a snapshot
+
+	t.Logf("taking snapshot of node5")
+
+	snapshot, err := srv5.fsm.Snapshot()
+	must.NoError(t, err)
+
+	defer snapshot.Release()
+
+	// Persist so we can read it back
+	buf := bytes.NewBuffer(nil)
+	sink := &MockSink{buf, false}
+	must.NoError(t, snapshot.Persist(sink))
+
+	must.NoError(t, srv5.fsm.Restore(sink))
+
+	// rotate the key
+
+	err = msgpackrpc.CallWithCodec(codec, "Keyring.Rotate", rotateReq, &rotateResp)
+	require.NoError(t, err)
+	keyID4 := rotateResp.Key.KeyID
+
+	require.Eventually(t, checkReplicationFn(keyID4),
+		time.Second*5, time.Second,
+		"expected new servers to get replicated keys after snapshot restore")
+
 }
 
 func TestEncrypter_EncryptDecrypt(t *testing.T) {