From a298895dd66a8df6ff834c5d646c38e0fc21e105 Mon Sep 17 00:00:00 2001
From: Nigel Brittain <108375408+nbaws@users.noreply.github.com>
Date: Fri, 9 Aug 2024 23:42:45 +1000
Subject: [PATCH] aws: fix test flake when no IMDS can be found (#35633)

Commit Message: aws: fix test flake when no IMDS can be found
Additional Description: Lack of IMDS (169.254.169.254 address) can cause
a race condition and crash during testing due to the shutdown of cluster
manager. This scenario should not occur normally, as cluster manager
will still exist and the lack of IMDS is handled gracefully.

Signed-off-by: Nigel Brittain <nbaws@amazon.com>
---
 .../extensions/common/aws/metadata_fetcher.cc |  6 +++++
 .../common/aws/metadata_fetcher_test.cc       | 26 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/source/extensions/common/aws/metadata_fetcher.cc b/source/extensions/common/aws/metadata_fetcher.cc
index f5dbe85eea55..aeaac64661d2 100644
--- a/source/extensions/common/aws/metadata_fetcher.cc
+++ b/source/extensions/common/aws/metadata_fetcher.cc
@@ -52,6 +52,12 @@ class MetadataFetcherImpl : public MetadataFetcher,
     ASSERT(!request_);
     complete_ = false;
     receiver_ = makeOptRef(receiver);
+
+    // Stop processing if we are shutting down
+    if (cm_.isShutdown()) {
+      return;
+    }
+
     const auto thread_local_cluster = cm_.getThreadLocalCluster(cluster_name_);
     if (thread_local_cluster == nullptr) {
       ENVOY_LOG(error, "{} AWS Metadata failed: [cluster = {}] not found", __func__, cluster_name_);
diff --git a/test/extensions/common/aws/metadata_fetcher_test.cc b/test/extensions/common/aws/metadata_fetcher_test.cc
index da4702211d00..32356aa883e1 100644
--- a/test/extensions/common/aws/metadata_fetcher_test.cc
+++ b/test/extensions/common/aws/metadata_fetcher_test.cc
@@ -77,6 +77,20 @@ MATCHER_P(OptionsHasRetryPolicy, policyMatcher, "") {
 class MetadataFetcherTest : public testing::Test {
 public:
   void setupFetcher() {
+
+    mock_factory_ctx_.server_factory_context_.cluster_manager_.initializeThreadLocalClusters(
+        {"cluster_name"});
+    fetcher_ = MetadataFetcher::create(mock_factory_ctx_.server_factory_context_.cluster_manager_,
+                                       "cluster_name");
+    EXPECT_TRUE(fetcher_ != nullptr);
+  }
+
+  void setupFetcherShutDown() {
+    ON_CALL(mock_factory_ctx_.server_factory_context_.cluster_manager_, getThreadLocalCluster(_))
+        .WillByDefault(Return(nullptr));
+    ON_CALL(mock_factory_ctx_.server_factory_context_.cluster_manager_, isShutdown())
+        .WillByDefault(Return(true));
+
     mock_factory_ctx_.server_factory_context_.cluster_manager_.initializeThreadLocalClusters(
         {"cluster_name"});
     fetcher_ = MetadataFetcher::create(mock_factory_ctx_.server_factory_context_.cluster_manager_,
@@ -103,6 +117,18 @@ TEST_F(MetadataFetcherTest, TestGetSuccess) {
   fetcher_->fetch(message, parent_span_, receiver);
 }
 
+TEST_F(MetadataFetcherTest, TestClusterShutdown) {
+  // Setup
+  setupFetcherShutDown();
+  Http::RequestMessageImpl message;
+  MockMetadataReceiver receiver;
+  EXPECT_CALL(receiver, onMetadataSuccess(_)).Times(0);
+  EXPECT_CALL(receiver, onMetadataError(_)).Times(0);
+
+  // Act
+  fetcher_->fetch(message, parent_span_, receiver);
+}
+
 TEST_F(MetadataFetcherTest, TestRequestMatchAndSpanPassedDown) {
   // Setup
   setupFetcher();