diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriter.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriter.java index e16fcd7ec..d5de7014f 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriter.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriter.java @@ -72,34 +72,40 @@ public class TransactionAwareBufferedWriter extends Writer { TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronizationAdapter() { @Override public void afterCompletion(int status) { - if (status == STATUS_COMMITTED) { - complete(); + try { + if (status == STATUS_COMMITTED) { + complete(); + } + } + catch (IOException e) { + throw new FlushFailedException("Could not write to output buffer", e); + } + finally { + clear(); } } - private void complete() { + private void complete() throws IOException { StringBuffer buffer = (StringBuffer) TransactionSynchronizationManager.getResource(bufferKey); if (buffer != null) { - try { - writer.write(buffer.toString()); - writer.flush(); - if (TransactionSynchronizationManager.hasResource(closeKey)) { - writer.close(); - closeCallback.run(); - } - } - catch (IOException e) { - throw new FlushFailedException("Could not write to output buffer", e); - } - finally { - TransactionSynchronizationManager.unbindResource(bufferKey); - if (TransactionSynchronizationManager.hasResource(closeKey)) { - TransactionSynchronizationManager.unbindResource(closeKey); - } + writer.write(buffer.toString()); + writer.flush(); + if (TransactionSynchronizationManager.hasResource(closeKey)) { + writer.close(); + closeCallback.run(); } } } + private void clear() { + if (TransactionSynchronizationManager.hasResource(bufferKey)) { + TransactionSynchronizationManager.unbindResource(bufferKey); + } + if (TransactionSynchronizationManager.hasResource(closeKey)) { + TransactionSynchronizationManager.unbindResource(closeKey); + } + } + }); } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriterTests.java index 51e9fd946..de834facc 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriterTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/support/transaction/TransactionAwareBufferedWriterTests.java @@ -174,4 +174,10 @@ public class TransactionAwareBufferedWriterTests { assertEquals("", stringWriter.toString()); } + @Test + public void testCleanUpAfterRollback() throws Exception { + testWriteWithRollback(); + testWriteWithCommit(); + } + }