package org.apache.cassandra.streaming;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Collection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.ning.compress.lzf.LZFOutputStream;
import org.apache.cassandra.io.sstable.Component;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.io.util.DataIntegrityMetadata;
import org.apache.cassandra.io.util.DataIntegrityMetadata.ChecksumValidator;
import org.apache.cassandra.io.util.DataOutputStreamPlus;
import org.apache.cassandra.io.util.RandomAccessReader;
import org.apache.cassandra.streaming.StreamManager.StreamRateLimiter;
import org.apache.cassandra.utils.FBUtilities;
import org.apache.cassandra.utils.Pair;
public class StreamWriter
{
private static final int DEFAULT_CHUNK_SIZE = 64 * 1024;
private static final Logger logger = LoggerFactory.getLogger(StreamWriter.class);
protected final SSTableReader sstable;
protected final Collection<Pair<Long, Long>> sections;
protected final StreamRateLimiter limiter;
protected final StreamSession session;
private OutputStream compressedOutput;
private byte[] transferBuffer;
public StreamWriter(SSTableReader sstable, Collection<Pair<Long, Long>> sections, StreamSession session)
{
this.session = session;
this.sstable = sstable;
this.sections = sections;
this.limiter = StreamManager.getRateLimiter(session.peer);
}
public void write(DataOutputStreamPlus output) throws IOException
{
long totalSize = totalSize();
logger.debug("[Stream #{}] Start streaming file {} to {}, repairedAt = {}, totalSize = {}", session.planId(),
sstable.getFilename(), session.peer, sstable.getSSTableMetadata().repairedAt, totalSize);
try(RandomAccessReader file = sstable.openDataReader();
ChecksumValidator validator = new File(sstable.descriptor.filenameFor(Component.CRC)).exists()
? DataIntegrityMetadata.checksumValidator(sstable.descriptor)
: null;)
{
transferBuffer = validator == null ? new byte[DEFAULT_CHUNK_SIZE] : new byte[validator.chunkSize];
compressedOutput = new LZFOutputStream(output);
long progress = 0L;
for (Pair<Long, Long> section : sections)
{
long start = validator == null ? section.left : validator.chunkStart(section.left);
int readOffset = (int) (section.left - start);
file.seek(start);
if (validator != null)
validator.seek(start);
long length = section.right - start;
long bytesRead = 0;
while (bytesRead < length)
{
long lastBytesRead = write(file, validator, readOffset, length, bytesRead);
bytesRead += lastBytesRead;
progress += (lastBytesRead - readOffset);
session.progress(sstable.descriptor.filenameFor(Component.DATA), ProgressInfo.Direction.OUT, progress, totalSize);
readOffset = 0;
}
compressedOutput.flush();
}
logger.debug("[Stream #{}] Finished streaming file {} to {}, bytesTransferred = {}, totalSize = {}",
session.planId(), sstable.getFilename(), session.peer, FBUtilities.prettyPrintMemory(progress), FBUtilities.prettyPrintMemory(totalSize));
}
}
protected long totalSize()
{
long size = 0;
for (Pair<Long, Long> section : sections)
size += section.right - section.left;
return size;
}
protected long write(RandomAccessReader reader, ChecksumValidator validator, int start, long length, long bytesTransferred) throws IOException
{
int toTransfer = (int) Math.min(transferBuffer.length, length - bytesTransferred);
int minReadable = (int) Math.min(transferBuffer.length, reader.length() - reader.getFilePointer());
reader.readFully(transferBuffer, 0, minReadable);
if (validator != null)
validator.validate(transferBuffer, 0, minReadable);
limiter.acquire(toTransfer - start);
compressedOutput.write(transferBuffer, start, (toTransfer - start));
return toTransfer;
}
}