From 526dc41af9ca12dc9659354a9f079b61aa0b13e0 Mon Sep 17 00:00:00 2001 From: Michael Minella Date: Fri, 22 Jul 2016 10:29:36 -0500 Subject: [PATCH] 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 --- .../partition/DeployerPartitionHandler.java | 66 ++++---------- .../EnvironmentVariablesProvider.java | 42 +++++++++ .../NoOpEnvironmentVariablesProvider.java | 43 +++++++++ .../SimpleEnvironmentVariablesProvider.java | 87 +++++++++++++++++++ .../DeployerPartitionHandlerTests.java | 30 ++++++- .../main/java/io/spring/JobConfiguration.java | 14 ++- 6 files changed, 226 insertions(+), 56 deletions(-) create mode 100644 spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/EnvironmentVariablesProvider.java create mode 100644 spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/NoOpEnvironmentVariablesProvider.java create mode 100644 spring-cloud-task-batch/src/main/java/org/springframework/cloud/task/batch/partition/SimpleEnvironmentVariablesProvider.java 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