<commit_before>package voldemort.server.protocol.admin;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.channels.Channels;
import java.nio.channels.FileChannel;
import java.nio.channels.WritableByteChannel;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;

import org.apache.log4j.Logger;

import voldemort.VoldemortException;
import voldemort.client.protocol.pb.ProtoUtils;
import voldemort.client.protocol.pb.VAdminProto;
import voldemort.server.StoreRepository;
import voldemort.server.VoldemortConfig;
import voldemort.server.protocol.StreamRequestHandler;
import voldemort.store.StoreDefinition;
import voldemort.store.metadata.MetadataStore;
import voldemort.store.readonly.ReadOnlyStorageConfiguration;
import voldemort.store.readonly.ReadOnlyStorageEngine;
import voldemort.store.stats.StreamStats;
import voldemort.utils.EventThrottler;
import voldemort.utils.Pair;

import com.google.common.collect.Maps;

public class FetchPartitionFileStreamRequestHandler implements StreamRequestHandler {

    private final VAdminProto.FetchPartitionFilesRequest request;

    private final EventThrottler throttler;

    private final File storeDir;

    private final Logger logger = Logger.getLogger(getClass());

    private final long blockSize;

    private final StreamStats stats;

    private final StreamStats.Handle handle;

    private final ReadOnlyStorageEngine storageEngine;

    private final HashMap<Integer, List<Integer>> replicaToPartitionList;

    protected FetchPartitionFileStreamRequestHandler(VAdminProto.FetchPartitionFilesRequest request,
                                                     MetadataStore metadataStore,
                                                     VoldemortConfig voldemortConfig,
                                                     StoreRepository storeRepository,
                                                     StreamStats stats) {
        this.request = request;

        StoreDefinition storeDef = metadataStore.getStoreDef(request.getStore());
        boolean isReadOnly = storeDef.getType().compareTo(ReadOnlyStorageConfiguration.TYPE_NAME) == 0;
        if(!isReadOnly) {
            throw new VoldemortException("Should be fetching partition files only for read-only stores");
        }

        HashMap<Integer, List<Integer>> localReplicaToPartitionList = ProtoUtils.decodePartitionTuple(request.getReplicaToPartitionList());

        // Filter the replica to partition mapping so as to include only till
        // the number of replicas
        this.replicaToPartitionList = Maps.newHashMap();
        for(int replicaType = 0; replicaType < storeDef.getReplicationFactor(); replicaType++) {
            if(localReplicaToPartitionList.containsKey(replicaType)) {
                this.replicaToPartitionList.put(replicaType,
                                                localReplicaToPartitionList.get(replicaType));
            }
        }
        this.storageEngine = AdminServiceRequestHandler.getReadOnlyStorageEngine(metadataStore,
                                                                                 storeRepository,
                                                                                 request.getStore());

        this.blockSize = voldemortConfig.getAllProps()
                                        .getLong("partition.buffer.size.bytes",
                                                 voldemortConfig.getAdminSocketBufferSize());
        this.storeDir = new File(storageEngine.getCurrentDirPath());
        this.throttler = new EventThrottler(voldemortConfig.getStreamMaxReadBytesPerSec());
        this.stats = stats;
        this.handle = stats.makeHandle(StreamStats.Operation.FETCH_FILE, replicaToPartitionList);
    }

    public StreamRequestDirection getDirection() {
        return StreamRequestDirection.WRITING;
    }

    public final void close(DataOutputStream outputStream) throws IOException {
        ProtoUtils.writeEndOfStream(outputStream);
    }

    public final void handleError(DataOutputStream outputStream, VoldemortException e)
            throws IOException {
        logger.error("handleFetchPartitionFilesEntries failed for request(" + request.toString()
                     + ")", e);
    }

