diff --git a/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandler.java b/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandler.java index 0119d2c9..9e764966 100644 --- a/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandler.java +++ b/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 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 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 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 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 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 getArguments(List arguments) { - Map 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 getCurrentEnvironmentProperties() { - Map currentEnvironment = new HashMap<>(); + @Override + public void afterPropertiesSet() throws Exception { + if(this.environmentVariablesProvider == null) { + this.environmentVariablesProvider = + new SimpleEnvironmentVariablesProvider(this.environment); - Set 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; } } diff --git a/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/EnvironmentVariablesProvider.java b/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/EnvironmentVariablesProvider.java new file mode 100644 index 00000000..7a7fb428 --- /dev/null +++ b/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 getEnvironmentVariables(ExecutionContext executionContext); +} diff --git a/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/NoOpEnvironmentVariablesProvider.java b/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/NoOpEnvironmentVariablesProvider.java new file mode 100644 index 00000000..1e42d0e7 --- /dev/null +++ b/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 getEnvironmentVariables(ExecutionContext executionContext) { + return new HashMap<>(0); + } +} diff --git a/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/SimpleEnvironmentVariablesProvider.java b/spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/SimpleEnvironmentVariablesProvider.java new file mode 100644 index 00000000..0451c334 --- /dev/null +++ b/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 environmentProperties 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 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 environmentProperties) { + this.environmentProperties = environmentProperties; + } + + @Override + public Map getEnvironmentVariables(ExecutionContext executionContext) { + + Map environmentProperties = new HashMap<>(this.environmentProperties.size()); + environmentProperties.putAll(getCurrentEnvironmentProperties()); + environmentProperties.putAll(this.environmentProperties); + + return environmentProperties; + } + + private Map getCurrentEnvironmentProperties() { + Map currentEnvironment = new HashMap<>(); + + Set 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; + } +} diff --git a/spring-cloud-task-batch/src/test/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandlerTests.java b/spring-cloud-task-batch/src/test/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandlerTests.java index 6d295831..417a4d1d 100644 --- a/spring-cloud-task-batch/src/test/java/org/springframework/cloud/task/batch/partition/DeployerPartitionHandlerTests.java +++ b/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 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 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 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 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 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 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 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 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 results = handler.handle(this.splitter, masterStepExecution); diff --git a/spring-cloud-task-samples/partitioned-batch-job/src/main/java/io/spring/JobConfiguration.java b/spring-cloud-task-samples/partitioned-batch-job/src/main/java/io/spring/JobConfiguration.java index 59c759da..4b9e7f81 100644 --- a/spring-cloud-task-samples/partitioned-batch-job/src/main/java/io/spring/JobConfiguration.java +++ b/spring-cloud-task-samples/partitioned-batch-job/src/main/java/io/spring/JobConfiguration.java @@ -42,10 +42,12 @@ import org.springframework.cloud.deployer.spi.local.LocalTaskLauncher; import org.springframework.cloud.deployer.spi.task.TaskLauncher; import org.springframework.cloud.task.batch.partition.DeployerPartitionHandler; import org.springframework.cloud.task.batch.partition.DeployerStepExecutionHandler; +import org.springframework.cloud.task.batch.partition.SimpleEnvironmentVariablesProvider; import org.springframework.context.ConfigurableApplicationContext; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.context.annotation.Profile; +import org.springframework.core.env.Environment; import org.springframework.core.io.Resource; /** @@ -72,6 +74,9 @@ public class JobConfiguration { @Autowired private DelegatingResourceLoader resourceLoader; + @Autowired + private Environment environment; + private static final int GRID_SIZE = 4; @Bean @@ -101,7 +106,10 @@ public class JobConfiguration { Map environmentProperties = new HashMap<>(); environmentProperties.put("spring.profiles.active", "worker"); - partitionHandler.setEnvironmentProperties(environmentProperties); + SimpleEnvironmentVariablesProvider environmentVariablesProvider = new SimpleEnvironmentVariablesProvider(this.environment); + environmentVariablesProvider.setEnvironmentProperties(environmentProperties); + partitionHandler.setEnvironmentVariablesProvider(environmentVariablesProvider); + partitionHandler.setMaxWorkers(2); return partitionHandler; @@ -130,9 +138,7 @@ public class JobConfiguration { @Bean @Profile("worker") public DeployerStepExecutionHandler stepExecutionHandler(JobExplorer jobExplorer) { - DeployerStepExecutionHandler handler = new DeployerStepExecutionHandler(this.context, jobExplorer, this.jobRepository); - - return handler; + return new DeployerStepExecutionHandler(this.context, jobExplorer, this.jobRepository); } @Bean