Added abstraction for defining environment variables

When launching workers as separate tasks, it can be useful to be able to
use additional logic to define environment variables.  This commit
provides an abstraction to allow for the customization of environment
variables on a per worker basis as well as two useful implementations:

* A no-op implementation (returns an empty Map).
* An implementation that moves the current environment variable handling
* out of the `DeployerPartitionHandler`.

Resolves spring-cloud/spring-cloud-task#181

spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandler.java

@@ -16,11 +16,8 @@
 package org.springframework.cloud.task.batch.partition;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -36,18 +33,17 @@ import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.core.explore.JobExplorer;
 import org.springframework.batch.core.partition.PartitionHandler;
 import org.springframework.batch.core.partition.StepExecutionSplitter;
+import org.springframework.batch.item.ExecutionContext;
 import org.springframework.batch.poller.DirectPoller;
 import org.springframework.batch.poller.Poller;
+import org.springframework.beans.factory.InitializingBean;
 import org.springframework.cloud.deployer.spi.core.AppDefinition;
 import org.springframework.cloud.deployer.spi.core.AppDeploymentRequest;
 import org.springframework.cloud.deployer.spi.task.TaskLauncher;
 import org.springframework.cloud.task.listener.annotation.BeforeTask;
 import org.springframework.cloud.task.repository.TaskExecution;
 import org.springframework.context.EnvironmentAware;
-import org.springframework.core.env.AbstractEnvironment;
 import org.springframework.core.env.Environment;
-import org.springframework.core.env.MapPropertySource;
-import org.springframework.core.env.PropertySource;
 import org.springframework.core.io.Resource;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -68,7 +64,7 @@ import org.springframework.util.CollectionUtils;
  *
  * @author Michael Minella
  */
-public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAware {
+public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAware, InitializingBean {
 
 	public static final String SPRING_CLOUD_TASK_JOB_EXECUTION_ID =
 			"spring.cloud.task.job-execution-id";
@@ -93,8 +89,6 @@ public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAw
 
 	private Resource resource;
 
-	private Map<String, String> environmentProperties = new HashMap<>();
-
 	private String stepName;
 
 	private Log logger = LogFactory.getLog(DeployerPartitionHandler.class);
@@ -107,6 +101,8 @@ public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAw
 
 	private Map<String, String> deploymentProperties;
 
+	private EnvironmentVariablesProvider environmentVariablesProvider;
+
 	public DeployerPartitionHandler(TaskLauncher taskLauncher,
 			JobExplorer jobExplorer,
 			Resource resource,
@@ -122,6 +118,10 @@ public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAw
 		this.stepName = stepName;
 	}
 
+	public void setEnvironmentVariablesProvider(EnvironmentVariablesProvider environmentVariablesProvider) {
+		this.environmentVariablesProvider = environmentVariablesProvider;
+	}
+
 	/**
 	 * The maximum number of workers to be executing at once.
 	 *
@@ -143,15 +143,6 @@ public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAw
 		this.gridSize = gridSize;
 	}
 
-	/**
-	 * System properties to be made available for all workers.
-	 *
-	 * @param environmentProperties Map of properties
-	 */
-	public void setEnvironmentProperties(Map<String, String> environmentProperties) {
-		this.environmentProperties = environmentProperties;
-	}
-
 	/**
 	 * The interval to check the job repository for completed steps.
 	 *
@@ -173,7 +164,7 @@ public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAw
 	/**
 	 * Map of deployment properties to be used by the {@link TaskLauncher}
 	 *
-	 * @param deploymentProperties
+	 * @param deploymentProperties properites to be used by the {@link TaskLauncher}
 	 */
 	public void setDeploymentProperties(Map<String, String> deploymentProperties) {
 		this.deploymentProperties = deploymentProperties;
@@ -232,16 +223,14 @@ public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAw
 				String.valueOf(workerStepExecution.getId())));
 		arguments.add(formatArgument(SPRING_CLOUD_TASK_STEP_NAME, this.stepName));
 
-		Map<String, String> environmentProperties = new HashMap<>(this.environmentProperties.size());
-		environmentProperties.putAll(getCurrentEnvironmentProperties());
-		environmentProperties.putAll(this.environmentProperties);
+		ExecutionContext copyContext = new ExecutionContext(workerStepExecution.getExecutionContext());
 
 		AppDefinition definition =
 				new AppDefinition(String.format("%s:%s:%s",
 						taskExecution.getTaskName(),
 						workerStepExecution.getJobExecution().getJobInstance().getJobName(),
 						workerStepExecution.getStepName()),
-						environmentProperties);
+						this.environmentVariablesProvider.getEnvironmentVariables(copyContext));
 
 		AppDeploymentRequest request =
 				new AppDeploymentRequest(definition,
@@ -312,38 +301,17 @@ public class DeployerPartitionHandler implements PartitionHandler, EnvironmentAw
 		return status.equals(BatchStatus.COMPLETED) || status.isGreaterThan(BatchStatus.STARTED);
 	}
 
