Cache AWS Current Region Result

Calling Regions.getCurrentRegion() attempts to access the EC2 instance metadata endpoint, but doing so repeatedly can result in being rate-limited. When that happens, null is returned and other operations that access the metadata endpoint such as instance profile credential refreshes can fail. This updates calls to Regions.getCurrentRegion() through a new class that caches the first successful region lookup and fails loudly when current region resolution fails instead of returning null.
trinodb · Mar 6, 2020 · 4374841 · 4374841
1 parent 773f77e
commit 4374841
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 20 deletions.
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/aws/AwsCurrentRegionHolder.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/aws/AwsCurrentRegionHolder.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.hive.aws;
+
+import com.amazonaws.regions.Region;
+import com.amazonaws.regions.Regions;
+import com.google.common.base.Suppliers;
+
+import java.util.function.Supplier;
+
+/**
+ * Caches the result of calling {@link Regions#getCurrentRegion()} since accessing EC2 instance
+ * metadata repeatedly can result in being throttled and prevent other metadata accessing operations
+ * such as refreshing instance credentials from working normally
+ */
+public final class AwsCurrentRegionHolder
+{
+    private static final Supplier<Region> SUPPLIER = Suppliers.memoize(AwsCurrentRegionHolder::loadCurrentRegionOrThrowOnNull);
+
+    private AwsCurrentRegionHolder() {}
+
+    /**
+     * Attempts to resolve the current region from EC2's instance metadata through {@link Regions#getCurrentRegion()}. If
+     * no region is able to be resolved an exception is thrown
+     */
+    public static Region getCurrentRegionFromEC2Metadata()
+            throws IllegalStateException
+    {
+        return SUPPLIER.get();
+    }
+
+    /**
+     * @throws IllegalStateException when no region is resolved to avoid memoizing a transient failure
+     */
+    private static Region loadCurrentRegionOrThrowOnNull()
+            throws IllegalStateException
+    {
+        Region result = Regions.getCurrentRegion();
+        if (result == null) {
+            throw new IllegalStateException("Failed to resolve current AWS region from EC2 metadata");
+        }
+        return result;
+    }
+}
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java
@@ -21,8 +21,6 @@
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.auth.InstanceProfileCredentialsProvider;
 import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
-import com.amazonaws.regions.Region;
-import com.amazonaws.regions.Regions;
 import com.amazonaws.services.glue.AWSGlueAsync;
 import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder;
 import com.amazonaws.services.glue.model.AlreadyExistsException;
@@ -119,6 +117,7 @@
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static io.prestosql.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR;
 import static io.prestosql.plugin.hive.HiveErrorCode.HIVE_PARTITION_DROPPED_DURING_QUERY;
+import static io.prestosql.plugin.hive.aws.AwsCurrentRegionHolder.getCurrentRegionFromEC2Metadata;
 import static io.prestosql.plugin.hive.metastore.MetastoreUtil.makePartitionName;
 import static io.prestosql.plugin.hive.metastore.MetastoreUtil.verifyCanDropColumn;
 import static io.prestosql.plugin.hive.metastore.glue.GlueExpressionUtil.buildGlueExpression;
@@ -187,10 +186,7 @@ private static AWSGlueAsync createAsyncGlueClient(GlueHiveMetastoreConfig config
             asyncGlueClientBuilder.setRegion(config.getGlueRegion().get());
         }
         else if (config.getPinGlueClientToCurrentRegion()) {
-            Region currentRegion = Regions.getCurrentRegion();
-            if (currentRegion != null) {
-                asyncGlueClientBuilder.setRegion(currentRegion.getName());
-            }
+            asyncGlueClientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName());
         }
 
         asyncGlueClientBuilder.setCredentials(getAwsCredentialsProvider(config));

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/s3/PrestoS3FileSystem.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/s3/PrestoS3FileSystem.java
@@ -32,8 +32,6 @@
 import com.amazonaws.event.ProgressEventType;
 import com.amazonaws.event.ProgressListener;
 import com.amazonaws.metrics.RequestMetricCollector;
-import com.amazonaws.regions.Region;
-import com.amazonaws.regions.Regions;
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.AmazonS3Builder;
 import com.amazonaws.services.s3.AmazonS3Client;
@@ -116,6 +114,7 @@
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.Iterables.toArray;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
+import static io.prestosql.plugin.hive.aws.AwsCurrentRegionHolder.getCurrentRegionFromEC2Metadata;
 import static io.prestosql.plugin.hive.util.RetryDriver.retry;
 import static java.lang.Math.max;
 import static java.lang.Math.toIntExact;
@@ -740,11 +739,8 @@ private AmazonS3 createAmazonS3Client(Configuration hadoopConfig, ClientConfigur
 
         // use local region when running inside of EC2
         if (pinS3ClientToCurrentRegion) {
-            Region region = Regions.getCurrentRegion();
-            if (region != null) {
-                clientBuilder.setRegion(region.getName());
-                regionOrEndpointSet = true;
-            }
+            clientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName());
+            regionOrEndpointSet = true;
         }
 
         String endpoint = hadoopConfig.get(S3_ENDPOINT);

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/s3select/PrestoS3ClientFactory.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/s3select/PrestoS3ClientFactory.java
@@ -22,8 +22,6 @@
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.auth.InstanceProfileCredentialsProvider;
 import com.amazonaws.metrics.RequestMetricCollector;
-import com.amazonaws.regions.Region;
-import com.amazonaws.regions.Regions;
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.AmazonS3Builder;
 import com.amazonaws.services.s3.AmazonS3Client;
@@ -44,6 +42,7 @@
 import static com.amazonaws.regions.Regions.US_EAST_1;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.base.Verify.verify;
+import static io.prestosql.plugin.hive.aws.AwsCurrentRegionHolder.getCurrentRegionFromEC2Metadata;
 import static io.prestosql.plugin.hive.s3.PrestoS3FileSystem.S3_ACCESS_KEY;
 import static io.prestosql.plugin.hive.s3.PrestoS3FileSystem.S3_CONNECT_TIMEOUT;
 import static io.prestosql.plugin.hive.s3.PrestoS3FileSystem.S3_CREDENTIALS_PROVIDER;
@@ -127,11 +126,8 @@ private AmazonS3 createS3Client(Configuration config)
 
         // use local region when running inside of EC2
         if (pinS3ClientToCurrentRegion) {
-            Region region = Regions.getCurrentRegion();
-            if (region != null) {
-                clientBuilder.withRegion(region.getName());
-                regionOrEndpointSet = true;
-            }
+            clientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName());
+            regionOrEndpointSet = true;
         }
 
         if (!isNullOrEmpty(endpoint)) {