private static Environment getMockEnvironment(File[] tempDirs) { IOManager ioMan = mock(IOManager.class); when(ioMan.getSpillingDirectories()).thenReturn(tempDirs); Environment env = mock(Environment.class); when(env.getJobID()).thenReturn(new JobID()); when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader()); when(env.getIOManager()).thenReturn(ioMan); return env; }
@After public void afterTest() throws Exception { ioManager.shutdown(); Assert.assertTrue("IO Manager has not properly shut down.", ioManager.isProperlyShutDown()); Assert.assertTrue( "Not all memory was returned to the memory manager in the test.", memManager.verifyEmpty()); memManager.shutdown(); memManager = null; }
SpilledSubpartitionViewAsyncIO( ResultSubpartition parent, BufferProvider bufferProvider, IOManager ioManager, FileIOChannel.ID channelId, long initialSeekPosition, int readBatchSize) throws IOException { checkArgument(initialSeekPosition >= 0, "Initial seek position is < 0."); checkArgument(readBatchSize >= 1, "Batch read size < 1."); this.parent = checkNotNull(parent); this.bufferProvider = checkNotNull(bufferProvider); this.bufferAvailabilityListener = new BufferProviderCallback(this); this.asyncFileReader = ioManager.createBufferFileReader(channelId, new IOThreadCallback(this)); if (initialSeekPosition > 0) { asyncFileReader.seekToPosition(initialSeekPosition); } this.readBatchSize = readBatchSize; this.fileSize = asyncFileReader.getSize(); // Trigger the initial read requests readNextBatchAsync(); }
private static Environment getMockEnvironment(File[] tempDirs) { IOManager ioMan = mock(IOManager.class); when(ioMan.getSpillingDirectories()).thenReturn(tempDirs); Environment env = mock(Environment.class); when(env.getJobID()).thenReturn(new JobID()); when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader()); when(env.getIOManager()).thenReturn(ioMan); when(env.getTaskKvStateRegistry()) .thenReturn(new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); TaskInfo taskInfo = mock(TaskInfo.class); when(env.getTaskInfo()).thenReturn(taskInfo); when(taskInfo.getIndexOfThisSubtask()).thenReturn(0); return env; }
private final void testChannelWithSegments(int numSegments) throws Exception { final List<MemorySegment> memory = this.memManager.allocatePages(memoryOwner, numSegments); final Channel.ID channel = this.ioManager.createChannel(); BlockChannelWriter writer = null; BlockChannelReader reader = null; try { writer = this.ioManager.createBlockChannelWriter(channel); final ChannelWriterOutputView out = new ChannelWriterOutputView(writer, memory, this.memManager.getPageSize()); long writeStart = System.currentTimeMillis(); int valsLeft = NUM_INTS_WRITTEN; while (valsLeft-- > 0) { out.writeInt(valsLeft); } out.close(); final int numBlocks = out.getBlockCount(); writer.close(); writer = null; long writeElapsed = System.currentTimeMillis() - writeStart; // ---------------------------------------------------------------- reader = ioManager.createBlockChannelReader(channel); final ChannelReaderInputView in = new ChannelReaderInputView(reader, memory, numBlocks, false); long readStart = System.currentTimeMillis(); valsLeft = NUM_INTS_WRITTEN; while (valsLeft-- > 0) { in.readInt(); // Assert.assertTrue(rec.getValue() == valsLeft); } in.close(); reader.close(); long readElapsed = System.currentTimeMillis() - readStart; reader.deleteChannel(); reader = null; LOG.info( "IOManager with " + numSegments + " mem segments: write " + writeElapsed + " msecs, read " + readElapsed + " msecs."); memManager.release(memory); } finally { if (reader != null) { reader.closeAndDelete(); } if (writer != null) { writer.closeAndDelete(); } } }
@SuppressWarnings("resource") private final void speedTestNIO(int bufferSize, boolean direct) throws IOException { final Channel.ID tmpChannel = ioManager.createChannel(); File tempFile = null; FileChannel fs = null; try { tempFile = new File(tmpChannel.getPath()); RandomAccessFile raf = new RandomAccessFile(tempFile, "rw"); fs = raf.getChannel(); ByteBuffer buf = direct ? ByteBuffer.allocateDirect(bufferSize) : ByteBuffer.allocate(bufferSize); long writeStart = System.currentTimeMillis(); int valsLeft = NUM_INTS_WRITTEN; while (valsLeft-- > 0) { if (buf.remaining() < 4) { buf.flip(); fs.write(buf); buf.clear(); } buf.putInt(valsLeft); } if (buf.position() > 0) { buf.flip(); fs.write(buf); } fs.close(); raf.close(); fs = null; long writeElapsed = System.currentTimeMillis() - writeStart; // ---------------------------------------------------------------- raf = new RandomAccessFile(tempFile, "r"); fs = raf.getChannel(); buf.clear(); long readStart = System.currentTimeMillis(); fs.read(buf); buf.flip(); valsLeft = NUM_INTS_WRITTEN; while (valsLeft-- > 0) { if (buf.remaining() < 4) { buf.compact(); fs.read(buf); buf.flip(); } if (buf.getInt() != valsLeft) { throw new IOException(); } } fs.close(); raf.close(); long readElapsed = System.currentTimeMillis() - readStart; LOG.info( "NIO Channel with buffer " + bufferSize + ": write " + writeElapsed + " msecs, read " + readElapsed + " msecs."); } finally { // close if possible if (fs != null) { fs.close(); fs = null; } // try to delete the file if (tempFile != null) { tempFile.delete(); } } }
private final void speedTestStream(int bufferSize) throws IOException { final Channel.ID tmpChannel = ioManager.createChannel(); final IntegerRecord rec = new IntegerRecord(0); File tempFile = null; DataOutputStream daos = null; DataInputStream dais = null; try { tempFile = new File(tmpChannel.getPath()); FileOutputStream fos = new FileOutputStream(tempFile); daos = new DataOutputStream(new BufferedOutputStream(fos, bufferSize)); long writeStart = System.currentTimeMillis(); int valsLeft = NUM_INTS_WRITTEN; while (valsLeft-- > 0) { rec.setValue(valsLeft); rec.write(new OutputViewDataOutputStreamWrapper(daos)); } daos.close(); daos = null; long writeElapsed = System.currentTimeMillis() - writeStart; // ---------------------------------------------------------------- FileInputStream fis = new FileInputStream(tempFile); dais = new DataInputStream(new BufferedInputStream(fis, bufferSize)); long readStart = System.currentTimeMillis(); valsLeft = NUM_INTS_WRITTEN; while (valsLeft-- > 0) { rec.read(new InputViewDataInputStreamWrapper(dais)); } dais.close(); dais = null; long readElapsed = System.currentTimeMillis() - readStart; LOG.info( "File-Stream with buffer " + bufferSize + ": write " + writeElapsed + " msecs, read " + readElapsed + " msecs."); } finally { // close if possible if (daos != null) { daos.close(); } if (dais != null) { dais.close(); } // try to delete the file if (tempFile != null) { tempFile.delete(); } } }
/** * Closes all structures and deletes all temporary files. Even in the presence of failures, this * method will try and continue closing files and deleting temporary files. * * @throws IOException Thrown if an error occurred while closing/deleting the files. */ public void close() throws IOException { // we go on closing and deleting files in the presence of failures. // we remember the first exception to occur and re-throw it later Throwable ex = null; synchronized (this) { if (closed) { return; } closed = true; // close the writers if (recordsOutFile != null) { try { recordsOutFile.close(); recordsOutFile = null; } catch (Throwable t) { LOG.error("Cannot close the large records spill file.", t); ex = ex == null ? t : ex; } } if (keysOutFile != null) { try { keysOutFile.close(); keysOutFile = null; } catch (Throwable t) { LOG.error("Cannot close the large records key spill file.", t); ex = ex == null ? t : ex; } } // close the readers if (recordsReader != null) { try { recordsReader.close(); recordsReader = null; } catch (Throwable t) { LOG.error("Cannot close the large records reader.", t); ex = ex == null ? t : ex; } } if (keysReader != null) { try { keysReader.close(); keysReader = null; } catch (Throwable t) { LOG.error("Cannot close the large records key reader.", t); ex = ex == null ? t : ex; } } // delete the spill files if (recordsChannel != null) { try { ioManager.deleteChannel(recordsChannel); recordsChannel = null; } catch (Throwable t) { LOG.error("Cannot delete the large records spill file.", t); ex = ex == null ? t : ex; } } if (keysChannel != null) { try { ioManager.deleteChannel(keysChannel); keysChannel = null; } catch (Throwable t) { LOG.error("Cannot delete the large records key spill file.", t); ex = ex == null ? t : ex; } } // close the key sorter if (keySorter != null) { try { keySorter.close(); keySorter = null; } catch (Throwable t) { LOG.error("Cannot properly dispose the key sorter and clean up its temporary files.", t); ex = ex == null ? t : ex; } } memManager.release(memory); recordCounter = 0; } // re-throw the exception, if necessary if (ex != null) { throw new IOException( "An error occurred cleaning up spill files in the large record handler.", ex); } }
public MutableObjectIterator<T> finishWriteAndSortKeys(List<MemorySegment> memory) throws IOException { if (recordsOutFile == null || keysOutFile == null) { throw new IllegalStateException("The LargeRecordHandler has not spilled any records"); } // close the writers and final int lastBlockBytesKeys; final int lastBlockBytesRecords; recordsOutFile.close(); keysOutFile.close(); lastBlockBytesKeys = keysOutFile.getBytesInLatestSegment(); lastBlockBytesRecords = recordsOutFile.getBytesInLatestSegment(); recordsOutFile = null; keysOutFile = null; final int pagesForReaders = Math.max( 3 * MIN_SEGMENTS_FOR_KEY_SPILLING, Math.min(2 * MAX_SEGMENTS_FOR_KEY_SPILLING, memory.size() / 50)); final int pagesForKeyReader = Math.min(pagesForReaders - MIN_SEGMENTS_FOR_KEY_SPILLING, MAX_SEGMENTS_FOR_KEY_SPILLING); final int pagesForRecordReader = pagesForReaders - pagesForKeyReader; // grab memory for the record reader ArrayList<MemorySegment> memForRecordReader = new ArrayList<MemorySegment>(); ArrayList<MemorySegment> memForKeysReader = new ArrayList<MemorySegment>(); for (int i = 0; i < pagesForRecordReader; i++) { memForRecordReader.add(memory.remove(memory.size() - 1)); } for (int i = 0; i < pagesForKeyReader; i++) { memForKeysReader.add(memory.remove(memory.size() - 1)); } keysReader = new FileChannelInputView( ioManager.createBlockChannelReader(keysChannel), memManager, memForKeysReader, lastBlockBytesKeys); InputViewIterator<Tuple> keyIterator = new InputViewIterator<Tuple>(keysReader, keySerializer); keySorter = new UnilateralSortMerger<Tuple>( memManager, memory, ioManager, keyIterator, memoryOwner, keySerializerFactory, keyComparator, 1, maxFilehandles, 1.0f, false); // wait for the sorter to sort the keys MutableObjectIterator<Tuple> result; try { result = keySorter.getIterator(); } catch (InterruptedException e) { throw new IOException(e); } recordsReader = new SeekableFileChannelInputView( ioManager, recordsChannel, memManager, memForRecordReader, lastBlockBytesRecords); return new FetchingIterator<T>(serializer, result, recordsReader, keySerializer, numKeyFields); }
@SuppressWarnings("unchecked") public long addRecord(T record) throws IOException { if (recordsOutFile == null) { if (closed) { throw new IllegalStateException("The large record handler has been closed."); } if (recordsReader != null) { throw new IllegalStateException("The handler has already switched to sorting."); } LOG.debug("Initializing the large record spilling..."); // initialize the utilities { final TypeComparator<?>[] keyComps = comparator.getFlatComparators(); numKeyFields = keyComps.length; Object[] keyHolder = new Object[numKeyFields]; comparator.extractKeys(record, keyHolder, 0); TypeSerializer<?>[] keySers = new TypeSerializer<?>[numKeyFields]; TypeSerializer<?>[] tupleSers = new TypeSerializer<?>[numKeyFields + 1]; int[] keyPos = new int[numKeyFields]; for (int i = 0; i < numKeyFields; i++) { keyPos[i] = i; keySers[i] = createSerializer(keyHolder[i], i); tupleSers[i] = keySers[i]; } // add the long serializer for the offset tupleSers[numKeyFields] = LongSerializer.INSTANCE; keySerializer = new TupleSerializer<Tuple>( (Class<Tuple>) Tuple.getTupleClass(numKeyFields + 1), tupleSers); keyComparator = new TupleComparator<Tuple>(keyPos, keyComps, keySers); keySerializerFactory = new RuntimeSerializerFactory<Tuple>(keySerializer, keySerializer.getTupleClass()); keyTuple = keySerializer.createInstance(); } // initialize the spilling final int totalNumSegments = memory.size(); final int segmentsForKeys = (totalNumSegments >= 2 * MAX_SEGMENTS_FOR_KEY_SPILLING) ? MAX_SEGMENTS_FOR_KEY_SPILLING : Math.max( MIN_SEGMENTS_FOR_KEY_SPILLING, totalNumSegments - MAX_SEGMENTS_FOR_KEY_SPILLING); List<MemorySegment> recordsMemory = new ArrayList<MemorySegment>(); List<MemorySegment> keysMemory = new ArrayList<MemorySegment>(); for (int i = 0; i < segmentsForKeys; i++) { keysMemory.add(memory.get(i)); } for (int i = segmentsForKeys; i < totalNumSegments; i++) { recordsMemory.add(memory.get(i)); } recordsChannel = ioManager.createChannel(); keysChannel = ioManager.createChannel(); recordsOutFile = new FileChannelOutputView( ioManager.createBlockChannelWriter(recordsChannel), memManager, recordsMemory, memManager.getPageSize()); keysOutFile = new FileChannelOutputView( ioManager.createBlockChannelWriter(keysChannel), memManager, keysMemory, memManager.getPageSize()); } final long offset = recordsOutFile.getWriteOffset(); if (offset < 0) { throw new RuntimeException("wrong offset"); } Object[] keyHolder = new Object[numKeyFields]; comparator.extractKeys(record, keyHolder, 0); for (int i = 0; i < numKeyFields; i++) { keyTuple.setField(keyHolder[i], i); } keyTuple.setField(offset, numKeyFields); keySerializer.serialize(keyTuple, keysOutFile); serializer.serialize(record, recordsOutFile); recordCounter++; return offset; }