package com.github.davidmoten.rx.internal.operators;
import java.io.DataInputStream;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.RandomAccessFile;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.channels.FileChannel.MapMode;
import java.util.concurrent.atomic.AtomicInteger;
import com.github.davidmoten.rx.buffertofile.DataSerializer;
import com.github.davidmoten.util.ByteArrayOutputStreamNoCopyUnsynchronized;
import com.github.davidmoten.util.Preconditions;
public class FileBasedSPSCQueueMemoryMappedReaderWriter<T> {
private volatile RandomAccessFile f;
private volatile FileChannel channel;
private DataInputStream input;
private DataOutputStream output;
private MappedByteBuffer read;
private MappedByteBuffer write;
private final DataSerializer<T> serializer;
private final File file;
private final int fileSize;
private final DataOutput buffer;
// TODO can be passed in to constructor for reuse
private final ByteArrayOutputStreamNoCopyUnsynchronized bytes;
private final AtomicInteger status = new AtomicInteger(WRITTEN_READ);
private final Object markerLock = new Object();
static final int WRITTEN_READ = 0;
static final int WRITTEN_READ_NOT_STARTED = 1;
static final int WRITTEN_READING = 2;
static final int WRITING_NOT_READING = 3;
static final int WRITING_READING = 4;
public FileBasedSPSCQueueMemoryMappedReaderWriter(File file, int fileSize, DataSerializer<T> serializer) {
Preconditions.checkArgument(serializer.size() == 0 || serializer.size() <= fileSize - 2 * MARKER_HEADER_SIZE,
"serializer.size() must be less than or equal to file based queue size - 2");
this.file = file;
this.serializer = serializer;
this.fileSize = fileSize;
this.bytes = new ByteArrayOutputStreamNoCopyUnsynchronized();
this.buffer = new DataOutputStream(bytes);
}
public FileBasedSPSCQueueMemoryMappedReaderWriter<T> openForRead() {
// System.out.println("openForRead " + file);
if (!status.compareAndSet(WRITTEN_READ_NOT_STARTED, WRITTEN_READING))
status.compareAndSet(WRITING_NOT_READING, WRITING_READING);
while (true) {
int st = status.get();
int newStatus;
if (st == WRITTEN_READ_NOT_STARTED)
newStatus = WRITTEN_READING;
else if (st == WRITING_NOT_READING)
newStatus = WRITING_READING;
else
newStatus = st;
if (status.compareAndSet(st, newStatus)) {
checkClose(newStatus);
break;
}
}
try {
if (f == null) {
f = new RandomAccessFile(file, "r");
}
if (channel == null) {
channel = f.getChannel();
}
read = channel.map(MapMode.READ_ONLY, 0, channel.size());
input = new DataInputStream(new MappedByteBufferInputStream(read));
// System.out.println("opened for read " + file.getName());
} catch (IOException e) {
throw new RuntimeException(e);
}
return this;
}
public void closeForRead() {
// System.out.println("closeForRead " + file);
while (true) {
int st = status.get();
int newStatus;
if (st == WRITTEN_READING)
newStatus = WRITTEN_READ;
else
newStatus = st;
if (status.compareAndSet(st, newStatus)) {
checkClose(newStatus);
break;
}
}
}
public FileBasedSPSCQueueMemoryMappedReaderWriter<T> openForWrite() {
// System.out.println("openForWrite " + file);
while (true) {
int st = status.get();
int newStatus;
if (st == WRITTEN_READ)
newStatus = WRITING_NOT_READING;
else
newStatus = st;
if (status.compareAndSet(st, newStatus)) {
checkClose(newStatus);
break;
}
}
try {
if (f == null) {
f = new RandomAccessFile(file, "rw");
}
if (channel == null) {
channel = f.getChannel();
}
write = channel.map(MapMode.READ_WRITE, 0, fileSize);
output = new DataOutputStream(new MappedByteBufferOutputStream(write));
synchronized (markerLock) {
output.write(MARKER_END_OF_QUEUE);
}
// System.out.println("opened for write " + file.getName());
return this;
} catch (IOException e) {
throw new RuntimeException(e);
}
}
public void closeForWrite() {
// System.out.println("closeForWrite " + file);
while (true) {
int st = status.get();
int newStatus;
if (st == WRITING_READING)
newStatus = WRITTEN_READING;
else if (st == WRITING_NOT_READING)
newStatus = WRITTEN_READ_NOT_STARTED;
newStatus = st;
if (status.compareAndSet(st, newStatus)) {
checkClose(newStatus);
break;
}
}
}
private void checkClose(int newStatus) {
// System.out.println("close status = " + newStatus + " for " +
// file.getName());
if (newStatus == WRITTEN_READ) {
try {
channel.close();
channel = null;
read = null;
input = null;
f.close();
f = null;
} catch (IOException e) {
throw new RuntimeException(e);
}
}
}
private static class MappedByteBufferOutputStream extends OutputStream {
private final MappedByteBuffer write;
MappedByteBufferOutputStream(MappedByteBuffer write) {
this.write = write;
}
@Override
public void write(int b) throws IOException {
write.put((byte) b);
}
}
private static class MappedByteBufferInputStream extends InputStream {
private final MappedByteBuffer read;
MappedByteBufferInputStream(MappedByteBuffer read) {
this.read = read;
}
@Override
public int read() throws IOException {
return toUnsignedInteger(read.get());
}
}
private static int toUnsignedInteger(byte b) {
return b & 0x000000FF;
}
private static final EOFRuntimeException EOF = new EOFRuntimeException();
static final class EOFRuntimeException extends RuntimeException {
private static final long serialVersionUID = -6943467453336359472L;
}
// markers must be powers of 2 so that we can detect
// a partial write of the byte (which we will treat as END_OF_QUEUE)
static final byte MARKER_END_OF_QUEUE = 0;
static final byte MARKER_END_OF_FILE = 1;
static final byte MARKER_ITEM_PRESENT = 2;
static final int MARKER_HEADER_SIZE = 1;
static final int UNKNOWN_LENGTH = 0;
public T poll() {
int position = read.position();
byte marker;
synchronized (markerLock) {
marker = read.get();
}
if (marker == MARKER_END_OF_QUEUE) {
read.position(position);
return null;
} else if (marker == MARKER_END_OF_FILE) {
throw EOF;
} else if (marker == MARKER_ITEM_PRESENT) {
try {
T t = serializer.deserialize(input);
if (t == null) {
// this is a trick that we can get away with due to type
// erasure in java as long as the return value of poll() is
// checked using NullSentinel.isNullSentinel(t) (?)
return NullSentinel.instance();
} else {
return t;
}
} catch (IOException e) {
throw new RuntimeException(e);
}
} else {
throw new RuntimeException("unexpected");
}
}
/**
* Returns true if value written to file or false if not enough space
* (writes and end-of-file marker in the fixed-length memory mapped file).
*
* @param t
* value to write to the serialized queue
* @return true if written, false if not enough space
*/
public boolean offer(T t) {
// the current position will be just past the length bytes for this
// item (length bytes will be 0 at the moment)
int serializedLength = serializer.size();
if (serializedLength == UNKNOWN_LENGTH) {
return offerUnknownLength(t);
} else {
return offerKnownLength(t, serializedLength);
}
}
private boolean offerKnownLength(T t, int serializedLength) {
try {
if (notEnoughSpace(serializedLength)) {
markFileAsCompletedAndClose();
return false;
}
int position = write.position();
// serialize the object t to the file
serializer.serialize(output, t);
int length = write.position() - position;
checkLength(serializedLength, length);
updateMarkers(serializedLength);
return true;
} catch (IOException e) {
throw new RuntimeException(e);
}
}
private boolean offerUnknownLength(T t) {
try {
bytes.reset();
// serialize to an in-memory buffer to calculate length
serializer.serialize(buffer, t);
int serializedLength = bytes.size();
if (notEnoughSpace(serializedLength)) {
markFileAsCompletedAndClose();
return false;
} else {
write.put(bytes.toByteArrayNoCopy(), 0, bytes.size());
updateMarkers(serializedLength);
return true;
}
} catch (IOException e) {
throw new RuntimeException(e);
}
}
private void checkLength(int serializedLength, int length) {
if (length > serializedLength) {
throw new IllegalArgumentException(
"serialized length of value being offered to file queue was greater than serializer.size() value (which was non-zero)");
}
}
private void markFileAsCompletedAndClose() {
write.position(write.position() - MARKER_HEADER_SIZE);
synchronized (markerLock) {
write.put(MARKER_END_OF_FILE);
}
closeForWrite();
}
private boolean notEnoughSpace(int serializedLength) {
if (serializedLength > fileSize - 2 * MARKER_HEADER_SIZE)
throw new RuntimeException("serialized length is larger than can fit in one file");
return serializedLength + MARKER_HEADER_SIZE > write.remaining();
}
private void updateMarkers(int serializedLength) throws IOException {
// write the marker for the next item
write.put(MARKER_END_OF_QUEUE);
// remember the position where the next write starts
int newWritePosition = write.position();
// rewind and update the length for the current item
write.position(write.position() - serializedLength - 2 * MARKER_HEADER_SIZE);
// now indicate to the reader that it can read this item
synchronized (markerLock) {
write.put(MARKER_ITEM_PRESENT);
}
// and update the position to the write position for the
// next item
write.position(newWritePosition);
}
public void close() {
try {
f.close();
} catch (IOException e) {
throw new RuntimeException(e);
}
}
}