/* * 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; } }