From eed42487dedb189ba5626a7b3ad66b1580752686 Mon Sep 17 00:00:00 2001
From: Michael Schurter <mschurter@hashicorp.com>
Date: Thu, 28 Sep 2023 11:56:11 -0700
Subject: [PATCH] client: prevent watching stale alloc state

When waiting on a previous alloc we must query against the leader before
switching to a stale query with index set.

Also check to ensure the response is fresh before using it like #18269
---
 client/allocwatcher/alloc_watcher.go | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/client/allocwatcher/alloc_watcher.go b/client/allocwatcher/alloc_watcher.go
index c85ec1c9f93b..a4866c11ec8a 100644
--- a/client/allocwatcher/alloc_watcher.go
+++ b/client/allocwatcher/alloc_watcher.go
@@ -353,9 +353,11 @@ func (p *remotePrevAlloc) Wait(ctx context.Context) error {
 	req := structs.AllocSpecificRequest{
 		AllocID: p.prevAllocID,
 		QueryOptions: structs.QueryOptions{
-			Region:     p.config.Region,
-			AllowStale: true,
-			AuthToken:  p.config.Node.SecretID,
+			Region:    p.config.Region,
+			AuthToken: p.config.Node.SecretID,
+
+			// Initially get response from leader, then switch to stale
+			AllowStale: false,
 		},
 	}
 
@@ -381,6 +383,22 @@ func (p *remotePrevAlloc) Wait(ctx context.Context) error {
 				return ctx.Err()
 			}
 		}
+
+		if req.AllowStale && resp.Index <= req.MinQueryIndex {
+			retry := getRemoteRetryIntv + helper.RandomStagger(getRemoteRetryIntv)
+			p.logger.Warn("received stale alloc; retrying",
+				"req_index", req.MinQueryIndex,
+				"resp_index", resp.Index,
+				"wait", retry,
+			)
+			select {
+			case <-time.After(retry):
+				continue
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+		}
+
 		if resp.Alloc == nil {
 			p.logger.Debug("blocking alloc was GC'd")
 			return nil
@@ -392,6 +410,7 @@ func (p *remotePrevAlloc) Wait(ctx context.Context) error {
 
 		// Update the query index and requery.
 		if resp.Index > req.MinQueryIndex {
+			req.AllowStale = true
 			req.MinQueryIndex = resp.Index
 		}
 	}