-	private Map<String, String> getArguments(List<String> arguments) {
-		Map<String, String> argumentMap = new HashMap<>(arguments.size());
-
-		for (String argument : arguments) {
-			String[] pieces = argument.split("=");
-			argumentMap.put(pieces[0], pieces[1]);
-		}
-
-		return argumentMap;
-	}
-
 	@Override
 	public void setEnvironment(Environment environment) {
 		this.environment = environment;
 	}
 
-	private Map<String, String> getCurrentEnvironmentProperties() {
-		Map<String, String> currentEnvironment = new HashMap<>();
+	@Override
+	public void afterPropertiesSet() throws Exception {
+		if(this.environmentVariablesProvider == null) {
+			this.environmentVariablesProvider =
+					new SimpleEnvironmentVariablesProvider(this.environment);
 
-		Set<String> keys = new HashSet<>();
-
-		for (Iterator it = ((AbstractEnvironment) this.environment).getPropertySources().iterator(); it.hasNext(); ) {
-			PropertySource propertySource = (PropertySource) it.next();
-			if (propertySource instanceof MapPropertySource) {
-				keys.addAll(Arrays.asList(((MapPropertySource) propertySource).getPropertyNames()));
-			}
 		}
-
-		for (String key : keys) {
-			currentEnvironment.put(key, this.environment.getProperty(key));
-		}
-
-		return currentEnvironment;
 	}
 }

spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/EnvironmentVariablesProvider.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2016 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
+ *
+ *      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 org.springframework.cloud.task.batch.partition;
+
+import java.util.Map;
+
+import org.springframework.batch.item.ExecutionContext;
+
+/**
+ * Strategy interface to allow for advanced configuration of environment variables for
+ * each worker in a partitioned job.
+ *
+ * @author Michael Minella
+ *
+ * @since 1.0.2
+ */
+public interface EnvironmentVariablesProvider {
+
+	/**
+	 * Provides a {@link Map} of Strings to be used as environment variables.  This method
+	 * will be called for each worker step.  For example, if there are 5 partitions, this
+	 * method will be called 5 times.
+	 *
+	 * @param executionContext the {@link ExecutionContext} associated with the worker's
+	 * 			step
+	 * @return A {@link Map} of values to be used as environment variables
+	 */
+	Map<String, String> getEnvironmentVariables(ExecutionContext executionContext);
+}

spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/NoOpEnvironmentVariablesProvider.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2016 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
+ *
+ *      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 org.springframework.cloud.task.batch.partition;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.springframework.batch.item.ExecutionContext;
+
+/**
+ * A simple no-op implementation of the {@link EnvironmentVariablesProvider}.  It returns
+ * an empty {@link Map}.
+ *
+ * @author Michael Minella
+ *
+ * @since 1.0.2
+ */
+public class NoOpEnvironmentVariablesProvider implements EnvironmentVariablesProvider {
+
+	/**
+	 *
+	 * @param executionContext the {@link ExecutionContext} associated with the worker's
+	 * 			step
+	 * @return an empty {@link Map}
+	 */
+	@Override
+	public Map<String, String> getEnvironmentVariables(ExecutionContext executionContext) {
+		return new HashMap<>(0);
+	}
+}

spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/SimpleEnvironmentVariablesProvider.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2016 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
+ *
+ *      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 org.springframework.cloud.task.batch.partition;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+import org.springframework.batch.item.ExecutionContext;
+import org.springframework.core.env.AbstractEnvironment;
+import org.springframework.core.env.Environment;
+import org.springframework.core.env.MapPropertySource;
+import org.springframework.core.env.PropertySource;
+
+/**
+ * Copies all existing environment variables as made available in the {@link Environment}.
+ * The <code>environmentProperties</code> option provides the ability to override any
+ * specific values on an as needed basis.
+ *
+ * @author Michael Minella
+ *
+ * @since 1.0.2
+ */
+public class SimpleEnvironmentVariablesProvider implements EnvironmentVariablesProvider {
+
+	private Environment environment;
+
+	private Map<String, String> environmentProperties = new HashMap<>(0);
+
+	/**
+	 * @param environment The {@link Environment} for this context
+	 */
+	public SimpleEnvironmentVariablesProvider(Environment environment) {
+		this.environment = environment;
+	}
+
+	/**
+	 * @param environmentProperties a {@link Map} of properties used to override any values
+	 * configured in the current {@link Environment}
+	 */
+	public void setEnvironmentProperties(Map<String, String> environmentProperties) {
+		this.environmentProperties = environmentProperties;
+	}
+
+	@Override
+	public Map<String, String> getEnvironmentVariables(ExecutionContext executionContext) {
+
+		Map<String, String> environmentProperties = new HashMap<>(this.environmentProperties.size());
+		environmentProperties.putAll(getCurrentEnvironmentProperties());
+		environmentProperties.putAll(this.environmentProperties);
+
+		return environmentProperties;
+	}
+
+	private Map<String, String> getCurrentEnvironmentProperties() {
+		Map<String, String> currentEnvironment = new HashMap<>();
+
+		Set<String> keys = new HashSet<>();
+
+		for (PropertySource<?> propertySource : ((AbstractEnvironment) this.environment).getPropertySources()) {
+			if (propertySource instanceof MapPropertySource) {
+				keys.addAll(Arrays.asList(((MapPropertySource) propertySource).getPropertyNames()));
+			}
+		}
+
+		for (String key : keys) {
+			currentEnvironment.put(key, this.environment.getProperty(key));
+		}
+
+		return currentEnvironment;
+	}
+}

spring-cloud-task-batch/src/test/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandlerTests.java

@@ -134,7 +134,10 @@ public class DeployerPartitionHandlerTests {
 
 		when(this.jobExplorer.getStepExecution(1L, 4L)).thenReturn(workerStepExecutionFinish);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
+
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);
 
 		verify(this.taskLauncher).launch(this.appDeploymentRequestArgumentCaptor.capture());
@@ -190,6 +193,8 @@ public class DeployerPartitionHandlerTests {
 		when(this.jobExplorer.getStepExecution(1L, 5L)).thenReturn(workerStepExecutionFinish2);
 		when(this.jobExplorer.getStepExecution(1L, 6L)).thenReturn(workerStepExecutionFinish3);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);
 
@@ -236,6 +241,8 @@ public class DeployerPartitionHandlerTests {
 		when(this.jobExplorer.getStepExecution(1L, 5L)).thenReturn(workerStepExecutionFinish2);
 		when(this.jobExplorer.getStepExecution(1L, 6L)).thenReturn(workerStepExecutionFinish3);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);
 
@@ -282,6 +289,8 @@ public class DeployerPartitionHandlerTests {
 		when(this.jobExplorer.getStepExecution(1L, 5L)).thenReturn(workerStepExecutionFinish2);
 		when(this.jobExplorer.getStepExecution(1L, 6L)).thenReturn(workerStepExecutionFinish3);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);
 
@@ -319,13 +328,14 @@ public class DeployerPartitionHandlerTests {
 		StepExecution workerStepExecutionFinish = getStepExecutionFinish(workerStepExecutionStart, BatchStatus.COMPLETED);
 
 		DeployerPartitionHandler handler = new DeployerPartitionHandler(this.taskLauncher, this.jobExplorer, this.resource, "step1");
-		handler.setEnvironment(this.environment);
 
 		Map<String, String> environmentParameters = new HashMap<>(2);
 		environmentParameters.put("foo", "bar");
 		environmentParameters.put("baz", "qux");
 
-		handler.setEnvironmentProperties(environmentParameters);
+		SimpleEnvironmentVariablesProvider environmentVariablesProvider = new SimpleEnvironmentVariablesProvider(this.environment);
+		environmentVariablesProvider.setEnvironmentProperties(environmentParameters);
+		handler.setEnvironmentVariablesProvider(environmentVariablesProvider);
 
 		TaskExecution taskExecution = new TaskExecution();
 		taskExecution.setTaskName("partitionedJobTask");
@@ -336,6 +346,8 @@ public class DeployerPartitionHandlerTests {
 
 		when(this.jobExplorer.getStepExecution(1L, 4L)).thenReturn(workerStepExecutionFinish);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);
 
@@ -384,7 +396,9 @@ public class DeployerPartitionHandlerTests {
 		environmentParameters.put("foo", "bar");
 		environmentParameters.put("baz", "qux");
 
-		handler.setEnvironmentProperties(environmentParameters);
+		SimpleEnvironmentVariablesProvider environmentVariablesProvider = new SimpleEnvironmentVariablesProvider(this.environment);
+		environmentVariablesProvider.setEnvironmentProperties(environmentParameters);
+		handler.setEnvironmentVariablesProvider(environmentVariablesProvider);
 
 		TaskExecution taskExecution = new TaskExecution();
 		taskExecution.setTaskName("partitionedJobTask");
@@ -395,6 +409,8 @@ public class DeployerPartitionHandlerTests {
 
 		when(this.jobExplorer.getStepExecution(1L, 4L)).thenReturn(workerStepExecutionFinish);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);
 
@@ -450,6 +466,8 @@ public class DeployerPartitionHandlerTests {
 		when(this.jobExplorer.getStepExecution(1L, 4L)).thenReturn(workerStepExecutionFinish1);
 		when(this.jobExplorer.getStepExecution(1L, 5L)).thenReturn(workerStepExecutionFinish2);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 
 		Date startTime = new Date();
@@ -497,6 +515,8 @@ public class DeployerPartitionHandlerTests {
 		when(this.jobExplorer.getStepExecution(1L, 4L)).thenReturn(workerStepExecutionFinish1);
 		when(this.jobExplorer.getStepExecution(1L, 5L)).thenReturn(workerStepExecutionFinish2);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 
 		handler.handle(this.splitter, masterStepExecution);
@@ -530,6 +550,8 @@ public class DeployerPartitionHandlerTests {
 		when(this.jobExplorer.getStepExecution(1L, 4L)).thenReturn(workerStepExecutionFinish1);
 		when(this.jobExplorer.getStepExecution(1L, 5L)).thenReturn(workerStepExecutionFinish2);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);
@@ -570,6 +592,8 @@ public class DeployerPartitionHandlerTests {
 
 		when(this.jobExplorer.getStepExecution(1L, 4L)).thenReturn(workerStepExecutionFinish);
 
+		handler.afterPropertiesSet();
+
 		handler.beforeTask(taskExecution);
 		Collection<StepExecution> results = handler.handle(this.splitter, masterStepExecution);