    public StreamRequestHandlerState handleRequest(DataInputStream inputStream,
                                                   DataOutputStream outputStream)
            throws IOException {
        HashMap<Object, Integer> bucketToNumChunks = storageEngine.getChunkedFileSet()
                                                                  .getChunkIdToNumChunks();

        for(Entry<Integer, List<Integer>> entry: replicaToPartitionList.entrySet()) {

            int replicaType = entry.getKey();

            // Go over every partition for this replica type
            for(int partitionId: entry.getValue()) {

                // Check if this bucket exists
                if(!bucketToNumChunks.containsKey(Pair.create(partitionId, replicaType))) {
                    throw new VoldemortException("Bucket [ partition = " + partitionId
                                                 + ", replica = " + replicaType
                                                 + " ] does not exist for store "
                                                 + request.getStore());
                }

                // Get number of chunks for this bucket
                int numChunks = bucketToNumChunks.get(Pair.create(partitionId, replicaType));

                for(int chunkId = 0; chunkId < numChunks; chunkId++) {
                    String fileName = Integer.toString(partitionId) + "_"
                                      + Integer.toString(replicaType) + "_"
                                      + Integer.toString(chunkId);
                    File index = new File(this.storeDir, fileName + ".index");
                    File data = new File(this.storeDir, fileName + ".data");

                    // Both files in chunk exist, start streaming...
                    logger.info("Streaming " + data.getAbsolutePath());
                    streamFile(data, outputStream);
                    logger.info("Completed streaming " + data.getAbsolutePath());
                    logger.info("Streaming " + index.getAbsolutePath());
                    streamFile(index, outputStream);
                    logger.info("Completed streaming " + index.getAbsolutePath());
                    handle.incrementEntriesScanned();
                }

            }
        }
        stats.closeHandle(handle);
        return StreamRequestHandlerState.COMPLETE;
    }

    void streamFile(File fileToStream, DataOutputStream stream) throws IOException {
        FileChannel dataChannel = new FileInputStream(fileToStream).getChannel();
        try {
            VAdminProto.FileEntry response = VAdminProto.FileEntry.newBuilder()
                                                                  .setFileName(fileToStream.getName())
                                                                  .setFileSizeBytes(dataChannel.size())
                                                                  .build();
            // Write header
            ProtoUtils.writeMessage(stream, response);
            throttler.maybeThrottle(response.getSerializedSize());

            // Write rest of file
            WritableByteChannel channelOut = Channels.newChannel(stream);

            // Send chunks to help with throttling
            boolean completedFile = false;
            long chunkSize = 0;
            for(long chunkStart = 0; chunkStart < dataChannel.size() && !completedFile; chunkStart += blockSize) {
                if(dataChannel.size() - chunkStart < blockSize) {
                    chunkSize = dataChannel.size() - chunkStart;
                    completedFile = true;
                } else {
                    chunkSize = blockSize;
                }
                dataChannel.transferTo(chunkStart, chunkSize, channelOut);
                throttler.maybeThrottle((int) chunkSize);
            }
        } finally {
            if(dataChannel != null) {
                dataChannel.close();
            }
        }
    }
}<commit_msg>reslove conflicts after merging with Roshan's new code on rebalance<commit_after>package voldemort.server.protocol.admin;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.nio.channels.Channels;
import java.nio.channels.FileChannel;
import java.nio.channels.WritableByteChannel;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;

import org.apache.log4j.Logger;

import voldemort.VoldemortException;
import voldemort.client.protocol.pb.ProtoUtils;
import voldemort.client.protocol.pb.VAdminProto;
import voldemort.server.StoreRepository;
import voldemort.server.VoldemortConfig;
import voldemort.server.protocol.StreamRequestHandler;
import voldemort.store.StoreDefinition;
import voldemort.store.metadata.MetadataStore;
import voldemort.store.readonly.ReadOnlyStorageConfiguration;
import voldemort.store.readonly.ReadOnlyStorageEngine;
import voldemort.store.stats.StreamStats;
import voldemort.utils.EventThrottler;
import voldemort.utils.Pair;
import voldemort.utils.RebalanceUtils;

import com.google.common.collect.Lists;

public class FetchPartitionFileStreamRequestHandler implements StreamRequestHandler {

    private final VAdminProto.FetchPartitionFilesRequest request;

    private final EventThrottler throttler;

    private final File storeDir;

    private final Logger logger = Logger.getLogger(getClass());

    private final long blockSize;

    private final StreamStats stats;

    private final StreamStats.Handle handle;

    private final Iterator<Pair<Integer, Integer>> partitionIterator;

    private FetchStatus fetchStatus;

    private int currentChunkId;

    private Pair<Integer, Integer> currentPair;

    private File indexFile;

    private File dataFile;

    private ChunkedFileWriter chunkedFileWriter;

    private final List<Pair<Integer, Integer>> replicaToPartitionList;

