Update YAML Rest tests to check for product header on all responses (#83290) (#83993)

jbaiera · web-flow · commit 485b82cbe256 · 2022-02-16T13:21:21.000-05:00
This PR adds assertions to YAML Rest tests to ensure that product headers are always 
returned in rest responses. Additional work has been included to fix a number of misuses
of ThreadContext, mostly because of stashing listeners without their accompanying contexts.

BWC Rest tests have been disabled for a few cases while the fixes are backported.
diff --git a/docs/changelog/83290.yaml b/docs/changelog/83290.yaml
@@ -0,0 +1,5 @@
+pr: 83290
+summary: Update YAML Rest tests to check for product header on all responses
+area: Infra/REST API
+type: enhancement
+issues: []
diff --git a/modules/reindex/src/main/java/org/elasticsearch/reindex/AbstractAsyncBulkByScrollAction.java b/modules/reindex/src/main/java/org/elasticsearch/reindex/AbstractAsyncBulkByScrollAction.java
@@ -593,7 +593,7 @@ protected void finishHim(Exception failure) {
      */
     protected void finishHim(Exception failure, List<Failure> indexingFailures, List<SearchFailure> searchFailures, boolean timedOut) {
         logger.debug("[{}]: finishing without any catastrophic failures", task.getId());
-        scrollSource.close(() -> {
+        scrollSource.close(threadPool.getThreadContext().preserveContext(() -> {
             if (failure == null) {
                 BulkByScrollResponse response = buildResponse(
                     timeValueNanos(System.nanoTime() - startTime.get()),
@@ -605,7 +605,7 @@ protected void finishHim(Exception failure, List<Failure> indexingFailures, List
             } else {
                 listener.onFailure(failure);
             }
-        });
+        }));
     }
 
     /**
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/cat.snapshots/10_basic.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/cat.snapshots/10_basic.yml
@@ -23,6 +23,10 @@
                $/
 ---
 "Test cat snapshots output":
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
+
   - do:
       snapshot.create_repository:
         repository: test_cat_snapshots_1
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.clone/10_basic.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.clone/10_basic.yml
@@ -1,5 +1,8 @@
 ---
 setup:
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
 
   - do:
       snapshot.create_repository:
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.create/10_basic.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.create/10_basic.yml
@@ -1,5 +1,8 @@
 ---
 setup:
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
 
   - do:
       snapshot.create_repository:
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.get/10_basic.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.get/10_basic.yml
@@ -1,5 +1,8 @@
 ---
 setup:
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
 
   - do:
       snapshot.create_repository:
@@ -61,6 +64,7 @@ setup:
 
 ---
 "Get snapshot info when verbose is false":
+
   - do:
       indices.create:
         index: test_index
@@ -198,7 +202,6 @@ setup:
   - skip:
       version: " - 7.12.99"
       reason: "Introduced in 7.13.0"
-
   - do:
       indices.create:
         index: test_index
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.get_repository/20_repository_uuid.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.get_repository/20_repository_uuid.yml
@@ -1,4 +1,9 @@
 ---
+setup:
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
+---
 "Get repository returns UUID":
   - skip:
       version: " - 7.12.99"
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.restore/10_basic.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.restore/10_basic.yml
@@ -1,5 +1,8 @@
 ---
 setup:
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
 
   - do:
       snapshot.create_repository:
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.status/10_basic.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/snapshot.status/10_basic.yml
@@ -1,5 +1,8 @@
 ---
 setup:
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
 
   - do:
       snapshot.create_repository:
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/tsdb/30_snapshot.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/tsdb/30_snapshot.yml
@@ -1,5 +1,9 @@
 ---
 setup:
+  - skip:
+      version: " - 8.1.99"
+      reason: "Pause BWC tests until #83290 is backported"
+
   - do:
       snapshot.create_repository:
         repository: test_repo
diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/RestoreClusterStateListener.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/RestoreClusterStateListener.java
@@ -16,10 +16,13 @@
 import org.elasticsearch.cluster.RestoreInProgress;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.snapshots.RestoreInfo;
 import org.elasticsearch.snapshots.RestoreService;
 
+import java.util.function.Supplier;
+
 import static org.elasticsearch.snapshots.RestoreService.restoreInProgress;
 
 public class RestoreClusterStateListener implements ClusterStateListener {
@@ -29,43 +32,48 @@ public class RestoreClusterStateListener implements ClusterStateListener {
     private final ClusterService clusterService;
     private final String uuid;
     private final ActionListener<RestoreSnapshotResponse> listener;
+    private final Supplier<ThreadContext.StoredContext> contextSupplier;
 
     private RestoreClusterStateListener(
         ClusterService clusterService,
         RestoreService.RestoreCompletionResponse response,
-        ActionListener<RestoreSnapshotResponse> listener
+        ActionListener<RestoreSnapshotResponse> listener,
+        Supplier<ThreadContext.StoredContext> contextSupplier
     ) {
         this.clusterService = clusterService;
         this.uuid = response.getUuid();
         this.listener = listener;
+        this.contextSupplier = contextSupplier;
     }
 
     @Override
     public void clusterChanged(ClusterChangedEvent changedEvent) {
-        final RestoreInProgress.Entry prevEntry = restoreInProgress(changedEvent.previousState(), uuid);
-        final RestoreInProgress.Entry newEntry = restoreInProgress(changedEvent.state(), uuid);
-        if (prevEntry == null) {
-            // When there is a master failure after a restore has been started, this listener might not be registered
-            // on the current master and as such it might miss some intermediary cluster states due to batching.
-            // Clean up listener in that case and acknowledge completion of restore operation to client.
-            clusterService.removeListener(this);
-            listener.onResponse(new RestoreSnapshotResponse((RestoreInfo) null));
-        } else if (newEntry == null) {
-            clusterService.removeListener(this);
-            ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards = prevEntry.shards();
-            assert prevEntry.state().completed() : "expected completed snapshot state but was " + prevEntry.state();
-            assert RestoreService.completed(shards) : "expected all restore entries to be completed";
-            RestoreInfo ri = new RestoreInfo(
-                prevEntry.snapshot().getSnapshotId().getName(),
-                prevEntry.indices(),
-                shards.size(),
-                shards.size() - RestoreService.failedShards(shards)
-            );
-            RestoreSnapshotResponse response = new RestoreSnapshotResponse(ri);
-            logger.debug("restore of [{}] completed", prevEntry.snapshot().getSnapshotId());
-            listener.onResponse(response);
-        } else {
-            // restore not completed yet, wait for next cluster state update
+        try (ThreadContext.StoredContext stored = contextSupplier.get()) {
+            final RestoreInProgress.Entry prevEntry = restoreInProgress(changedEvent.previousState(), uuid);
+            final RestoreInProgress.Entry newEntry = restoreInProgress(changedEvent.state(), uuid);
+            if (prevEntry == null) {
+                // When there is a master failure after a restore has been started, this listener might not be registered
+                // on the current master and as such it might miss some intermediary cluster states due to batching.
+                // Clean up listener in that case and acknowledge completion of restore operation to client.
+                clusterService.removeListener(this);
+                listener.onResponse(new RestoreSnapshotResponse((RestoreInfo) null));
+            } else if (newEntry == null) {
+                clusterService.removeListener(this);
+                ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards = prevEntry.shards();
+                assert prevEntry.state().completed() : "expected completed snapshot state but was " + prevEntry.state();
+                assert RestoreService.completed(shards) : "expected all restore entries to be completed";
+                RestoreInfo ri = new RestoreInfo(
+                    prevEntry.snapshot().getSnapshotId().getName(),
+                    prevEntry.indices(),
+                    shards.size(),
+                    shards.size() - RestoreService.failedShards(shards)
+                );
+                RestoreSnapshotResponse response = new RestoreSnapshotResponse(ri);
+                logger.debug("restore of [{}] completed", prevEntry.snapshot().getSnapshotId());
+                listener.onResponse(response);
+            } else {
+                // restore not completed yet, wait for next cluster state update
+            }
         }
     }
 
@@ -76,8 +84,11 @@ public void clusterChanged(ClusterChangedEvent changedEvent) {
     public static void createAndRegisterListener(
         ClusterService clusterService,
         RestoreService.RestoreCompletionResponse response,
-        ActionListener<RestoreSnapshotResponse> listener
+        ActionListener<RestoreSnapshotResponse> listener,
+        ThreadContext threadContext
     ) {
-        clusterService.addListener(new RestoreClusterStateListener(clusterService, response, listener));
+        clusterService.addListener(
+            new RestoreClusterStateListener(clusterService, response, listener, threadContext.newRestorableContext(true))
+        );
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/TransportRestoreSnapshotAction.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/TransportRestoreSnapshotAction.java
@@ -72,7 +72,12 @@ protected void masterOperation(
     ) {
         restoreService.restoreSnapshot(request, listener.delegateFailure((delegatedListener, restoreCompletionResponse) -> {
             if (restoreCompletionResponse.getRestoreInfo() == null && request.waitForCompletion()) {
-                RestoreClusterStateListener.createAndRegisterListener(clusterService, restoreCompletionResponse, delegatedListener);
+                RestoreClusterStateListener.createAndRegisterListener(
+                    clusterService,
+                    restoreCompletionResponse,
+                    delegatedListener,
+                    threadPool.getThreadContext()
+                );
             } else {
                 delegatedListener.onResponse(new RestoreSnapshotResponse(restoreCompletionResponse.getRestoreInfo()));
             }
diff --git a/server/src/main/java/org/elasticsearch/cluster/service/ClusterApplierService.java b/server/src/main/java/org/elasticsearch/cluster/service/ClusterApplierService.java
@@ -213,7 +213,7 @@ public void removeApplier(ClusterStateApplier applier) {
     }
 
     /**
-     * Add a listener for updated cluster states
+     * Add a listener for updated cluster states. Listeners are executed in the system thread context.
      */
     public void addListener(ClusterStateListener listener) {
         clusterStateListeners.add(listener);
@@ -222,7 +222,7 @@ public void addListener(ClusterStateListener listener) {
     /**
      * Removes a listener for updated cluster states.
      */
-    public void removeListener(ClusterStateListener listener) {
+    public void removeListener(final ClusterStateListener listener) {
         clusterStateListeners.remove(listener);
     }
 
diff --git a/server/src/main/java/org/elasticsearch/snapshots/SnapshotsService.java b/server/src/main/java/org/elasticsearch/snapshots/SnapshotsService.java
@@ -21,6 +21,7 @@
 import org.elasticsearch.action.admin.cluster.snapshots.create.CreateSnapshotRequest;
 import org.elasticsearch.action.admin.cluster.snapshots.delete.DeleteSnapshotRequest;
 import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.action.support.GroupedActionListener;
 import org.elasticsearch.action.support.master.TransportMasterNodeAction;
 import org.elasticsearch.cluster.ClusterChangedEvent;
@@ -2957,7 +2958,8 @@ static Map<String, DataStreamAlias> filterDataStreamAliases(
      * @param listener listener
      */
     private void addListener(Snapshot snapshot, ActionListener<Tuple<RepositoryData, SnapshotInfo>> listener) {
-        snapshotCompletionListeners.computeIfAbsent(snapshot, k -> new CopyOnWriteArrayList<>()).add(listener);
+        snapshotCompletionListeners.computeIfAbsent(snapshot, k -> new CopyOnWriteArrayList<>())
+            .add(ContextPreservingActionListener.wrapPreservingContext(listener, threadPool.getThreadContext()));
     }
 
     @Override
diff --git a/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/ClientYamlTestResponse.java b/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/ClientYamlTestResponse.java
@@ -87,13 +87,20 @@ public String getReasonPhrase() {
      * Get a list of all of the values of all warning headers returned in the response.
      */
     public List<String> getWarningHeaders() {
-        List<String> warningHeaders = new ArrayList<>();
+        return getHeaders("Warning");
+    }
+
+    /**
+     * Get a list of all the values of a given header returned in the response.
+     */
+    public List<String> getHeaders(String name) {
+        List<String> headers = new ArrayList<>();
         for (Header header : response.getHeaders()) {
-            if (header.getName().equals("Warning")) {
-                warningHeaders.add(header.getValue());
+            if (header.getName().equalsIgnoreCase(name)) {
+                headers.add(header.getValue());
             }
         }
-        return warningHeaders;
+        return headers;
     }
 
     /**
diff --git a/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java b/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java
@@ -367,6 +367,7 @@ public void execute(ClientYamlTestExecutionContext executionContext) throws IOEx
             final String testPath = executionContext.getClientYamlTestCandidate() != null
                 ? executionContext.getClientYamlTestCandidate().getTestPath()
                 : null;
+            checkElasticProductHeader(response.getHeaders("X-elastic-product"));
             checkWarningHeaders(response.getWarningHeaders(), testPath);
         } catch (ClientYamlTestResponseException e) {
             ClientYamlTestResponse restTestResponse = e.getRestTestResponse();
@@ -392,6 +393,31 @@ public void execute(ClientYamlTestExecutionContext executionContext) throws IOEx
         }
     }
 
+    void checkElasticProductHeader(final List<String> productHeaders) {
+        if (productHeaders.isEmpty()) {
+            fail("Response is missing required X-Elastic-Product response header");
+        }
+        boolean headerPresent = false;
+        final List<String> unexpected = new ArrayList<>();
+        for (String header : productHeaders) {
+            if (header.equals("Elasticsearch")) {
+                headerPresent = true;
+                break;
+            } else {
+                unexpected.add(header);
+            }
+        }
+        if (headerPresent == false) {
+            StringBuilder failureMessage = new StringBuilder();
+            appendBadHeaders(
+                failureMessage,
+                unexpected,
+                "did not get expected product header [Elasticsearch], found header" + (unexpected.size() > 1 ? "s" : "")
+            );
+            fail(failureMessage.toString());
+        }
+    }
+
     void checkWarningHeaders(final List<String> warningHeaders) {
         checkWarningHeaders(warningHeaders, null);
     }
diff --git a/test/framework/src/test/java/org/elasticsearch/test/rest/yaml/section/DoSectionTests.java b/test/framework/src/test/java/org/elasticsearch/test/rest/yaml/section/DoSectionTests.java
@@ -605,6 +605,7 @@ public void testNodeSelectorByVersion() throws IOException {
                 doSection.getApiCallSection().getNodeSelector()
             )
         ).thenReturn(mockResponse);
+        when(mockResponse.getHeaders("X-elastic-product")).thenReturn(List.of("Elasticsearch"));
         doSection.execute(context);
         verify(context).callApi(
             "indices.get_field_mapping",
diff --git a/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportPutFollowAction.java b/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportPutFollowAction.java
@@ -268,7 +268,8 @@ public void onFailure(Exception e) {
                     assert restoreInfo.failedShards() > 0 : "Should have failed shards";
                     delegatedListener.onResponse(new PutFollowAction.Response(true, false, false));
                 }
-            })
+            }),
+            threadPool.getThreadContext()
         );
     }
 
diff --git a/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportUnfollowAction.java b/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportUnfollowAction.java
@@ -16,6 +16,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.action.support.GroupedActionListener;
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.action.support.master.AcknowledgedTransportMasterNodeAction;
@@ -178,10 +179,16 @@ private void removeRetentionLeaseForShard(
             ) {
                 logger.trace("{} removing retention lease [{}] while unfollowing leader index", followerShardId, retentionLeaseId);
                 final ThreadContext threadContext = threadPool.getThreadContext();
+                // We're about to stash the thread context for this retention lease removal. The listener will be completed while the
+                // context is stashed. The context needs to be restored in the listener when it is completing or else it is simply wiped.
+                final ActionListener<ActionResponse.Empty> preservedListener = new ContextPreservingActionListener<>(
+                    threadContext.newRestorableContext(true),
+                    listener
+                );
                 try (ThreadContext.StoredContext ignore = threadPool.getThreadContext().stashContext()) {
                     // we have to execute under the system context so that if security is enabled the removal is authorized
                     threadContext.markAsSystemContext();
-                    CcrRetentionLeases.asyncRemoveRetentionLease(leaderShardId, retentionLeaseId, remoteClient, listener);
+                    CcrRetentionLeases.asyncRemoveRetentionLease(leaderShardId, retentionLeaseId, remoteClient, preservedListener);
                 }
             }
 
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/action/TransportXPackUsageAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/action/TransportXPackUsageAction.java
diff --git a/x-pack/plugin/watcher/src/main/java/org/elasticsearch/xpack/watcher/WatcherUsageTransportAction.java b/x-pack/plugin/watcher/src/main/java/org/elasticsearch/xpack/watcher/WatcherUsageTransportAction.java
