CachingShuffleBatchReader implements ShuffleBatchReader { @Override public Batch read(@Nullable ShufflePosition startPosition, @Nullable ShufflePosition endPosition) throws IOException { final BatchRange range = new BatchRange(startPosition, endPosition); try { return cache.get(range); } catch (RuntimeException | ExecutionException e) { Throwables.propagateIfPossible(e, IOException.class); throw new RuntimeException("unexpected", e); } } CachingShuffleBatchReader(
      ShuffleBatchReader shuffleReader, int maximumBatches, long expireAfterAccessMillis);  CachingShuffleBatchReader(ShuffleBatchReader shuffleReader);  CachingShuffleBatchReader(ShuffleBatchReader shuffleReader, int maximumBatches); @Override Batch read(@Nullable ShufflePosition startPosition, @Nullable ShufflePosition endPosition);  }
@Test public void readerShouldCacheReads() throws IOException { ShuffleBatchReader base = mock(ShuffleBatchReader.class); CachingShuffleBatchReader reader = new CachingShuffleBatchReader(base); when(base.read(null, null)).thenReturn(testBatch); ShuffleBatchReader.Batch read = reader.read(null, null); assertThat(read, equalTo(testBatch)); assertThat(reader.read(null, null), equalTo(testBatch)); assertThat(reader.read(null, null), equalTo(testBatch)); assertThat(reader.read(null, null), equalTo(testBatch)); assertThat(reader.read(null, null), equalTo(testBatch)); verify(base, times(1)).read(null, null); }
@Test public void readerShouldNotCacheExceptions() throws IOException { ShuffleBatchReader base = mock(ShuffleBatchReader.class); CachingShuffleBatchReader reader = new CachingShuffleBatchReader(base); when(base.read(null, null)).thenThrow(new IOException("test")).thenReturn(testBatch); try { reader.read(null, null); fail("expected an IOException"); } catch (Exception e) { } assertThat(reader.read(null, null), equalTo(testBatch)); verify(base, times(2)).read(null, null); }
@Test public void readerShouldRereadEvictedBatches() throws IOException, ExecutionException { ShuffleBatchReader base = mock(ShuffleBatchReader.class); CachingShuffleBatchReader reader = new CachingShuffleBatchReader(base); when(base.read(null, null)).thenReturn(testBatch); ShuffleBatchReader.Batch read = reader.read(null, null); assertThat(read, equalTo(testBatch)); verify(base, times(1)).read(null, null); CachingShuffleBatchReader.BatchRange range = new CachingShuffleBatchReader.BatchRange(null, null); CachingShuffleBatchReader.Batch batch = reader.cache.get(range); assertThat(batch, notNullValue()); reader.cache.invalidateAll(); read = reader.read(null, null); assertThat(read, equalTo(testBatch)); verify(base, times(2)).read(null, null); }