RecordIterator.java

/*
 * Copyright (c) 2011 LinkedIn, Inc
 * 
 * Licensed 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 com.flaptor.indextank.storage;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.nio.channels.Channels;

import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TTransportException;

import com.flaptor.indextank.rpc.LogRecord;
import com.google.common.collect.AbstractIterator;

public final class RecordIterator extends AbstractIterator<LogRecord> {

    private final TProtocol protocol;
    private final CountingReadTransport transport;
    private long safelyRead = 0;
    private long totalRead = 0;
    private long end = -1;
    private final String description;

    RecordIterator(CountingReadTransport transport, TProtocol protocol, String description) {
        this.transport = transport;
        this.protocol = protocol;
        this.description = description;
    }

    public RecordIterator(CountingReadTransport transport, TProtocol protocol, long end, String description) {
        this(transport, protocol, description);
        this.end = end;
    }

    @Override
    protected LogRecord computeNext() {
        if (end >= 0 && transport.getBytesRead() >= end) {
            if (transport != null) {
                totalRead = transport.getBytesRead();
                transport.close();
            }
            return endOfData();
        }
        LogRecord record = new LogRecord();
        try {
            ((TBinaryProtocol)protocol).setReadLength(10000000);
            record.read(protocol);
            if (transport != null) {
                safelyRead = transport.getBytesRead();
            }
        } catch (TTransportException e) {
            switch (e.getType()) {
                case TTransportException.END_OF_FILE:
                    if (transport != null) {
                        totalRead = transport.getBytesRead();
                        transport.close();
                    }
                    return endOfData();
                case TTransportException.UNKNOWN:
                    if (e.getMessage().startsWith("Cannot read. Remote side has closed")) {
                        if (transport != null) {
                            totalRead = transport.getBytesRead();
                            transport.close();
                        }
                        return endOfData();
                    }
                default:
                    transport.close();
                    throw new RuntimeException("Failed while iterating: " + description, e);
            }
        } catch (TException e) {
            transport.close();
            throw new RuntimeException("Failed while iterating: " + description, e);
        }
        return record;
    }
    
    public long getSafelyRead() {
        return safelyRead;
    }
    
    public long getTotalRead() {
        return totalRead;
    }
    
    public static RecordIterator forFiles(File... files) throws IOException {
        return forFiles(8192, files);
    }
    public static RecordIterator forFiles(int bufferingSize, File... files) throws IOException {
        return forFiles(-1, -1, bufferingSize, files);
    }
    public static RecordIterator forFiles(long start, long end, int bufferingSize, File... files) throws IOException {
        FileInputStream in; 
        int i = 0;
        long fileLength = 0;
        while (true) {
            try {
                fileLength = files[i].length();
                in = new FileInputStream(files[i++]);
                break;
            } catch (FileNotFoundException e) {
                if (i == files.length) {
                    throw e;
                }
            }
        }
        InputStream is = start >= 0 ? Channels.newInputStream(in.getChannel().position(start)) : in;
        BufferedInputStream bis = new BufferedInputStream(is, bufferingSize);
        CountingReadTransport transport = new CountingReadTransport(bis);
        TProtocol protocol = new TBinaryProtocol.Factory().getProtocol(transport);
        /*if (fileLength < (long)Integer.MAX_VALUE) {
            ((TBinaryProtocol)protocol).setReadLength((int) fileLength);
        }*/
        long length = -1;
        if (end >= 0) {
            if (start >= 0) {
                length = end - start;
            } else {
                length = end;
            }
        }
        return new RecordIterator(transport, protocol, length, files[i-1].getAbsolutePath());
    }
    
    public String getDescription() {
        return description;
    }
    
}