BATCH-2434 Improve TransactionAwareBufferedWriter

TransactionAwareBufferedWriter offers a number of optimization
potentials. First it buffers at the char level. Buffering at the byte
level instead saves about 50% memory usage in common cases. Second it
does not overwrite any of the  #write(String) methods leading to
unnecessary intermediate copies.

 * buffer in a ByteArrayOutputStream instead of StringBuilder
 * overwrite #write(String) methods to avoid copies

Together these two changes should help to reduce both live set size and
allocation rate.

Issue: BATCH-2434
https://jira.spring.io/browse/BATCH-2434

Author: marschall

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/FlatFileItemWriter.java

@@ -506,23 +506,13 @@ public void close() {
 
 		private void closeStream() {
 			try {
-				if (fileChannel != null) {
-					fileChannel.close();
+				if (os != null) {
+					os.close();
 				}
 			}
 			catch (IOException ioe) {
 				throw new ItemStreamException("Unable to close the the ItemWriter", ioe);
 			}
-			finally {
-				try {
-					if (os != null) {
-						os.close();
-					}
-				}
-				catch (IOException ioe) {
-					throw new ItemStreamException("Unable to close the the ItemWriter", ioe);
-				}
-			}
 		}
 
 		/**
@@ -561,7 +551,7 @@ private void initializeBufferedWriter() throws IOException {
 			os = new FileOutputStream(file.getAbsolutePath(), true);
 			fileChannel = os.getChannel();
 
-			outputBufferedWriter = getBufferedWriter(fileChannel, encoding);
+			outputBufferedWriter = getBufferedWriter(os, encoding);
 			outputBufferedWriter.flush();
 
 			if (append) {
@@ -591,11 +581,11 @@ public boolean isInitialized() {
 		 * Returns the buffered writer opened to the beginning of the file
 		 * specified by the absolute path name contained in absoluteFileName.
 		 */
-		private Writer getBufferedWriter(FileChannel fileChannel, String encoding) {
+		private Writer getBufferedWriter(FileOutputStream outputStream, String encoding) {
 			try {
 				final FileChannel channel = fileChannel;
 				if (transactional) {
-					TransactionAwareBufferedWriter writer = new TransactionAwareBufferedWriter(channel, new Runnable() {
+					TransactionAwareBufferedWriter writer = new TransactionAwareBufferedWriter(outputStream, new Runnable() {
 						@Override
 						public void run() {
 							closeStream();

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxEventItemWriter.java

@@ -456,7 +456,7 @@ private void open(long position) {
 		try {
 			final FileChannel channel = fileChannel;
 			if (transactional) {
-				TransactionAwareBufferedWriter writer = new TransactionAwareBufferedWriter(channel, new Runnable() {
+				TransactionAwareBufferedWriter writer = new TransactionAwareBufferedWriter(os, new Runnable() {
 					@Override
 					public void run() {
 						closeStream();

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

@@ -15,10 +15,14 @@
  */
 package org.springframework.batch.support.transaction;
 
+import java.io.BufferedOutputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
 import java.io.UnsupportedEncodingException;
 import java.io.Writer;
-import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 
 import org.springframework.batch.item.WriteFailedException;
@@ -41,7 +45,7 @@ public class TransactionAwareBufferedWriter extends Writer {
 
 	private final Object closeKey;
 
-	private FileChannel channel;
+	private FileOutputStream outputStream;
 
 	private final Runnable closeCallback;
 
@@ -57,12 +61,12 @@ public class TransactionAwareBufferedWriter extends Writer {
 	 * to execute on close. The callback should clean up related resources like
 	 * output streams or channels.
 	 *
-	 * @param channel channel used to do the actual file IO
+	 * @param outputStream channel used to do the actual file IO
 	 * @param closeCallback callback to execute on close
 	 */
-	public TransactionAwareBufferedWriter(FileChannel channel, Runnable closeCallback) {
+	public TransactionAwareBufferedWriter(FileOutputStream outputStream, Runnable closeCallback) {
 		super();
-		this.channel = channel;
+		this.outputStream = outputStream;
 		this.closeCallback = closeCallback;
 		this.bufferKey = new Object();
 		this.closeKey = new Object();
@@ -71,6 +75,47 @@ public TransactionAwareBufferedWriter(FileChannel channel, Runnable closeCallbac
 	public void setEncoding(String encoding) {
 		this.encoding = encoding;
 	}
+	
+	/**
+	 * Buffers the output on a {@code byte[]} level. Compared to
+	 * buffering on a {@code char[]} level this saves about 50% memory
+	 * for common encodings and input.
+	 */
+	static final class Buffer {
+
+		private ByteArrayOutputStream outputStream;
+
+		private OutputStreamWriter writer;
+		
+		Buffer(String encoding) throws UnsupportedEncodingException {
+			this.outputStream = new ByteArrayOutputStream();
+			this.writer = new OutputStreamWriter(outputStream, encoding);
+		}
+
+		int length() throws IOException {
+			this.writer.flush();
+			return this.outputStream.size();
+		}
+
+		void append(char[] cbuf, int off, int len) throws IOException {
+			this.writer.write(cbuf, off, len);
+		}
+
+		void append(String str, int off, int len) throws IOException {
+			this.writer.write(str, off, len);
+		}
+		
+		byte[] getBytes() throws IOException {
+			this.writer.flush();
+			return this.outputStream.toByteArray();
+		}
+		
+		void writeTo(OutputStream stream) throws IOException {
+			this.writer.flush();
+			this.outputStream.writeTo(stream);
+		}
+		
+	}
 
 	/**
 	 * Flag to indicate that changes should be force-synced to disk on flush.
@@ -88,11 +133,15 @@ public void setForceSync(boolean forceSync) {
 	/**
 	 * @return
 	 */
-	private StringBuilder getCurrentBuffer() {
+	private Buffer getCurrentBuffer() {
 
 		if (!TransactionSynchronizationManager.hasResource(bufferKey)) {
 
-			TransactionSynchronizationManager.bindResource(bufferKey, new StringBuilder());
+			try {
+				TransactionSynchronizationManager.bindResource(bufferKey, new Buffer(this.encoding));
+			} catch (UnsupportedEncodingException e) {
+				throw new WriteFailedException("Could not determine buffer size because of unsupported encoding: " + encoding, e);
+			}
 
 			TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronizationAdapter() {
 				@Override
@@ -113,18 +162,11 @@ public void beforeCommit(boolean readOnly) {
 				}
 
 				private void complete() throws IOException {
-					StringBuilder buffer = (StringBuilder) TransactionSynchronizationManager.getResource(bufferKey);
+					Buffer buffer = (Buffer) TransactionSynchronizationManager.getResource(bufferKey);
 					if (buffer != null) {
-						String string = buffer.toString();
-						byte[] bytes = string.getBytes(encoding);
-						int bufferLength = bytes.length;
-						ByteBuffer bb = ByteBuffer.wrap(bytes);
-						int bytesWritten = channel.write(bb);
-						if(bytesWritten != bufferLength) {
-							throw new IOException("All bytes to be written were not successfully written");
-						}
+						buffer.writeTo(outputStream);
 						if (forceSync) {
-							channel.force(false);
+							outputStream.getChannel().force(false);
 						}
 						if (TransactionSynchronizationManager.hasResource(closeKey)) {
 							closeCallback.run();
@@ -145,7 +187,7 @@ private void clear() {
 
 		}
 
-		return (StringBuilder) TransactionSynchronizationManager.getResource(bufferKey);
+		return (Buffer) TransactionSynchronizationManager.getResource(bufferKey);
 
 	}
 
@@ -160,9 +202,9 @@ public long getBufferSize() {
 			return 0L;
 		}
 		try {
-			return getCurrentBuffer().toString().getBytes(encoding).length;
-		} catch (UnsupportedEncodingException e) {
-			throw new WriteFailedException("Could not determine buffer size because of unsupported encoding: " + encoding, e);
+			return getCurrentBuffer().length();
+		} catch (IOException e) {
+			throw new WriteFailedException("Could not determine buffer size because of underlying error", e);
 		}
 	}
 
@@ -197,8 +239,30 @@ public void close() throws IOException {
 	@Override
 	public void flush() throws IOException {
 		if (!transactionActive() && forceSync) {
-			channel.force(false);
+			outputStream.getChannel().force(false);
+		}
+	}
+	
+	@Override
+	public void write(String str, int off, int len) throws IOException {
+		if (!transactionActive()) {
+			if (len > 8192) {
+				// avoid creation of large byte[]
+				BufferedOutputStream bos = new BufferedOutputStream(outputStream);
+				OutputStreamWriter writer = new OutputStreamWriter(bos, this.encoding);
+				writer.write(str, off, len);
+				writer.flush();
+				// do not call #close(), do not close underlying stream
+			}
+			else {
+				byte[] bytes = str.substring(off, len).getBytes(encoding);
+				outputStream.write(bytes);
+			}
+			return;
 		}
+
+		Buffer buffer = getCurrentBuffer();
+		buffer.append(str, off, len);
 	}
 
 	/*
@@ -210,19 +274,22 @@ public void flush() throws IOException {
 	public void write(char[] cbuf, int off, int len) throws IOException {
 
 		if (!transactionActive()) {
-			char [] subArray = new char[len];
-			System.arraycopy(cbuf, off, subArray, 0, len);
-			byte[] bytes = new String(subArray).getBytes(encoding);
-			int length = bytes.length;
-			ByteBuffer bb = ByteBuffer.wrap(bytes);
-			int bytesWritten = channel.write(bb);
-			if(bytesWritten != length) {
-				throw new IOException("Unable to write all data.  Bytes to write: " + len + ".  Bytes written: " + bytesWritten);
+			if (len > 8192) {
+				// avoid creation of large byte[]
+				BufferedOutputStream bos = new BufferedOutputStream(outputStream);
+				OutputStreamWriter writer = new OutputStreamWriter(bos, this.encoding);
+				writer.write(cbuf, off, len);
+				writer.flush();
+				// do not call #close(), do not close underlying stream
+			}
+			else {
+				byte[] bytes = new String(cbuf, off, len).getBytes(encoding);
+				outputStream.write(bytes);
 			}
 			return;
 		}
 
-		StringBuilder buffer = getCurrentBuffer();
+		Buffer buffer = getCurrentBuffer();
 		buffer.append(cbuf, off, len);
 	}
 }