    protected FetchPartitionFileStreamRequestHandler(VAdminProto.FetchPartitionFilesRequest request,
                                                     MetadataStore metadataStore,
                                                     VoldemortConfig voldemortConfig,
                                                     StoreRepository storeRepository,
                                                     StreamStats stats) {
        this.request = request;
        StoreDefinition storeDef = metadataStore.getStoreDef(request.getStore());
        boolean isReadOnly = storeDef.getType().compareTo(ReadOnlyStorageConfiguration.TYPE_NAME) == 0;
        if(!isReadOnly) {
            throw new VoldemortException("Should be fetching partition files only for read-only stores");
        }

        ReadOnlyStorageEngine storageEngine = AdminServiceRequestHandler.getReadOnlyStorageEngine(metadataStore,
                                                                                                  storeRepository,
                                                                                                  request.getStore());

        HashMap<Integer, List<Integer>> localReplicaToPartitionList = ProtoUtils.decodePartitionTuple(request.getReplicaToPartitionList());

        // Filter the replica to partition mapping so as to include only till
        // the number of replicas
        this.replicaToPartitionList = Lists.newArrayList();
        for(int replicaType = 0; replicaType < storeDef.getReplicationFactor(); replicaType++) {
            if(localReplicaToPartitionList.containsKey(replicaType)) {
                List<Integer> partitionList = localReplicaToPartitionList.get(replicaType);
                for(Iterator<Integer> it = partitionList.iterator(); it.hasNext();) {
                    this.replicaToPartitionList.add(new Pair<Integer, Integer>(replicaType,
                                                                               it.next()));
                }
            }
        }

        this.blockSize = voldemortConfig.getAllProps()
                                        .getLong("partition.buffer.size.bytes",
                                                 voldemortConfig.getAdminSocketBufferSize());
        this.storeDir = new File(storageEngine.getCurrentDirPath());
        this.throttler = new EventThrottler(voldemortConfig.getStreamMaxReadBytesPerSec());
        this.stats = stats;
        this.handle = stats.makeHandle(StreamStats.Operation.FETCH_FILE,
                                       RebalanceUtils.flattenPartitionTuples(new HashSet<Pair<Integer, Integer>>(replicaToPartitionList)));
        this.partitionIterator = Collections.unmodifiableList(replicaToPartitionList).iterator();
        this.fetchStatus = FetchStatus.next_partition;
        this.currentChunkId = 0;
        this.indexFile = null;
        this.dataFile = null;
        this.chunkedFileWriter = null;
    }

    public StreamRequestDirection getDirection() {
        return StreamRequestDirection.WRITING;
    }

    public final void close(DataOutputStream outputStream) throws IOException {
        ProtoUtils.writeEndOfStream(outputStream);
    }

    public final void handleError(DataOutputStream outputStream, VoldemortException e)
            throws IOException {
        logger.error("handleFetchPartitionFilesEntries failed for request(" + request.toString()
                     + ")", e);
    }

    public StreamRequestHandlerState handleRequest(DataInputStream inputStream,
                                                   DataOutputStream outputStream)
            throws IOException {
        StreamRequestHandlerState handlerState = StreamRequestHandlerState.WRITING;

        switch(fetchStatus) {
            case next_partition:
                handlerState = handleNextPartition();
                break;
            case send_data_file:
                handleSendDataFile(outputStream);
                break;
            case send_index_file:
                handleSendIndexFile();
                break;
            default:
                throw new VoldemortException("Invalid fetch status: " + fetchStatus);
        }

        return handlerState;
    }

    private void handleSendIndexFile() throws IOException {
        if(0 == chunkedFileWriter.streamFile()) {
            // we are done with the index file, move to next chunk
            logger.info("Completed streaming " + indexFile.getAbsolutePath());
            this.chunkedFileWriter.close();
            currentChunkId++;
            dataFile = indexFile = null;
            handle.incrementEntriesScanned();
            fetchStatus = FetchStatus.send_data_file;
        } else {
            // index file not done yet, keep sending
            fetchStatus = FetchStatus.send_index_file;
        }
    }

