/*
* 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 org.apache.cassandra.io.util;
import java.io.Closeable;
import java.io.File;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.RandomAccessFile;
import java.util.concurrent.atomic.AtomicBoolean;
import static org.apache.cassandra.utils.Throwables.maybeFail;
import static org.apache.cassandra.utils.Throwables.merge;
/**
* Adds mark/reset functionality to another input stream by caching read bytes to a memory buffer and
* spilling to disk if necessary.
*
* When the stream is marked via {@link #mark()} or {@link #mark(int)}, up to
* <code>maxMemBufferSize</code> will be cached in memory (heap). If more than
* <code>maxMemBufferSize</code> bytes are read while the stream is marked, the
* following bytes are cached on the <code>spillFile</code> for up to <code>maxDiskBufferSize</code>.
*
* Please note that successive calls to {@link #mark()} and {@link #reset()} will write
* sequentially to the same <code>spillFile</code> until <code>maxDiskBufferSize</code> is reached.
* At this point, if less than <code>maxDiskBufferSize</code> bytes are currently cached on the
* <code>spillFile</code>, the remaining bytes are written to the beginning of the file,
* treating the <code>spillFile</code> as a circular buffer.
*
* If more than <code>maxMemBufferSize + maxDiskBufferSize</code> are cached while the stream is marked,
* the following {@link #reset()} invocation will throw a {@link IllegalStateException}.
*
*/
public class RewindableDataInputStreamPlus extends FilterInputStream implements RewindableDataInput, Closeable
{
private boolean marked = false;
private boolean exhausted = false;
private AtomicBoolean closed = new AtomicBoolean(false);
protected int memAvailable = 0;
protected int diskTailAvailable = 0;
protected int diskHeadAvailable = 0;
private final File spillFile;
private final int initialMemBufferSize;
private final int maxMemBufferSize;
private final int maxDiskBufferSize;
private volatile byte memBuffer[];
private int memBufferSize;
private RandomAccessFile spillBuffer;
private final DataInputPlus dataReader;
public RewindableDataInputStreamPlus(InputStream in, int initialMemBufferSize, int maxMemBufferSize,
File spillFile, int maxDiskBufferSize)
{
super(in);
dataReader = new DataInputStreamPlus(this);
this.initialMemBufferSize = initialMemBufferSize;
this.maxMemBufferSize = maxMemBufferSize;
this.spillFile = spillFile;
this.maxDiskBufferSize = maxDiskBufferSize;
}
/* RewindableDataInput methods */
/**
* Marks the current position of a stream to return to this position later via the {@link #reset(DataPosition)} method.
* @return An empty @link{DataPosition} object
*/
public DataPosition mark()
{
mark(0);
return new RewindableDataInputPlusMark();
}
/**
* Rewinds to the previously marked position via the {@link #mark()} method.
* @param mark it's not possible to return to a custom position, so this parameter is ignored.
* @throws IOException if an error ocurs while resetting
*/
public void reset(DataPosition mark) throws IOException
{
reset();
}
public long bytesPastMark(DataPosition mark)
{
return maxMemBufferSize - memAvailable + (diskTailAvailable == -1? 0 : maxDiskBufferSize - diskHeadAvailable - diskTailAvailable);
}
protected static class RewindableDataInputPlusMark implements DataPosition
{
}
/* InputStream methods */
public boolean markSupported()
{
return true;
}
/**
* Marks the current position of a stream to return to this position
* later via the {@link #reset()} method.
* @param readlimit the maximum amount of bytes to cache
*/
public synchronized void mark(int readlimit)
{
if (marked)
throw new IllegalStateException("Cannot mark already marked stream.");
if (memAvailable > 0 || diskHeadAvailable > 0 || diskTailAvailable > 0)
throw new IllegalStateException("Can only mark stream after reading previously marked data.");
marked = true;
memAvailable = maxMemBufferSize;
diskHeadAvailable = -1;
diskTailAvailable = -1;
}
public synchronized void reset() throws IOException
{
if (!marked)
throw new IOException("Must call mark() before calling reset().");
if (exhausted)
throw new IOException(String.format("Read more than capacity: %d bytes.", maxMemBufferSize + maxDiskBufferSize));
memAvailable = maxMemBufferSize - memAvailable;
memBufferSize = memAvailable;
if (diskTailAvailable == -1)
{
diskHeadAvailable = 0;
diskTailAvailable = 0;
}
else
{
int initialPos = diskTailAvailable > 0 ? 0 : (int)getIfNotClosed(spillBuffer).getFilePointer();
int diskMarkpos = initialPos + diskHeadAvailable;
getIfNotClosed(spillBuffer).seek(diskMarkpos);
diskHeadAvailable = diskMarkpos - diskHeadAvailable;
diskTailAvailable = (maxDiskBufferSize - diskTailAvailable) - diskMarkpos;
}
marked = false;
}
public int available() throws IOException
{
return super.available() + (marked? 0 : memAvailable + diskHeadAvailable + diskTailAvailable);
}
public int read() throws IOException
{
int read = readOne();
if (read == -1)
return read;
if (marked)
{
//mark exhausted
if (isExhausted(1))
{
exhausted = true;
return read;
}
writeOne(read);
}
return read;
}
public int read(byte[] b, int off, int len) throws IOException
{
int readBytes = readMulti(b, off, len);
if (readBytes == -1)
return readBytes;
if (marked)
{
//check we have space on buffer
if (isExhausted(readBytes))
{
exhausted = true;
return readBytes;
}
writeMulti(b, off, readBytes);
}
return readBytes;
}
private void maybeCreateDiskBuffer() throws IOException
{
if (spillBuffer == null)
{
if (!spillFile.getParentFile().exists())
spillFile.getParentFile().mkdirs();
spillFile.createNewFile();
this.spillBuffer = new RandomAccessFile(spillFile, "rw");
}
}
private int readOne() throws IOException
{
if (!marked)
{
if (memAvailable > 0)
{
int pos = memBufferSize - memAvailable;
memAvailable--;
return getIfNotClosed(memBuffer)[pos] & 0xff;
}
if (diskTailAvailable > 0 || diskHeadAvailable > 0)
{
int read = getIfNotClosed(spillBuffer).read();
if (diskTailAvailable > 0)
diskTailAvailable--;
else if (diskHeadAvailable > 0)
diskHeadAvailable++;
if (diskTailAvailable == 0)
spillBuffer.seek(0);
return read;
}
}
return getIfNotClosed(in).read();
}
private boolean isExhausted(int readBytes)
{
return exhausted || readBytes > memAvailable + (long)(diskTailAvailable == -1? maxDiskBufferSize : diskTailAvailable + diskHeadAvailable);
}
private int readMulti(byte[] b, int off, int len) throws IOException
{
int readBytes = 0;
if (!marked)
{
if (memAvailable > 0)
{
readBytes += memAvailable < len ? memAvailable : len;
int pos = memBufferSize - memAvailable;
System.arraycopy(memBuffer, pos, b, off, readBytes);
memAvailable -= readBytes;
off += readBytes;
len -= readBytes;
}
if (len > 0 && diskTailAvailable > 0)
{
int readFromTail = diskTailAvailable < len? diskTailAvailable : len;
readFromTail = getIfNotClosed(spillBuffer).read(b, off, readFromTail);
readBytes += readFromTail;
diskTailAvailable -= readFromTail;
off += readFromTail;
len -= readFromTail;
if (diskTailAvailable == 0)
spillBuffer.seek(0);
}
if (len > 0 && diskHeadAvailable > 0)
{
int readFromHead = diskHeadAvailable < len? diskHeadAvailable : len;
readFromHead = getIfNotClosed(spillBuffer).read(b, off, readFromHead);
readBytes += readFromHead;
diskHeadAvailable -= readFromHead;
off += readFromHead;
len -= readFromHead;
}
}
if (len > 0)
readBytes += getIfNotClosed(in).read(b, off, len);
return readBytes;
}
private void writeMulti(byte[] b, int off, int len) throws IOException
{
if (memAvailable > 0)
{
if (memBuffer == null)
memBuffer = new byte[initialMemBufferSize];
int pos = maxMemBufferSize - memAvailable;
int memWritten = memAvailable < len? memAvailable : len;
if (pos + memWritten >= getIfNotClosed(memBuffer).length)
growMemBuffer(pos, memWritten);
System.arraycopy(b, off, memBuffer, pos, memWritten);
off += memWritten;
len -= memWritten;
memAvailable -= memWritten;
}
if (len > 0)
{
if (diskTailAvailable == -1)
{
maybeCreateDiskBuffer();
diskHeadAvailable = (int)spillBuffer.getFilePointer();
diskTailAvailable = maxDiskBufferSize - diskHeadAvailable;
}
if (len > 0 && diskTailAvailable > 0)
{
int diskTailWritten = diskTailAvailable < len? diskTailAvailable : len;
getIfNotClosed(spillBuffer).write(b, off, diskTailWritten);
off += diskTailWritten;
len -= diskTailWritten;
diskTailAvailable -= diskTailWritten;
if (diskTailAvailable == 0)
spillBuffer.seek(0);
}
if (len > 0 && diskTailAvailable > 0)
{
int diskHeadWritten = diskHeadAvailable < len? diskHeadAvailable : len;
getIfNotClosed(spillBuffer).write(b, off, diskHeadWritten);
}
}
}
private void writeOne(int value) throws IOException
{
if (memAvailable > 0)
{
if (memBuffer == null)
memBuffer = new byte[initialMemBufferSize];
int pos = maxMemBufferSize - memAvailable;
if (pos == getIfNotClosed(memBuffer).length)
growMemBuffer(pos, 1);
getIfNotClosed(memBuffer)[pos] = (byte)value;
memAvailable--;
return;
}
if (diskTailAvailable == -1)
{
maybeCreateDiskBuffer();
diskHeadAvailable = (int)spillBuffer.getFilePointer();
diskTailAvailable = maxDiskBufferSize - diskHeadAvailable;
}
if (diskTailAvailable > 0 || diskHeadAvailable > 0)
{
getIfNotClosed(spillBuffer).write(value);
if (diskTailAvailable > 0)
diskTailAvailable--;
else if (diskHeadAvailable > 0)
diskHeadAvailable--;
if (diskTailAvailable == 0)
spillBuffer.seek(0);
return;
}
}
public int read(byte[] b) throws IOException
{
return read(b, 0, b.length);
}
private void growMemBuffer(int pos, int writeSize)
{
int newSize = Math.min(2 * (pos + writeSize), maxMemBufferSize);
byte newBuffer[] = new byte[newSize];
System.arraycopy(memBuffer, 0, newBuffer, 0, pos);
memBuffer = newBuffer;
}
public long skip(long n) throws IOException
{
long skipped = 0;
if (marked)
{
//if marked, we need to cache skipped bytes
while (n-- > 0 && read() != -1)
{
skipped++;
}
return skipped;
}
if (memAvailable > 0)
{
skipped += memAvailable < n ? memAvailable : n;
memAvailable -= skipped;
n -= skipped;
}
if (n > 0 && diskTailAvailable > 0)
{
int skipFromTail = diskTailAvailable < n? diskTailAvailable : (int)n;
getIfNotClosed(spillBuffer).skipBytes(skipFromTail);
diskTailAvailable -= skipFromTail;
skipped += skipFromTail;
n -= skipFromTail;
if (diskTailAvailable == 0)
spillBuffer.seek(0);
}
if (n > 0 && diskHeadAvailable > 0)
{
int skipFromHead = diskHeadAvailable < n? diskHeadAvailable : (int)n;
getIfNotClosed(spillBuffer).skipBytes(skipFromHead);
diskHeadAvailable -= skipFromHead;
skipped += skipFromHead;
n -= skipFromHead;
}
if (n > 0)
skipped += getIfNotClosed(in).skip(n);
return skipped;
}
private <T> T getIfNotClosed(T in) throws IOException
{
if (closed.get())
throw new IOException("Stream closed");
return in;
}
public void close() throws IOException
{
close(true);
}
public void close(boolean closeUnderlying) throws IOException
{
if (closed.compareAndSet(false, true))
{
Throwable fail = null;
if (closeUnderlying)
{
try
{
super.close();
}
catch (IOException e)
{
fail = merge(fail, e);
}
}
try
{
if (spillBuffer != null)
{
this.spillBuffer.close();
this.spillBuffer = null;
}
} catch (IOException e)
{
fail = merge(fail, e);
}
try
{
if (spillFile.exists())
{
spillFile.delete();
}
}
catch (Throwable e)
{
fail = merge(fail, e);
}
maybeFail(fail, IOException.class);
}
}
/* DataInputPlus methods */
public void readFully(byte[] b) throws IOException
{
dataReader.readFully(b);
}
public void readFully(byte[] b, int off, int len) throws IOException
{
dataReader.readFully(b, off, len);
}
public int skipBytes(int n) throws IOException
{
return dataReader.skipBytes(n);
}
public boolean readBoolean() throws IOException
{
return dataReader.readBoolean();
}
public byte readByte() throws IOException
{
return dataReader.readByte();
}
public int readUnsignedByte() throws IOException
{
return dataReader.readUnsignedByte();
}
public short readShort() throws IOException
{
return dataReader.readShort();
}
public int readUnsignedShort() throws IOException
{
return dataReader.readUnsignedShort();
}
public char readChar() throws IOException
{
return dataReader.readChar();
}
public int readInt() throws IOException
{
return dataReader.readInt();
}
public long readLong() throws IOException
{
return dataReader.readLong();
}
public float readFloat() throws IOException
{
return dataReader.readFloat();
}
public double readDouble() throws IOException
{
return dataReader.readDouble();
}
public String readLine() throws IOException
{
return dataReader.readLine();
}
public String readUTF() throws IOException
{
return dataReader.readUTF();
}
}