io.milton.common.BufferingOutputStream Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 io.milton.common;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.output.ByteArrayOutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.*;
/**
* An output stream which will buffer data, initially using memory up to
* maxMemorySize, and then overflowing to a temporary file.
*
* To use this class you will write to it, and then close it, and then call
* getInputStream to read the data.
*
* The temporary file, if it was created, will be deleted when the input stream
* is closed.
*
* @author brad
*/
public class BufferingOutputStream extends OutputStream {
private static final Logger log = LoggerFactory.getLogger(BufferingOutputStream.class);
private ByteArrayOutputStream tempMemoryBuffer = new ByteArrayOutputStream();
private final int maxMemorySize;
private File tempFile;
private FileOutputStream fout;
private BufferedOutputStream bufOut;
private Runnable runnable;
private long size;
private boolean closed;
public BufferingOutputStream(int maxMemorySize) {
this.maxMemorySize = maxMemorySize;
}
public InputStream getInputStream() {
if (!closed) {
throw new IllegalStateException("this output stream is not yet closed");
}
if (tempMemoryBuffer == null) {
FileDeletingInputStream fin;
try {
fin = new FileDeletingInputStream(tempFile);
} catch (FileNotFoundException ex) {
throw new RuntimeException(tempFile.getAbsolutePath(), ex);
}
return new BufferedInputStream(fin);
} else {
return new ByteArrayInputStream(tempMemoryBuffer.toByteArray());
}
}
@Override
public void write(byte[] b) throws IOException {
size += b.length;
if (tempMemoryBuffer != null) {
tempMemoryBuffer.write(b);
} else {
bufOut.write(b);
}
checkSize();
}
@Override
public void write(int b) throws IOException {
size++;
if (tempMemoryBuffer != null) {
tempMemoryBuffer.write(b);
} else {
bufOut.write(b);
}
checkSize();
}
@Override
public void write(byte[] b, int off, int len) throws IOException {
size += len;
if (tempMemoryBuffer != null) {
tempMemoryBuffer.write(b, off, len);
} else {
bufOut.write(b, off, len);
}
checkSize();
}
private void checkSize() throws IOException {
if (log.isTraceEnabled()) {
log.trace("checkSize: {}", size);
}
if (tempMemoryBuffer == null) {
return;
}
if (tempMemoryBuffer.size() < maxMemorySize) {
return;
}
tempFile = File.createTempFile("" + System.currentTimeMillis(), ".buffer");
fout = new FileOutputStream(tempFile);
bufOut = new BufferedOutputStream(fout);
bufOut.write(tempMemoryBuffer.toByteArray());
tempMemoryBuffer = null;
}
@Override
public void flush() throws IOException {
if (tempMemoryBuffer != null) {
tempMemoryBuffer.flush();
} else {
bufOut.flush();
fout.flush();
}
}
@Override
public void close() throws IOException {
if (!closed) {
closed = true;
if (tempMemoryBuffer != null) {
tempMemoryBuffer.close();
} else {
bufOut.close();
fout.close();
}
if (runnable != null) {
runnable.run();
}
}
}
/**
* Returns size of the buffer,
* @return size of the buffer.
*/
public long getSize() {
return size;
}
/**
* Returns underlying temporary file if exists.
* @return underlying temporary file if exists.
*/
File getTempFile() {
return tempFile;
}
/**
* Returns temporary in-memory buffer.
* @return temporary in-memory buffer.
*/
ByteArrayOutputStream getTempMemoryBuffer() {
return tempMemoryBuffer;
}
/**
* Sets callback which will be executed on stream close.
* @param r - callback for stream close.
*/
public void setOnClose(Runnable r) {
this.runnable = r;
}
/**
* returns true if the data is completely held in memory
*
* @return
*/
public boolean isCompleteInMemory() {
return tempFile == null;
}
/**
* Gets the data currently held in memory
*
* @return
*/
public byte[] getInMemoryData() {
return this.tempMemoryBuffer.toByteArray();
}
// BM: deleting is taken care of by FileDeletingInputStream
// @Override
// protected void finalize() throws Throwable {
// deleteTempFileIfExists();
// super.finalize();
// }
/**
* If this is called before the inputstream is used, then the inputstream
* will fail to open (because it needs the file!!) So should only use in
* exception handlers
*/
public void deleteTempFileIfExists() {
if (bufOut != null) {
IOUtils.closeQuietly(bufOut);
}
if (fout != null) {
IOUtils.closeQuietly(fout);
}
if (tempFile != null && tempFile.exists()) {
log.error("temporary file was not deleted. Was close called on the inputstream? Will attempt to delete");
if (!tempFile.delete()) {
log.error("Still couldnt delete temporary file: {}", tempFile.getAbsolutePath());
}
}
}
}