Enable parallel fetching of property sources in AwsS3EnvironmentRepository

This change introduces an optimization to fetch configuration files from S3 concurrently using CompletableFuture. It includes a new 'pool-size' property to configure the level of parallelism, while maintaining the original precedence order and providing backward-compatible sequential execution by default.

Fixes gh-3222

Signed-off-by: Noah Hanka <noah.hanka@veeva.com>

Author: noah-hanka-veeva

spring-cloud-config-server/src/main/java/org/springframework/cloud/config/server/environment/AwsS3EnvironmentProperties.java

@@ -21,6 +21,7 @@
 
 /**
  * @author Clay McCoy
+ * @author Noah Hanka
  */
 @ConfigurationProperties("spring.cloud.config.server.awss3")
 public class AwsS3EnvironmentProperties implements EnvironmentRepositoryProperties {
@@ -41,11 +42,17 @@ public class AwsS3EnvironmentProperties implements EnvironmentRepositoryProperti
 	private String bucket;
 
 	/**
-	 * Use application name as intermediate directory. Analogous to `searchPaths:
-	 * {application}` from Git backend.
+	 * Use application name as intermediate directory. Analogous to \`searchPaths:
+	 * {application}\` from Git backend.
 	 */
 	private boolean useDirectoryLayout;
 
+	/**
+	 * Thread pool size for fetching properties from S3 in parallel. If set to 0 or less,
+	 * fetching will be sequential.
+	 */
+	private int poolSize = 0;
+
 	private int order = DEFAULT_ORDER;
 
 	public String getRegion() {
@@ -80,6 +87,14 @@ public void setUseDirectoryLayout(boolean useDirectoryLayout) {
 		this.useDirectoryLayout = useDirectoryLayout;
 	}
 
+	public int getPoolSize() {
+		return poolSize;
+	}
+
+	public void setPoolSize(int poolSize) {
+		this.poolSize = poolSize;
+	}
+
 	public int getOrder() {
 		return order;
 	}

spring-cloud-config-server/src/main/java/org/springframework/cloud/config/server/environment/AwsS3EnvironmentRepository.java

@@ -22,7 +22,9 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
-import java.util.function.Consumer;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.stream.Collectors;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -45,6 +47,7 @@
 
 /**
  * @author Clay McCoy
+ * @author Noah Hanka
  * @author Scott Frederick
  * @author Daniel Aiken
  */
@@ -64,18 +67,26 @@ public class AwsS3EnvironmentRepository implements EnvironmentRepository, Ordere
 
 	private final boolean useApplicationAsDirectory;
 
+	private final Executor executor;
+
 	protected int order = Ordered.LOWEST_PRECEDENCE;
 
 	public AwsS3EnvironmentRepository(S3Client s3Client, String bucketName, ConfigServerProperties server) {
-		this(s3Client, bucketName, false, server);
+		this(s3Client, bucketName, false, server, null);
 	}
 
 	public AwsS3EnvironmentRepository(S3Client s3Client, String bucketName, boolean useApplicationAsDirectory,
 			ConfigServerProperties server) {
+		this(s3Client, bucketName, useApplicationAsDirectory, server, null);
+	}
+
+	public AwsS3EnvironmentRepository(S3Client s3Client, String bucketName, boolean useApplicationAsDirectory,
+			ConfigServerProperties server, Executor executor) {
 		this.s3Client = s3Client;
 		this.bucketName = bucketName;
 		this.serverProperties = server;
 		this.useApplicationAsDirectory = useApplicationAsDirectory;
+		this.executor = (executor != null) ? executor : Runnable::run;
 	}
 
 	@Override
@@ -126,77 +137,70 @@ public Environment findOne(String specifiedApplication, String specifiedProfiles
 
 	private void addPropertySources(Environment environment, List<String> apps, String[] profiles,
 			List<String> labels) {
+		List<S3ConfigFile> allConfigs = new ArrayList<>();
 		for (String label : labels) {
 			// If we have profiles, add property sources with those profiles
 			for (String profile : profiles) {
-				addPropertySourcesForApps(apps,
-						app -> addProfileSpecificPropertySource(environment, app, profile, label));
+				apps.forEach(app -> allConfigs.addAll(getProfileSpecificS3ConfigFiles(app, profile, label)));
 			}
 		}
 
 		// If we have no profiles just add property sources for all apps
 		if (profiles.length == 0) {
 			for (String label : labels) {
-				addPropertySourcesForApps(apps,
-						app -> addNonProfileSpecificPropertySource(environment, app, null, label));
+				apps.forEach(app -> allConfigs.addAll(getNonProfileSpecificS3ConfigFiles(app, null, label)));
 			}
 		}
 		else {
 			for (String label : labels) {
 				// If we have profiles, we still need to add property sources from files
-				// that
-				// are not profile specific but we pass
-				// along the profiles as well so we can check if any non-profile specific
-				// YAML
-				// files have profile specific documents
-				// within them
+				// that are not profile specific but we pass along the profiles as well
+				// so we can check if any non-profile specific YAML files have profile
+				// specific documents within them
 				for (String profile : profiles) {
-					addPropertySourcesForApps(apps,
-							app -> addNonProfileSpecificPropertySource(environment, app, profile, label));
+					apps.forEach(app -> allConfigs.addAll(getNonProfileSpecificS3ConfigFiles(app, profile, label)));
 				}
 			}
 		}
-	}
 
-	private void addPropertySourcesForApps(List<String> apps, Consumer<String> addPropertySource) {
-		apps.forEach(addPropertySource);
+		List<CompletableFuture<Void>> futures = allConfigs.stream()
+			.map(config -> CompletableFuture.runAsync(config::read, executor))
+			.collect(Collectors.toList());
+		CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
+
+		for (S3ConfigFile s3ConfigFile : allConfigs) {
+			addPropertySource(environment, s3ConfigFile);
+		}
 	}
 
-	private void addProfileSpecificPropertySource(Environment environment, String app, String profile, String label) {
-		List<S3ConfigFile> s3ConfigFiles = getS3ConfigFile(app, profile, label, this::getS3PropertiesOrJsonConfigFile,
+	private List<S3ConfigFile> getProfileSpecificS3ConfigFiles(String app, String profile, String label) {
+		return getS3ConfigFile(app, profile, label, this::getS3PropertiesOrJsonConfigFile,
 				this::getProfileSpecificS3ConfigFileYaml);
-		addPropertySource(environment, s3ConfigFiles);
 	}
 
-	private void addNonProfileSpecificPropertySource(Environment environment, String app, String profile,
-			String label) {
-		List<S3ConfigFile> s3ConfigFiles = getS3ConfigFile(app, profile, label,
-				this::getNonProfileSpecificPropertiesOrJsonConfigFile, this::getNonProfileSpecificS3ConfigFileYaml);
-		addPropertySource(environment, s3ConfigFiles);
-	}
-
-	private void addPropertySource(Environment environment, List<S3ConfigFile> s3ConfigFiles) {
-		for (S3ConfigFile s3ConfigFile : s3ConfigFiles) {
-			final Properties config = s3ConfigFile.read();
-			// This logic handles the case where the s3 file is a YAML file that is
-			// not profile specific (ie it does not have -<profile> in the name)
-			// and does not have any profile specific documents in it. In this case we do
-			// not want to include this
-			// property source we only want to include the document for the default
-			// profile. When we create
-			// the S3ConfigFile for this file we set the
-			// shouldIncludeWithEmptyProperties to false
-			// in ProfileSpecificYamlDocumentS3ConfigFile for this specific case.
-			if (config != null) {
-				if (!config.isEmpty() || s3ConfigFile.isShouldIncludeWithEmptyProperties()) {
-					environment.setVersion(s3ConfigFile.getVersion());
-					config.putAll(serverProperties.getOverrides());
-					PropertySource propertySource = new PropertySource(s3ConfigFile.getName(), config);
-					if (LOG.isDebugEnabled()) {
-						LOG.debug("Adding property source to environment " + propertySource);
-					}
-					environment.add(propertySource);
+	private List<S3ConfigFile> getNonProfileSpecificS3ConfigFiles(String app, String profile, String label) {
+		return getS3ConfigFile(app, profile, label, this::getNonProfileSpecificPropertiesOrJsonConfigFile,
+				this::getNonProfileSpecificS3ConfigFileYaml);
+	}
+
+	private void addPropertySource(Environment environment, S3ConfigFile s3ConfigFile) {
+		final Properties config = s3ConfigFile.read();
+		// This logic handles the case where the s3 file is a YAML file that is
+		// not profile specific (ie it does not have -<profile> in the name)
+		// and does not have any profile specific documents in it. In this case we do
+		// not want to include this property source we only want to include the
+		// document for the default profile. When we create the S3ConfigFile for
+		// this file we set the shouldIncludeWithEmptyProperties to false
+		// in ProfileSpecificYamlDocumentS3ConfigFile for this specific case.
+		if (config != null) {
+			if (!config.isEmpty() || s3ConfigFile.isShouldIncludeWithEmptyProperties()) {
+				environment.setVersion(s3ConfigFile.getVersion());
+				config.putAll(serverProperties.getOverrides());
+				PropertySource propertySource = new PropertySource(s3ConfigFile.getName(), config);
+				if (LOG.isDebugEnabled()) {
+					LOG.debug("Adding property source to environment " + propertySource);
 				}
+				environment.add(propertySource);
 			}
 		}
 	}
@@ -348,8 +352,8 @@ protected S3ConfigFile(String application, String profile, String label, String
 		this.profile = profile;
 		this.label = label;
 		this.bucketName = bucketName;
-		this.s3Client = s3Client;
 		this.useApplicationAsDirectory = useApplicationAsDirectory;
+		this.s3Client = s3Client;
 	}
 
 	String getVersion() {
@@ -437,11 +441,10 @@ class PropertyS3ConfigFile extends S3ConfigFile {
 	PropertyS3ConfigFile(String application, String profile, String label, String bucketName,
 			boolean useApplicationAsDirectory, S3Client s3Client) {
 		super(application, profile, label, bucketName, useApplicationAsDirectory, s3Client);
-		this.properties = read();
 	}
 
 	@Override
-	public Properties read() {
+	public synchronized Properties read() {
 		if (this.properties != null) {
 			return this.properties;
 		}
@@ -453,6 +456,7 @@ public Properties read() {
 			LOG.warn("Exception thrown when reading property file", e);
 			throw new IllegalStateException("Cannot load environment", e);
 		}
+		this.properties = props;
 		return props;
 	}
 
@@ -478,8 +482,6 @@ class YamlS3ConfigFile extends S3ConfigFile {
 			final YamlProcessor.DocumentMatcher... documentMatchers) {
 		super(application, profile, label, bucketName, useApplicationAsDirectory, s3Client);
 		this.documentMatchers = documentMatchers;
-		this.properties = read();
-
 	}
 
 	protected static boolean profileMatchesActivateProperty(String profile, Properties properties) {
@@ -493,15 +495,16 @@ protected static boolean onProfilePropertyExists(Properties properties) {
 	}
 
 	@Override
-	public Properties read() {
+	public synchronized Properties read() {
 		if (properties != null) {
 			return properties;
 		}
 		final YamlPropertiesFactoryBean yaml = new YamlPropertiesFactoryBean();
 		try (InputStream in = getObject()) {
 			yaml.setResources(new InputStreamResource(in));
 			yaml.setDocumentMatchers(documentMatchers);
-			return yaml.getObject();
+			this.properties = yaml.getObject();
+			return this.properties;
 		}
 		catch (Exception e) {
 			LOG.warn("Could not read YAML file", e);
@@ -567,7 +570,6 @@ class JsonS3ConfigFile extends YamlS3ConfigFile {
 	JsonS3ConfigFile(String application, String profile, String label, String bucketName,
 			boolean useApplicationAsDirectory, S3Client s3Client) {
 		super(application, profile, label, bucketName, useApplicationAsDirectory, s3Client);
-		this.properties = read();
 	}
 
 	@Override

spring-cloud-config-server/src/main/java/org/springframework/cloud/config/server/environment/AwsS3EnvironmentRepositoryFactory.java

@@ -16,6 +16,9 @@
 
 package org.springframework.cloud.config.server.environment;
 
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
 
@@ -38,8 +41,13 @@ public AwsS3EnvironmentRepository build(AwsS3EnvironmentProperties environmentPr
 		configureClientBuilder(clientBuilder, environmentProperties.getRegion(), environmentProperties.getEndpoint());
 		final S3Client client = clientBuilder.build();
 
+		Executor executor = null;
+		if (environmentProperties.getPoolSize() > 0) {
+			executor = Executors.newFixedThreadPool(environmentProperties.getPoolSize());
+		}
+
 		AwsS3EnvironmentRepository repository = new AwsS3EnvironmentRepository(client,
-				environmentProperties.getBucket(), environmentProperties.isUseDirectoryLayout(), server);
+				environmentProperties.getBucket(), environmentProperties.isUseDirectoryLayout(), server, executor);
 		repository.setOrder(environmentProperties.getOrder());
 		return repository;
 	}

spring-cloud-config-server/src/test/java/org/springframework/cloud/config/server/environment/ParallelAwsS3EnvironmentRepositoryTests.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2013-present the original author or authors.
+ *
+ * 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
+ *
+ *      https://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 org.springframework.cloud.config.server.environment;
+
+import java.io.ByteArrayInputStream;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.core.ResponseInputStream;
+import software.amazon.awssdk.http.AbortableInputStream;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.GetObjectRequest;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+
+import org.springframework.cloud.config.environment.Environment;
+import org.springframework.cloud.config.server.config.ConfigServerProperties;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * @author Noah Hanka
+ */
+public class ParallelAwsS3EnvironmentRepositoryTests {
+
+	@Test
+	public void testParallelFetching() {
+		S3Client s3Client = mock(S3Client.class);
+		ConfigServerProperties server = new ConfigServerProperties();
+		Executor executor = Executors.newFixedThreadPool(2);
+
+		AwsS3EnvironmentRepository repo = new AwsS3EnvironmentRepository(s3Client, "bucket", false, server, executor);
+
+		String content = "foo: bar";
+		GetObjectResponse response = GetObjectResponse.builder().build();
+
+		when(s3Client.getObject(any(GetObjectRequest.class))).thenAnswer(invocation -> {
+			return new ResponseInputStream<>(response,
+					AbortableInputStream.create(new ByteArrayInputStream(content.getBytes())));
+		});
+
+		// Request with 2 profiles
+		Environment env = repo.findOne("app", "p1,p2", "label");
+
+		assertThat(env.getPropertySources()).isNotEmpty();
+		verify(s3Client, atLeastOnce()).getObject(any(GetObjectRequest.class));
+	}
+
+	@Test
+	public void testSequentialFetchingByDefault() {
+		S3Client s3Client = mock(S3Client.class);
+		ConfigServerProperties server = new ConfigServerProperties();
+
+		AwsS3EnvironmentRepository repo = new AwsS3EnvironmentRepository(s3Client, "bucket", server);
+
+		String content = "foo: bar";
+		GetObjectResponse response = GetObjectResponse.builder().build();
+
+		when(s3Client.getObject(any(GetObjectRequest.class))).thenAnswer(invocation -> {
+			return new ResponseInputStream<>(response,
+					AbortableInputStream.create(new ByteArrayInputStream(content.getBytes())));
+		});
+
+		Environment env = repo.findOne("app", "p1", "label");
+
+		assertThat(env.getPropertySources()).isNotEmpty();
+		verify(s3Client, atLeastOnce()).getObject(any(GetObjectRequest.class));
+	}
+
+}