Merge pull request #168 from jpraet/BATCH-2018

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java

@@ -46,15 +46,14 @@ import org.springframework.util.StringUtils;
  */
 public class MongoItemWriter<T> implements ItemWriter<T>, InitializingBean {
 
-	private static final String BUFFER_KEY_PREFIX = MongoItemWriter.class.getName() + ".BUFFER_KEY";
 	private MongoOperations template;
-	private final String bufferKey;
+	private final Object bufferKey;
 	private String collection;
 	private boolean delete = false;
 
 	public MongoItemWriter() {
 		super();
-		this.bufferKey = BUFFER_KEY_PREFIX + "." + hashCode();
+		this.bufferKey = new Object();
 	}
 
 	/**

spring-batch-infrastructure/src/main/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriter.java

@@ -37,37 +37,33 @@ import org.springframework.transaction.support.TransactionSynchronizationManager
  */
 public class TransactionAwareBufferedWriter extends Writer {
 
-	private static final String BUFFER_KEY_PREFIX = TransactionAwareBufferedWriter.class.getName() + ".BUFFER_KEY";
+	private final Object bufferKey;
 
-	private static final String CLOSE_KEY_PREFIX = TransactionAwareBufferedWriter.class.getName() + ".CLOSE_KEY";
-
-	private final String bufferKey;
-
-	private final String closeKey;
+	private final Object closeKey;
 
 	private FileChannel channel;
 
 	private final Runnable closeCallback;
-	
+
 	// default encoding for writing to output files - set to UTF-8.
 	private static final String DEFAULT_CHARSET = "UTF-8";
-	
+
 	private String encoding = DEFAULT_CHARSET;
-	
+
 	/**
 	 * Create a new instance with the underlying file channel provided, and a callback
 	 * to execute on close. The callback should clean up related resources like
 	 * output streams or channels.
 	 * 
-	 * @param channel channel used to do the actuall file IO
+	 * @param channel channel used to do the actual file IO
 	 * @param closeCallback callback to execute on close
 	 */
 	public TransactionAwareBufferedWriter(FileChannel channel, Runnable closeCallback) {
 		super();
 		this.channel = channel;
 		this.closeCallback = closeCallback;
-		this.bufferKey = BUFFER_KEY_PREFIX + "." + hashCode();
-		this.closeKey = CLOSE_KEY_PREFIX + "." + hashCode();
+		this.bufferKey = new Object();
+		this.closeKey = new Object();
 	}
 
 	public void setEncoding(String encoding) {
@@ -88,7 +84,7 @@ public class TransactionAwareBufferedWriter extends Writer {
 				public void afterCompletion(int status) {
 					clear();
 				}
-				
+
 				@Override
 				public void beforeCommit(boolean readOnly) {
 					try {
@@ -145,7 +141,7 @@ public class TransactionAwareBufferedWriter extends Writer {
 		if (!transactionActive()) {
 			return 0L;
 		}
-		try {			
+		try {
 			return getCurrentBuffer().toString().getBytes(encoding).length;
 		} catch (UnsupportedEncodingException e) {
 			throw new WriteFailedException("Could not determine buffer size because of unsupported encoding: " + encoding, e);

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemWriterTests.java

@@ -2,17 +2,28 @@ package org.springframework.batch.item.data;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.doAnswer;
 
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import org.springframework.batch.support.transaction.ResourcelessTransactionManager;
+import org.springframework.batch.support.transaction.TransactionAwareBufferedWriter;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.transaction.PlatformTransactionManager;
 import org.springframework.transaction.TransactionStatus;
@@ -232,4 +243,52 @@ public class MongoItemWriterTests {
 		verify(template).remove(items.get(0), "collection");
 		verify(template).remove(items.get(1), "collection");
 	}
+	
+	// BATCH-2018
+	@Test
+	public void testResourceKeyCollision() throws Exception {
+		final int limit = 5000;
+		final MongoItemWriter<String>[] writers = new MongoItemWriter[limit];
+		final String[] results = new String[limit];
+		for(int i = 0; i< limit; i++) {
+			final int index = i;
+			MongoOperations mongoOperations = mock(MongoOperations.class);
+			
+			doAnswer(new Answer<Void>() {
+				@Override
+				public Void answer(InvocationOnMock invocation)
+						throws Throwable {
+					String val = (String) invocation.getArguments()[0];
+					if(results[index] == null) {
+						results[index] = val;
+					} else {
+						results[index] += val;
+					}
+					return null;
+				}
+			}).when(mongoOperations).save(any(String.class));
+			writers[i] = new MongoItemWriter<String>();
+			writers[i].setTemplate(mongoOperations);
+		}
+		
+		new TransactionTemplate(transactionManager).execute(new TransactionCallback() {
+            @Override
+			public Object doInTransaction(TransactionStatus status) {
+				try {
+					for(int i=0; i< limit; i++) {
+						writers[i].write(Collections.singletonList(String.valueOf(i)));
+					}
+				}
+				catch (Exception e) {
+					throw new IllegalStateException("Unexpected Exception", e);
+				}
+				return null;
+			}
+		});		
+		
+		for(int i=0; i< limit; i++) {
+			assertEquals(String.valueOf(i), results[i]);
+		}				
+	}
+	
 }

spring-batch-infrastructure/src/test/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriterTests.java

@@ -20,15 +20,20 @@ import static org.junit.Assert.fail;
 import static org.mockito.Matchers.anyObject;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.any;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 
+import javax.validation.constraints.AssertTrue;
+
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import org.springframework.transaction.PlatformTransactionManager;
 import org.springframework.transaction.TransactionStatus;
 import org.springframework.transaction.support.TransactionCallback;
@@ -296,6 +301,53 @@ public class TransactionAwareBufferedWriterTests {
 			assertEquals("Could not write to output buffer", ffe.getMessage());
 		}
 	}
+	
+	// BATCH-2018
+	@Test
+	@SuppressWarnings({"unchecked", "rawtypes"})
+	public void testResourceKeyCollision() throws Exception {
+		final int limit = 5000;
+		final TransactionAwareBufferedWriter[] writers = new TransactionAwareBufferedWriter[limit];
+		final String[] results = new String[limit];
+		for(int i = 0; i< limit; i++) {
+			final int index = i;
+			FileChannel fileChannel = mock(FileChannel.class);
+			when(fileChannel.write(any(ByteBuffer.class))).thenAnswer(new Answer<Integer>() {
+				@Override
+				public Integer answer(InvocationOnMock invocation)
+						throws Throwable {
+					ByteBuffer buffer = (ByteBuffer) invocation.getArguments()[0];
+					String val = new String(buffer.array(), "UTF-8");
+					if(results[index] == null) {
+						results[index] = val;
+					} else {
+						results[index] += val;
+					}
+					return buffer.limit();
+				}
+			});
+			writers[i] = new TransactionAwareBufferedWriter(fileChannel, null);
+		}
+		
+		new TransactionTemplate(transactionManager).execute(new TransactionCallback() {
+            @Override
+			public Object doInTransaction(TransactionStatus status) {
+				try {
+					for(int i=0; i< limit; i++) {
+						writers[i].write(String.valueOf(i));
+					}
+				}
+				catch (IOException e) {
+					throw new IllegalStateException("Unexpected IOException", e);
+				}
+				return null;
+			}
+		});		
+		
+		for(int i=0; i< limit; i++) {
+			assertEquals(String.valueOf(i), results[i]);
+		}				
+	}
 
 	private String getStringFromByteBuffer(ByteBuffer bb) {
 		byte[] bytearr = new byte[bb.remaining()];