    private void handleSendDataFile(DataOutputStream outputStream) throws IOException {
        if(null == dataFile && null == indexFile) {
            // first time enter here, create files based on partition and
            // chunk ids
            String fileName = currentPair.getSecond().toString() + "_"
                              + currentPair.getFirst().toString() + "_"
                              + Integer.toString(currentChunkId);
            dataFile = new File(this.storeDir, fileName + ".data");
            indexFile = new File(this.storeDir, fileName + ".index");
            validateFiles();
            if(isPartitionFinished()) {
                // finished with this partition, move to next one
                fetchStatus = FetchStatus.next_partition;
                return;
            }

            // create a new writer for data file
            this.chunkedFileWriter = new ChunkedFileWriter(dataFile, outputStream);
            logger.info("Streaming " + dataFile.getAbsolutePath());
            this.chunkedFileWriter.writeHeader();
        }

        if(0 == chunkedFileWriter.streamFile()) {
            // we are done with the data file, move to index file
            logger.info("Completed streaming " + dataFile.getAbsolutePath());
            this.chunkedFileWriter.close();
            this.chunkedFileWriter = new ChunkedFileWriter(indexFile, outputStream);
            logger.info("Streaming " + indexFile.getAbsolutePath());
            this.chunkedFileWriter.writeHeader();
            fetchStatus = FetchStatus.send_index_file;
        } else {
            // data file not done yet, keep sending
            fetchStatus = FetchStatus.send_data_file;
        }
    }

    private StreamRequestHandlerState handleNextPartition() {

        StreamRequestHandlerState handlerState = StreamRequestHandlerState.WRITING;

        if(partitionIterator.hasNext()) {
            // start a new partition
            currentPair = partitionIterator.next();
            currentChunkId = 0;
            dataFile = indexFile = null;
            fetchStatus = FetchStatus.send_data_file;
        } else {
            // we are done since we have gone through the entire
            // partition list
            logger.info("Finished streaming files for partitions: " + replicaToPartitionList);
            stats.closeHandle(handle);
            handlerState = StreamRequestHandlerState.COMPLETE;
        }

        return handlerState;
    }

    private void validateFiles() {
        if(!indexFile.exists() && !dataFile.exists() && 0 == currentChunkId) {
            throw new VoldemortException("Could not find any data for "
                                         + currentPair.getSecond().toString() + "_"
                                         + currentPair.getFirst().toString() + "_" + currentChunkId);
        }

        if(indexFile.exists() ^ dataFile.exists()) {
            throw new VoldemortException("One of the following does not exist: "
                                         + indexFile.toString() + " and " + dataFile.toString()
                                         + ".");
        }
    }

    private boolean isPartitionFinished() {
        boolean finished = false;
        if(!indexFile.exists() && !dataFile.exists() && currentChunkId > 0) {
            finished = true;
        }
        return finished;
    }

    enum FetchStatus {
        next_partition,
        send_data_file,
        send_index_file
    }

    class ChunkedFileWriter {

        private final File fileToWrite;
        private final DataOutputStream outStream;
        private final FileChannel dataChannel;
        private final WritableByteChannel outChannel;
        private long currentPos;

        ChunkedFileWriter(File fileToWrite, DataOutputStream stream) throws FileNotFoundException {
            this.fileToWrite = fileToWrite;
            this.outStream = stream;
            this.dataChannel = new FileInputStream(fileToWrite).getChannel();
            this.outChannel = Channels.newChannel(outStream);
            this.currentPos = 0;
        }

        public void close() throws IOException {
            dataChannel.close();
            outChannel.close();
        }

        void writeHeader() throws IOException {
            VAdminProto.FileEntry response = VAdminProto.FileEntry.newBuilder()
                                                                  .setFileName(fileToWrite.getName())
                                                                  .setFileSizeBytes(dataChannel.size())
                                                                  .build();

            ProtoUtils.writeMessage(outStream, response);
            throttler.maybeThrottle(response.getSerializedSize());
        }

        // this function returns the number of bytes left, 0 means done
        long streamFile() throws IOException {
            long bytesRemaining = dataChannel.size() - currentPos;
            if(0 < bytesRemaining) {
                long bytesToWrite = Math.min(bytesRemaining, blockSize);
                long bytesWritten = dataChannel.transferTo(currentPos, bytesToWrite, outChannel);
                currentPos += bytesWritten;
                logger.info(bytesWritten + " of bytes sent");
                throttler.maybeThrottle((int) bytesWritten);
            }
            bytesRemaining = dataChannel.size() - currentPos;
            return bytesRemaining;
        }
    }
}
{'fraction_non_alphanumeric': 0.07666650553487697, 'fraction_numerical': 0.0008701116643302558, 'mean_word_length': 1.9788336933045356, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 76, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5715049', 'n_tokens_mistral': 4781, 'n_tokens_neox': 4653, 'n_words': 1024}