Java tutorial
/* * 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.addthis.hydra.data.io; import java.io.BufferedInputStream; import java.io.BufferedOutputStream; import java.io.DataInputStream; import java.io.DataOutputStream; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.Comparator; import java.util.Iterator; import java.util.List; import java.util.ListIterator; import java.util.PriorityQueue; import com.addthis.basis.util.LessBytes; import com.addthis.basis.util.LessFiles; import com.addthis.basis.util.MemoryCounter; import com.addthis.basis.util.Parameter; import com.google.common.base.Throwables; import org.apache.commons.lang3.tuple.Pair; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * An implementation of disk-backed list tuned for good sorting performance. * Incoming values are split into a number of chunks. Each chunk is small enough * to fit entirely within memory. * * @param <K> a codable object type */ public class DiskBackedList2<K> implements List<K> { public static interface ItemCodec<K> { public K decode(byte[] row) throws IOException; public byte[] encode(K row) throws IOException; } private static final Logger log = LoggerFactory.getLogger(DiskBackedList2.class); private ItemCodec<K> codec; private List<DBLChunk> chunks; private DBLChunk currentChunk; private int totalItems; private final long maxChunkSizeBytes; private final long maxTotalSizeBytes = Parameter.longValue("max.total.query.size.bytes", 10 * 1000 * 1000 * 1000); private static final int defaultChunkSizeBytes = Parameter.intValue("default.chunk.size.bytes", 16 * 1000 * 1000); private final String filePrefix = "dbl2file-"; private final String fileSuffix = ".dat"; private final File directory; public DiskBackedList2(ItemCodec<K> codec) throws IOException { this(codec, defaultChunkSizeBytes, LessFiles.createTempDir()); } public DiskBackedList2(ItemCodec<K> codec, long maxChunkSizeBytes, File directory) throws IOException { this.codec = codec; this.maxChunkSizeBytes = maxChunkSizeBytes; this.directory = directory; this.chunks = new ArrayList<>(); this.currentChunk = addChunk(); this.totalItems = 0; } public void setCodec(ItemCodec<K> codec) { this.codec = codec; } @Override protected void finalize() { chunks.clear(); currentChunk = null; } public String toString() { return "DBL: rows=" + totalItems + ", numChunks=" + chunks.size() + ", maxChunkSize=" + maxChunkSizeBytes; } // Which chunk is this element on and where is the element in that list? private Pair<Integer, Integer> indicesForElement(int elementIndex) { int chunkIndex = 0; int remaining = elementIndex; for (DBLChunk chunk : chunks) { int numEntries = chunk.getNumEntries(); if (remaining >= numEntries) { remaining -= numEntries; chunkIndex += 1; } else { break; } } return Pair.of(chunkIndex, remaining); } private void loadChunk(int chunkIndex) throws IOException { if (currentChunk == null) { loadChunkFully(chunkIndex); } else if (currentChunk.getIndex() != chunkIndex) { currentChunk.saveToFile(); currentChunk.clear(); loadChunkFully(chunkIndex); } } private void installChunk(int index, DBLChunk chunk) { try { while (index >= chunks.size()) { chunks.add(new DBLChunk(chunks.size(), directory)); } chunks.set(index, chunk); totalItems += chunk.getNumEntries(); currentChunk = null; } catch (IOException io) { log.warn("exception during install chunk: ", io); } } @Override public boolean add(K element) { try { loadChunk(chunks.size() - 1); if (!currentChunk.hasRoom()) { currentChunk = addChunk(); } currentChunk.store(element); totalItems++; return true; } catch (IOException ex) { log.warn("[disk.backed.list] exception while adding " + element, ex); } return false; } @Override /** * This function adds an element to a particular location. * TODO does not presently respect maxChunkSize */ public void add(int elementIndex, K element) { Pair<Integer, Integer> indices = indicesForElement(elementIndex); try { loadChunk(indices.getLeft()); currentChunk.store(indices.getRight(), element); totalItems++; } catch (IOException io) { log.warn("[disk.backed.list] exception while adding " + element, io); } } @Override public boolean addAll(Collection<? extends K> c) { for (K k : c) { add(k); } return true; } @Override public boolean addAll(int index, Collection<? extends K> c) { for (K k : c) { add(index++, k); } return true; } @Override public void clear() { for (DBLChunk chunk : chunks) { chunk.deleteFile(); } } @Override public boolean contains(Object o) { return indexOf(o) >= 0; } @Override public boolean containsAll(Collection<?> c) { for (Iterator<?> iter = c.iterator(); iter.hasNext();) { if (!contains(iter.next())) { return false; } } return true; } @Override public K get(int elementIndex) { try { Pair<Integer, Integer> indices = indicesForElement(elementIndex); loadChunk(indices.getLeft()); return currentChunk.get(indices.getRight()); } catch (IOException ex) { log.warn("[disk.backed.list] exception while getting element " + elementIndex, ex); return null; } } @Override public int indexOf(Object o) { int index = 0; //noinspection unchecked for (K elt : (Iterable<K>) iterator()) { if (elt.equals(o)) { return index; } index++; } return -1; } @Override public boolean isEmpty() { return chunks.isEmpty(); } @Override public Iterator<K> iterator() { return listIterator(); } @Override public int lastIndexOf(Object o) { int index = 0; int found = -1; //noinspection unchecked for (K elt : (Iterable<K>) iterator()) { if (elt.equals(o)) { found = index; } index++; } return found; } @Override public ListIterator<K> listIterator() { return listIterator(0); } @Override public ListIterator<K> listIterator(final int ind) { try { if (currentChunk != null) { currentChunk.saveToFile(); } } catch (IOException ex) { log.warn("[disk.backed.list] exception during saving", ex); return null; } return new ListIterator<K>() { private int index = ind; @Override public boolean hasNext() { return index < totalItems; } @Override public K next() { return get(index++); } @Override public void remove() { throw new UnsupportedOperationException(); } @Override public void add(K e) { DiskBackedList2.this.add(e); } @Override public boolean hasPrevious() { return index > 0; } @Override public int nextIndex() { return index; } @Override public K previous() { return get(index--); } @Override public int previousIndex() { return index - 1; } @Override public void set(K e) { throw new UnsupportedOperationException(); } }; } @Override public boolean remove(Object o) { return remove(indexOf(o)) != null; } @Override public K remove(int elementIndex) { Pair<Integer, Integer> indices = indicesForElement(elementIndex); try { loadChunk(indices.getLeft()); } catch (IOException io) { log.warn("[disk.backed.list] io exception during remove operation", io); return null; } K elt = currentChunk.remove(indices.getRight().intValue()); totalItems -= 1; return elt; } @Override public boolean removeAll(Collection<?> c) { boolean success = true; for (Iterator<?> iter = c.iterator(); iter.hasNext();) { if (!remove(iter.next())) { success = false; } } return success; } @Override public boolean retainAll(Collection<?> c) { throw new UnsupportedOperationException(); } @Override public K set(int elementIndex, K element) { Pair<Integer, Integer> indices = indicesForElement(elementIndex); try { loadChunk(indices.getLeft()); } catch (IOException io) { log.warn("[disk.backed.list] io exception during remove operation", io); } return currentChunk.set(indices.getRight(), element); } @Override public int size() { return totalItems; } @Override public List<K> subList(int fromIndex, int toIndex) { throw new UnsupportedOperationException(); } @Override public Object[] toArray() { throw new UnsupportedOperationException(); } @Override public <T> T[] toArray(T[] a) { throw new UnsupportedOperationException(); } /** * clean up and close */ public void close() throws IOException { if (currentChunk != null) { currentChunk.saveToFile(); currentChunk = null; } } /** * Sort the collection of elements using a standard external sort algorithm: sort each chunk of elements, then * merge the chunks into a new list, then switch to the new list. */ public void sort(final Comparator<? super K> comp) { try { // Sort each chunk. Done if there is only one chunk. sortEachChunk(comp); if (chunks.size() <= 1) { return; } Comparator<Pair<K, Integer>> pairComp = new Comparator<Pair<K, Integer>>() { @Override public int compare(Pair<K, Integer> e1, Pair<K, Integer> e2) { return comp.compare(e1.getLeft(), e2.getLeft()); } }; // This heap stores the lowest remaining value from each chunk PriorityQueue<Pair<K, Integer>> heap = new PriorityQueue<>(chunks.size(), pairComp); ArrayList<Iterator> iterators = new ArrayList<>(chunks.size()); // Initialize the heap with one value per chunk close(); for (int i = 0; i < chunks.size(); i++) { Iterator<K> it = chunks.get(i).getChunkIterator(); iterators.add(i, it); if (it.hasNext()) { K elt = it.next(); if (elt != null) { heap.add(Pair.of(elt, i)); } } } // Make a new disk backed list to store sorted values. // When the number of chunks is large, the size of the output buffer needs to shrink to make up for the extra mem usage long storageMaxChunkSize = maxChunkSizeBytes / (1 + chunks.size() / 20); DiskBackedList2<K> storage = new DiskBackedList2<>(codec, storageMaxChunkSize, directory); // Repeatedly pull the smallest element from the heap while (!heap.isEmpty()) { Pair<K, Integer> leastElt = heap.poll(); storage.add(leastElt.getLeft()); @SuppressWarnings({ "unchecked" }) Iterator<K> polledIterator = iterators.get(leastElt.getRight()); if (polledIterator.hasNext()) { heap.add(Pair.of(polledIterator.next(), leastElt.getRight())); } } // Switch to the storage dbl's chunks storage.close(); chunks = storage.getChunks(); currentChunk = null; } catch (IOException io) { throw Throwables.propagate(io); } } // This function is used to switch to the sorted values after a sort private List<DBLChunk> getChunks() { return chunks; } private void sortEachChunk(final Comparator<? super K> comp) throws IOException { int startChunkIndex = currentChunk != null ? currentChunk.getIndex() : 0; int numChunks = chunks.size(); for (int i = 0; i < numChunks; i++) { int chunkIndex = (i + startChunkIndex) % numChunks; loadChunk(chunkIndex); Collections.sort(currentChunk, comp); } } public void addEncodedData(Iterator<byte[]> stream) throws Exception { while (stream.hasNext()) { byte[] bytes = stream.next(); add(codec.decode(bytes)); } } public Iterator<byte[]> getEncodedData() { return new Iterator<byte[]>() { final Iterator<K> iter = listIterator(); @Override public boolean hasNext() { return iter.hasNext(); } @Override public byte[] next() { try { return codec.encode(iter.next()); } catch (IOException e) { throw new RuntimeException(e); } } @Override public void remove() { throw new UnsupportedOperationException(); } }; } private DBLChunk addChunk() throws IOException { if (currentChunk != null) { currentChunk.saveToFile(); currentChunk.clear(); } int index = chunks.size(); if (maxTotalSizeBytes > 0 && index * maxChunkSizeBytes > maxTotalSizeBytes) { log.warn("[disk.backed.list] disk space limit exceeded: " + this.toString()); throw new RuntimeException( "Query exceeded the disk space limit on sort file storage. (Number of rows at time of failure: " + totalItems + "). Try using query operations to reduce the size and number of rows being sorted, so as not to stress the QueryMaster."); } DBLChunk chunk = new DBLChunk(index, directory); chunks.add(chunk); return chunk; } /** * Load the specified chunk from the disk. By design, only currentChunk can be loaded into memory, ensuring that * only one DBLChunk is loaded at any given time. */ private void loadChunkFully(int index) throws IOException { currentChunk = chunks.get(index); currentChunk.resetNumEntries(); currentChunk.readFromFile(); } public void loadAllChunksFromDirectory() throws IOException { if (directory != null && directory.isDirectory()) { for (int i = 0; i < directory.listFiles().length; i++) { File file = new File(directory, filePrefix + Integer.toString(i) + fileSuffix); if (file.exists()) { DBLChunk chunk = new DBLChunk(i, directory); chunk.readFromFile(); installChunk(i, chunk); } } } } public void saveAllChunksToDirectory() throws IOException { if (directory != null && directory.isDirectory()) { currentChunk.saveToFile(); } } /** * An ArrayList of entries that can be loaded and saved to disk. * Each DBLChunk is small enough that it can be loaded entirely into memory. */ private class DBLChunk extends ArrayList<K> { private final int index; private int numEntries; private long availableBytes; private File file; public String makeFileName() { return filePrefix + index; } public DBLChunk(int index, File directory) throws IOException { this.index = index; this.numEntries = 0; this.availableBytes = maxChunkSizeBytes; file = new File(directory, makeFileName() + fileSuffix); } public int getIndex() { return index; } public boolean hasRoom() { return availableBytes > 0; } public int getNumEntries() { return numEntries; } public boolean store(K element) { return store(this.size(), element); } public boolean store(int elementIndex, K element) { long totalBytes = 4 + MemoryCounter.estimateSize(element); availableBytes -= totalBytes; this.add(elementIndex, element); numEntries += 1; return true; } public void saveToFile() throws IOException { FileOutputStream fos = new FileOutputStream(file); BufferedOutputStream bos = new BufferedOutputStream(fos); DataOutputStream dos = new DataOutputStream(bos); dos.writeInt(this.size()); for (K val : this) { byte[] valBytes = codec.encode(val); if (log.isDebugEnabled()) log.debug("stf: encoding " + val + " as " + valBytes.length + " to " + file); dos.writeInt(valBytes.length); dos.write(valBytes); } dos.close(); } public boolean deleteFile() { return file != null && file.delete(); } public void readFromFile() throws IOException { FileInputStream fis = new FileInputStream(file); BufferedInputStream bis = new BufferedInputStream(fis); DataInputStream dis = new DataInputStream(bis); int newNum = dis.readInt(); for (int i = 0; i < newNum; i++) { int newLen = dis.readInt(); byte[] bytes = LessBytes.readBytes(dis, newLen); if (bytes == null || bytes.length == 0 || newLen == 0) { log.warn("read null/0 bytes @ i=" + i + " of " + newNum + " for index=" + index + " file=" + file); } else { if (log.isDebugEnabled()) log.debug("rff: decoding " + bytes.length + " bytes from " + file); K tv = codec.decode(bytes); this.store(tv); } } dis.close(); } public void resetNumEntries() { numEntries = 0; } public Iterator<K> getChunkIterator() throws IOException { return new ListIterator<K>() { final FileInputStream fis = new FileInputStream(file); final BufferedInputStream bis = new BufferedInputStream(fis); final DataInputStream dis = new DataInputStream(bis); int itemsRemaining = dis.readInt(); { if (log.isDebugEnabled()) log.debug("iter: items=" + itemsRemaining + " file=" + file); } @Override public boolean hasNext() { return itemsRemaining > 0; } @Override public K next() { try { int eltLength = dis.readInt(); byte[] bytes = LessBytes.readBytes(dis, eltLength); itemsRemaining--; if (itemsRemaining == 0) { dis.close(); } if (log.isDebugEnabled()) { log.debug("iter: decoding " + bytes.length + " bytes remain=" + itemsRemaining + " from " + file); } return codec.decode(bytes); } catch (IOException io) { log.warn("[disk.backed.list] io exception during chunk iteration", io); return null; } } @Override public void remove() { throw new UnsupportedOperationException(); } @Override public void add(K e) { throw new UnsupportedOperationException(); } @Override public boolean hasPrevious() { throw new UnsupportedOperationException(); } @Override public int nextIndex() { throw new UnsupportedOperationException(); } @Override public K previous() { throw new UnsupportedOperationException(); } @Override public int previousIndex() { throw new UnsupportedOperationException(); } @Override public void set(K e) { throw new UnsupportedOperationException(); } }; } } public File getDirectory() { return directory; } }