diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/SimpleStepContext.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/SimpleStepContext.java index b99e20971..a06f12a43 100644 --- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/SimpleStepContext.java +++ b/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/SimpleStepContext.java @@ -81,11 +81,8 @@ public class SimpleStepContext extends SynchronizedAttributeAccessor implements super.setAttribute(name, value); if (streamManager != null && (value instanceof ItemStream)) { ItemStream stream = (ItemStream) value; - streamManager.register(this, stream); stream.open(); - if (streamContext != null) { - stream.restoreFrom(streamContext); - } + streamManager.register(this, stream, streamContext); } } @@ -203,11 +200,10 @@ public class SimpleStepContext extends SynchronizedAttributeAccessor implements return streamManager.getStreamContext(this); } - /* - * (non-Javadoc) - * @see org.springframework.batch.execution.scope.StepContext#setInitialStreamContext(org.springframework.batch.item.StreamContext) + /* (non-Javadoc) + * @see org.springframework.batch.execution.scope.StepContext#restoreFrom(org.springframework.batch.item.StreamContext) */ - public void setInitialStreamContext(StreamContext streamContext) { + public void restoreFrom(StreamContext streamContext) { this.streamContext = streamContext; } diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/StepContext.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/StepContext.java index 457729b3c..9ca89bfa3 100644 --- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/StepContext.java +++ b/spring-batch-execution/src/main/java/org/springframework/batch/execution/scope/StepContext.java @@ -55,12 +55,13 @@ public interface StepContext extends AttributeAccessor, StreamContextProvider { void close(); /** - * Provide the stream context that will be needed to restore - * {@link ItemStream} instances. If this is not set the streams will simply - * not be initialised and repositioned for restart (which is sometimes - * desirable). + * Provide the stream context needed to restore {@link ItemStream} + * instances. Implementations are free to use this as necessary (e.g. lazily + * if all the streams are not available at once). If this is not set the + * streams will simply not be initialised and repositioned for restart + * (which is sometimes desirable). * * @param streamContext */ - void setInitialStreamContext(StreamContext streamContext); + void restoreFrom(StreamContext streamContext); } \ No newline at end of file diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/simple/SimpleStepExecutor.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/simple/SimpleStepExecutor.java index dcda44d86..76c126f67 100644 --- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/simple/SimpleStepExecutor.java +++ b/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/simple/SimpleStepExecutor.java @@ -190,7 +190,7 @@ public class SimpleStepExecutor { final boolean saveStreamContext = step.isSaveStreamContext(); if (saveStreamContext && isRestart) { - stepContext.setInitialStreamContext(stepInstance.getStreamContext()); + stepContext.restoreFrom(stepInstance.getStreamContext()); } try { diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/scope/SimpleStepContextTests.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/scope/SimpleStepContextTests.java index 4d69b7827..1dbce4663 100644 --- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/scope/SimpleStepContextTests.java +++ b/spring-batch-execution/src/test/java/org/springframework/batch/execution/scope/SimpleStepContextTests.java @@ -171,7 +171,7 @@ public class SimpleStepContextTests extends TestCase { public void open(Object key) { } - public void register(Object key, ItemStream stream) { + public void register(Object key, ItemStream stream, StreamContext streamContext) { map.put(key, stream); } diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/simple/SimpleStepExecutorTests.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/simple/SimpleStepExecutorTests.java index b1b27b4f9..1d5db23ea 100644 --- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/simple/SimpleStepExecutorTests.java +++ b/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/simple/SimpleStepExecutorTests.java @@ -98,10 +98,13 @@ public class SimpleStepExecutorTests extends TestCase { */ protected void setUp() throws Exception { super.setUp(); + ResourcelessTransactionManager transactionManager = new ResourcelessTransactionManager(); + stepConfiguration = new SimpleStep(); stepConfiguration.setTasklet(getTasklet(new String[] { "foo", "bar", "spam" })); stepConfiguration.setJobRepository(new JobRepositorySupport()); - stepConfiguration.setTransactionManager(new ResourcelessTransactionManager()); + stepConfiguration.setTransactionManager(transactionManager); + stepExecutor = (SimpleStepExecutor) stepConfiguration.createStepExecutor(); template = new RepeatTemplate(); template.setCompletionPolicy(new SimpleCompletionPolicy(1)); @@ -113,6 +116,11 @@ public class SimpleStepExecutorTests extends TestCase { jobInstance = new JobInstance(new Long(0), new JobParameters()); jobInstance.setJob(new JobSupport("FOO")); + + SimpleStreamManager streamManager = new SimpleStreamManager(transactionManager); + streamManager.setUseClassNameAsPrefix(false); + stepExecutor.setStreamManager(streamManager); + } public void testStepExecutor() throws Exception { @@ -304,6 +312,7 @@ public class SimpleStepExecutorTests extends TestCase { stepConfiguration.setSaveStreamContext(true); JobExecution jobExecutionContext = new JobExecution(jobInstance); StepExecution stepExecution = new StepExecution(step, jobExecutionContext); + stepExecution.getStep().setStreamContext( new StreamContext(PropertiesConverter.stringToProperties("foo=bar"))); diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/SimpleStreamManager.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/SimpleStreamManager.java index e822c3e23..34f9a792e 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/SimpleStreamManager.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/SimpleStreamManager.java @@ -20,8 +20,8 @@ import java.util.HashMap; import java.util.Iterator; import java.util.LinkedHashSet; import java.util.Map; -import java.util.Properties; import java.util.Set; +import java.util.Map.Entry; import org.springframework.batch.item.ItemStream; import org.springframework.batch.item.StreamContext; @@ -32,6 +32,7 @@ import org.springframework.transaction.support.DefaultTransactionDefinition; import org.springframework.transaction.support.TransactionSynchronization; import org.springframework.transaction.support.TransactionSynchronizationAdapter; import org.springframework.transaction.support.TransactionSynchronizationManager; +import org.springframework.util.ClassUtils; /** * Simple {@link StreamManager} that tries to resolve conflicts between key @@ -48,6 +49,8 @@ public class SimpleStreamManager implements StreamManager { private PlatformTransactionManager transactionManager; + private boolean useClassNameAsPrefix = true; + /** * @param transactionManager a {@link PlatformTransactionManager} */ @@ -63,6 +66,18 @@ public class SimpleStreamManager implements StreamManager { super(); } + /** + * Public setter for the flag. If this is true then the class name of the + * streams will be used as a prefix in the {@link StreamContext} in + * {@link #getStreamContext(Object)}. The default value is true, which + * gives the best chance of unique key names in the context. + * + * @param useClassNameAsPrefix the flag to set (default true). + */ + public void setUseClassNameAsPrefix(boolean useClassNameAsPrefix) { + this.useClassNameAsPrefix = useClassNameAsPrefix; + } + /** * Public setter for the {@link PlatformTransactionManager}. * @param transactionManager the {@link PlatformTransactionManager} to set @@ -72,42 +87,28 @@ public class SimpleStreamManager implements StreamManager { } /** - * Simple aggregate statistics provider for the contributions registered - * under the given key. + * Simple aggregate {@link StreamContext} provider for the contributions + * registered under the given key. * * @see org.springframework.batch.item.stream.StreamManager#getStreamContext(java.lang.Object) */ public StreamContext getStreamContext(Object key) { - Set set = new LinkedHashSet(); - synchronized (registry) { - Collection collection = (Collection) registry.get(key); - if (collection != null) { - set = new LinkedHashSet(collection); - } - } - return aggregate(set); - } - - /** - * @param list a list of {@link ItemStream}s - * @return aggregated streamcontext - */ - private StreamContext aggregate(Collection list) { - Properties result = new Properties(); - for (Iterator iterator = list.iterator(); iterator.hasNext();) { - ItemStream provider = (ItemStream) iterator.next(); - Properties properties = provider.getStreamContext().getProperties(); - if (properties != null) { - String prefix = ""; // ClassUtils.getShortClassName(provider.getClass()) - // + "."; - for (Iterator propiter = properties.keySet().iterator(); propiter.hasNext();) { - String key = (String) propiter.next(); - String value = properties.getProperty(key); - result.setProperty(prefix + key, value); + final StreamContext result = new StreamContext(); + iterate(key, new Callback() { + public void execute(ItemStream stream) { + StreamContext context = stream.getStreamContext(); + String prefix = ClassUtils.getQualifiedName(stream.getClass()) + "."; + if (!useClassNameAsPrefix) { + prefix = ""; + } + for (Iterator iterator = context.entrySet().iterator(); iterator.hasNext();) { + Entry entry = (Entry) iterator.next(); + String contextKey = prefix + entry.getKey(); + result.put(contextKey, entry.getValue()); } } - } - return new StreamContext(result); + }); + return result; } /** @@ -115,17 +116,41 @@ public class SimpleStreamManager implements StreamManager { * the provided key. * * @see org.springframework.batch.item.stream.StreamManager#register(java.lang.Object, - * org.springframework.batch.item.ItemStream) + * org.springframework.batch.item.ItemStream, StreamContext) */ - public void register(Object key, ItemStream provider) { + public void register(Object key, ItemStream stream, StreamContext streamContext) { synchronized (registry) { Set set = (Set) registry.get(key); if (set == null) { set = new LinkedHashSet(); registry.put(key, set); } - set.add(provider); + set.add(stream); } + if (streamContext != null) { + stream.restoreFrom(extract(stream, streamContext)); + } + } + + /** + * @param stream + * @param streamContext + * @return + */ + private StreamContext extract(ItemStream stream, StreamContext context) { + StreamContext result = new StreamContext(); + String prefix = ClassUtils.getQualifiedName(stream.getClass()) + "."; + if (!useClassNameAsPrefix) { + prefix = ""; + } + for (Iterator iterator = context.entrySet().iterator(); iterator.hasNext();) { + Entry entry = (Entry) iterator.next(); + String contextKey = (String) entry.getKey(); + if (contextKey.startsWith(prefix)) { + result.put(contextKey.substring(prefix.length()), entry.getValue()); + } + } + return result; } /** diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/StreamManager.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/StreamManager.java index fa1061e18..a8ae6b836 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/StreamManager.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/stream/StreamManager.java @@ -37,8 +37,9 @@ public interface StreamManager { * * @param key the key under which to add the provider * @param stream an {@link ItemStream} + * @param streamContext the context (may be null) to restore from on registration */ - void register(Object key, ItemStream stream); + void register(Object key, ItemStream stream, StreamContext streamContext); /** * Extract and aggregate the {@link StreamContext} from all streams under diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/stream/SimpleStreamManagerTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/stream/SimpleStreamManagerTests.java new file mode 100644 index 000000000..c7d9529d5 --- /dev/null +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/stream/SimpleStreamManagerTests.java @@ -0,0 +1,246 @@ +/* + * Copyright 2006-2007 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.batch.item.stream; + +import java.util.ArrayList; +import java.util.List; + +import junit.framework.TestCase; + +import org.springframework.batch.item.StreamContext; +import org.springframework.batch.item.StreamException; +import org.springframework.batch.support.PropertiesConverter; +import org.springframework.batch.support.transaction.ResourcelessTransactionManager; +import org.springframework.transaction.TransactionException; +import org.springframework.transaction.TransactionStatus; +import org.springframework.util.ClassUtils; + +/** + * @author Dave Syer + * + */ +public class SimpleStreamManagerTests extends TestCase { + + private SimpleStreamManager manager = new SimpleStreamManager(new ResourcelessTransactionManager()); + + private ItemStreamAdapter stream = new ItemStreamAdapterExtension(); + + private List list = new ArrayList(); + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#SimpleStreamManager(org.springframework.transaction.PlatformTransactionManager)}. + */ + public void testSimpleStreamManagerPlatformTransactionManager() { + manager = new SimpleStreamManager(); + try { + manager.getTransaction("foo"); + fail("Expected NullPointerException"); + } + catch (NullPointerException e) { + // expected; + } + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#setTransactionManager(org.springframework.transaction.PlatformTransactionManager)}. + */ + public void testSetTransactionManager() { + manager.setTransactionManager(new ResourcelessTransactionManager() { + protected Object doGetTransaction() throws TransactionException { + list.add("bar"); + return super.doGetTransaction(); + } + }); + manager.getTransaction("foo"); + assertEquals("bar", list.get(0)); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#getStreamContext(java.lang.Object)}. + */ + public void testGetStreamContextEmpty() { + StreamContext streamContext = manager.getStreamContext("foo"); + assertEquals(0, streamContext.entrySet().size()); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#getStreamContext(java.lang.Object)}. + */ + public void testGetStreamContextNotEmpty() { + manager.register("foo", stream, null); + StreamContext streamContext = manager.getStreamContext("foo"); + assertEquals(1, streamContext.entrySet().size()); + assertEquals("bar", streamContext.getString(ClassUtils.getQualifiedName(stream.getClass()) + ".foo")); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#getStreamContext(java.lang.Object)}. + */ + public void testGetStreamContextNotEmptyAndRestore() { + testGetStreamContextNotEmpty(); + StreamContext context = manager.getStreamContext("foo"); + // Register again, now with the context that was created from the same + // stream... + manager.register("foo", stream, context); + assertEquals(1, list.size()); + // The list should have the foo= map value from the sub-context + assertEquals("bar", list.get(0)); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#getStreamContext(java.lang.Object)}. + */ + public void testGetStreamContextNotEmptyAndRestoreWithNoPrefix() { + StreamContext context = new StreamContext(PropertiesConverter.stringToProperties("foo=bar")); + manager.setUseClassNameAsPrefix(false); + manager.register("foo", stream, context); + assertEquals(1, list.size()); + // The list should have the foo= map value from the sub-context + assertEquals("bar", list.get(0)); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#getStreamContext(java.lang.Object)}. + */ + public void testGetStreamContextWithNoPrefix() { + manager.setUseClassNameAsPrefix(false); + manager.register("foo", stream, null); + StreamContext context = manager.getStreamContext("foo"); + assertEquals(1, context.entrySet().size()); + // The list should have the foo= map value from the sub-context + assertEquals("bar", context.getString("foo")); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#getStreamContext(java.lang.Object)}. + */ + public void testGetStreamContextTwoRegistrations() { + manager.register("foo", new ItemStreamAdapter() { + public StreamContext getStreamContext() { + return new StreamContext(PropertiesConverter.stringToProperties("foo=bar")); + } + }, null); + manager.register("foo", new ItemStreamAdapter() { + public StreamContext getStreamContext() { + return new StreamContext(PropertiesConverter.stringToProperties("foo=spam")); + } + }, null); + StreamContext streamContext = manager.getStreamContext("foo"); + assertEquals(2, streamContext.entrySet().size()); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#close(java.lang.Object)}. + */ + public void testClose() { + manager.register("foo", new ItemStreamAdapter() { + public void close() throws StreamException { + list.add("bar"); + super.close(); + } + }, null); + manager.close("foo"); + assertEquals(1, list.size()); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#commit(org.springframework.transaction.TransactionStatus)}. + */ + public void testCommit() { + manager.register("foo", new ItemStreamAdapter() { + public boolean isMarkSupported() { + return true; + } + + public void mark(StreamContext streamContext) { + list.add("bar"); + } + }, null); + TransactionStatus status = manager.getTransaction("foo"); + manager.commit(status); + assertEquals(1, list.size()); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#commit(org.springframework.transaction.TransactionStatus)}. + */ + public void testCommitWithoutMark() { + manager.register("foo", new ItemStreamAdapter() { + public void mark(StreamContext streamContext) { + list.add("bar"); + } + }, null); + TransactionStatus status = manager.getTransaction("foo"); + manager.commit(status); + assertEquals(0, list.size()); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#rollback(org.springframework.transaction.TransactionStatus)}. + */ + public void testRollback() { + manager.register("foo", new ItemStreamAdapter() { + public boolean isMarkSupported() { + return true; + } + + public void reset(StreamContext streamContext) { + list.add("bar"); + } + }, null); + TransactionStatus status = manager.getTransaction("foo"); + manager.rollback(status); + assertEquals(1, list.size()); + } + + /** + * Test method for + * {@link org.springframework.batch.item.stream.SimpleStreamManager#rollback(org.springframework.transaction.TransactionStatus)}. + */ + public void testRollbackWithoutMark() { + manager.register("foo", new ItemStreamAdapter() { + public void reset(StreamContext streamContext) { + list.add("bar"); + } + }, null); + TransactionStatus status = manager.getTransaction("foo"); + manager.rollback(status); + assertEquals(0, list.size()); + } + + private final class ItemStreamAdapterExtension extends ItemStreamAdapter { + public StreamContext getStreamContext() { + return new StreamContext(PropertiesConverter.stringToProperties("foo=bar")); + } + + public void restoreFrom(StreamContext context) { + list.add(context.getString("foo")); + } + } + +}