Skip to content

Commit aa31aa3

Browse files
jeqojeqo
committed
feat: prefetch next segment
1 parent d4209a1 commit aa31aa3

10 files changed

Lines changed: 268 additions & 89 deletions

File tree

‎core/src/main/java/io/aiven/kafka/tieredstorage/RemoteStorageManager.java‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -428,6 +428,18 @@ public InputStream fetchLogSegment(final RemoteLogSegmentMetadata remoteLogSegme
428428
final int startPosition) throws RemoteStorageException {
429429
return this.fetchLogSegment(
430430
remoteLogSegmentMetadata,
431+
Optional.empty(),
432+
startPosition
433+
);
434+
}
435+
436+
@Override
437+
public InputStream fetchLogSegment(final RemoteLogSegmentMetadata remoteLogSegmentMetadata,
438+
final Optional<RemoteLogSegmentMetadata> nextRemoteLogSegmentMetadata,
439+
final int startPosition) throws RemoteStorageException {
440+
return this.fetchLogSegment(
441+
remoteLogSegmentMetadata,
442+
nextRemoteLogSegmentMetadata,
431443
startPosition,
432444
remoteLogSegmentMetadata.segmentSizeInBytes() - 1
433445
);
@@ -437,11 +449,23 @@ public InputStream fetchLogSegment(final RemoteLogSegmentMetadata remoteLogSegme
437449
public InputStream fetchLogSegment(final RemoteLogSegmentMetadata remoteLogSegmentMetadata,
438450
final int startPosition,
439451
final int endPosition) throws RemoteStorageException {
452+
return this.fetchLogSegment(
453+
remoteLogSegmentMetadata,
454+
Optional.empty(),
455+
startPosition,
456+
endPosition
457+
);
458+
}
459+
460+
@Override
461+
public InputStream fetchLogSegment(final RemoteLogSegmentMetadata remoteLogSegmentMetadata,
462+
final Optional<RemoteLogSegmentMetadata> nextRemoteLogSegmentMetadata,
463+
final int startPosition,
464+
final int endPosition) throws RemoteStorageException {
440465
try {
441-
final BytesRange range = BytesRange.of(
442-
startPosition,
443-
Math.min(endPosition, remoteLogSegmentMetadata.segmentSizeInBytes() - 1)
444-
);
466+
final var endOfFile = remoteLogSegmentMetadata.segmentSizeInBytes() - 1;
467+
final var actualEndPosition = Math.min(endPosition, endOfFile);
468+
final BytesRange range = BytesRange.of(startPosition, actualEndPosition);
445469

446470
log.trace("Fetching log segment {} with range: {}", remoteLogSegmentMetadata, range);
447471

@@ -451,10 +475,23 @@ public InputStream fetchLogSegment(final RemoteLogSegmentMetadata remoteLogSegme
451475

452476
final var segmentManifest = fetchSegmentManifest(remoteLogSegmentMetadata);
453477

454-
final var suffix = ObjectKeyFactory.Suffix.LOG;
455-
final var segmentKey = objectKey(remoteLogSegmentMetadata, suffix);
456-
return new FetchChunkEnumeration(chunkManager, segmentKey, segmentManifest, range)
457-
.toInputStream();
478+
final var segmentKey = objectKey(remoteLogSegmentMetadata, ObjectKeyFactory.Suffix.LOG);
479+
final var fetchChunkEnumeration = new FetchChunkEnumeration(
480+
chunkManager,
481+
segmentKey,
482+
segmentManifest,
483+
range,
484+
endOfFile,
485+
nextRemoteLogSegmentMetadata.map(next -> objectKey(next, ObjectKeyFactory.Suffix.LOG)),
486+
nextRemoteLogSegmentMetadata.map(next -> () -> {
487+
try {
488+
return fetchSegmentManifest(next);
489+
} catch (StorageBackendException | IOException e) {
490+
// ignore the error and let fetch of next segment to deal with it
491+
return null;
492+
}
493+
}));
494+
return fetchChunkEnumeration.toInputStream();
458495
} catch (final KeyNotFoundException | KeyNotFoundRuntimeException e) {
459496
throw new RemoteResourceNotFoundException(e);
460497
} catch (final Exception e) {

‎core/src/main/java/io/aiven/kafka/tieredstorage/chunkmanager/ChunkManager.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818

1919
import java.io.IOException;
2020
import java.io.InputStream;
21+
import java.util.Optional;
22+
import java.util.function.Supplier;
2123

2224
import io.aiven.kafka.tieredstorage.manifest.SegmentManifest;
2325
import io.aiven.kafka.tieredstorage.storage.ObjectKey;
@@ -27,5 +29,9 @@ public interface ChunkManager {
2729

2830
InputStream getChunk(final ObjectKey objectKey,
2931
final SegmentManifest manifest,
30-
final int chunkId) throws StorageBackendException, IOException;
32+
final int chunkId,
33+
final int endOfFile,
34+
final Optional<ObjectKey> maybeNextSegmentKey,
35+
final Optional<Supplier<SegmentManifest>> maybeNextSegmentManifestSupplier)
36+
throws StorageBackendException, IOException;
3137
}

‎core/src/main/java/io/aiven/kafka/tieredstorage/chunkmanager/DefaultChunkManager.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.io.InputStream;
2020
import java.util.List;
2121
import java.util.Optional;
22+
import java.util.function.Supplier;
2223

2324
import io.aiven.kafka.tieredstorage.Chunk;
2425
import io.aiven.kafka.tieredstorage.manifest.SegmentEncryptionMetadata;
@@ -47,8 +48,13 @@ public DefaultChunkManager(final ObjectFetcher fetcher, final AesEncryptionProvi
4748
*
4849
* @return an {@link InputStream} of the chunk, plain text (i.e., decrypted and decompressed).
4950
*/
50-
public InputStream getChunk(final ObjectKey objectKey, final SegmentManifest manifest,
51-
final int chunkId) throws StorageBackendException {
51+
public InputStream getChunk(final ObjectKey objectKey,
52+
final SegmentManifest manifest,
53+
final int chunkId,
54+
final int endOfFile,
55+
final Optional<ObjectKey> maybeNextSegmentKey,
56+
final Optional<Supplier<SegmentManifest>> maybeNextSegmentManifestSupplier)
57+
throws StorageBackendException {
5258
final Chunk chunk = manifest.chunkIndex().chunks().get(chunkId);
5359

5460
final InputStream chunkContent = fetcher.fetch(objectKey, chunk.range());

‎core/src/main/java/io/aiven/kafka/tieredstorage/chunkmanager/cache/ChunkCache.java‎

Lines changed: 54 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import java.io.IOException;
2020
import java.io.InputStream;
21+
import java.util.Optional;
2122
import java.util.concurrent.CompletableFuture;
2223
import java.util.concurrent.CompletionException;
2324
import java.util.concurrent.ExecutionException;
@@ -26,6 +27,7 @@
2627
import java.util.concurrent.TimeUnit;
2728
import java.util.concurrent.TimeoutException;
2829
import java.util.concurrent.atomic.AtomicReference;
30+
import java.util.function.Supplier;
2931

3032
import org.apache.kafka.common.Configurable;
3133

@@ -69,11 +71,23 @@ protected ChunkCache(final ChunkManager chunkManager) {
6971
* opened right when fetching from cache happens even if the actual value is removed from the cache,
7072
* the InputStream will still contain the data.
7173
*/
72-
public InputStream getChunk(final ObjectKey objectKey,
73-
final SegmentManifest manifest,
74-
final int chunkId) throws StorageBackendException, IOException {
74+
public InputStream getChunk(
75+
final ObjectKey objectKey,
76+
final SegmentManifest manifest,
77+
final int chunkId,
78+
final int endOfFile,
79+
final Optional<ObjectKey> maybeNextSegmentKey,
80+
final Optional<Supplier<SegmentManifest>> maybeNextSegmentManifestSupplier
81+
) throws StorageBackendException, IOException {
7582
final var currentChunk = manifest.chunkIndex().chunks().get(chunkId);
76-
startPrefetching(objectKey, manifest, currentChunk.originalPosition + currentChunk.originalSize);
83+
startPrefetching(
84+
objectKey,
85+
manifest,
86+
currentChunk.originalPosition + currentChunk.originalSize,
87+
endOfFile,
88+
maybeNextSegmentKey,
89+
maybeNextSegmentManifestSupplier
90+
);
7791
final ChunkKey chunkKey = new ChunkKey(objectKey.value(), chunkId);
7892
final AtomicReference<InputStream> result = new AtomicReference<>();
7993
try {
@@ -83,7 +97,8 @@ public InputStream getChunk(final ObjectKey objectKey,
8397
statsCounter.recordMiss();
8498
try {
8599
final InputStream chunk =
86-
chunkManager.getChunk(objectKey, manifest, chunkId);
100+
chunkManager.getChunk(objectKey, manifest, chunkId, endOfFile,
101+
maybeNextSegmentKey, maybeNextSegmentManifestSupplier);
87102
final T t = this.cacheChunk(chunkKey, chunk);
88103
result.getAndSet(cachedChunkToInputStream(t));
89104
return t;
@@ -147,8 +162,11 @@ protected AsyncCache<ChunkKey, T> buildCache(final ChunkCacheConfig config) {
147162
}
148163

149164
private void startPrefetching(final ObjectKey segmentKey,
150-
final SegmentManifest segmentManifest,
151-
final int startPosition) {
165+
final SegmentManifest segmentManifest,
166+
final int startPosition,
167+
final int endOfFile,
168+
final Optional<ObjectKey> maybeNextSegmentKey,
169+
final Optional<Supplier<SegmentManifest>> maybeNextSegmentManifestSupplier) {
152170
if (prefetchingSize > 0) {
153171
final BytesRange prefetchingRange;
154172
if (Integer.MAX_VALUE - startPosition < prefetchingSize) {
@@ -162,14 +180,41 @@ private void startPrefetching(final ObjectKey segmentKey,
162180
cache.asMap()
163181
.computeIfAbsent(chunkKey, key -> CompletableFuture.supplyAsync(() -> {
164182
try {
165-
final InputStream chunkStream =
166-
chunkManager.getChunk(segmentKey, segmentManifest, chunk.id);
183+
final InputStream chunkStream = chunkManager.getChunk(
184+
segmentKey,
185+
segmentManifest,
186+
chunk.id,
187+
endOfFile,
188+
maybeNextSegmentKey,
189+
maybeNextSegmentManifestSupplier
190+
);
167191
return this.cacheChunk(chunkKey, chunkStream);
168192
} catch (final StorageBackendException | IOException e) {
169193
throw new CompletionException(e);
170194
}
171195
}, executor));
172196
});
197+
if (maybeNextSegmentKey.isPresent()
198+
&& maybeNextSegmentManifestSupplier.isPresent()
199+
&& endOfFile - startPosition < prefetchingSize) {
200+
final ChunkKey chunkKey = new ChunkKey(maybeNextSegmentKey.get().value(), 0);
201+
cache.asMap()
202+
.computeIfAbsent(chunkKey, key -> CompletableFuture.supplyAsync(() -> {
203+
try {
204+
final InputStream chunkStream = chunkManager.getChunk(
205+
maybeNextSegmentKey.get(),
206+
maybeNextSegmentManifestSupplier.get().get(),
207+
0,
208+
endOfFile,
209+
Optional.empty(),
210+
Optional.empty()
211+
);
212+
return this.cacheChunk(chunkKey, chunkStream);
213+
} catch (final StorageBackendException | IOException e) {
214+
throw new CompletionException(e);
215+
}
216+
}, executor));
217+
}
173218
}
174219
}
175220
}

‎core/src/main/java/io/aiven/kafka/tieredstorage/transform/FetchChunkEnumeration.java‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import java.util.List;
2424
import java.util.NoSuchElementException;
2525
import java.util.Objects;
26+
import java.util.Optional;
27+
import java.util.function.Supplier;
2628

2729
import io.aiven.kafka.tieredstorage.Chunk;
2830
import io.aiven.kafka.tieredstorage.chunkmanager.ChunkManager;
@@ -40,27 +42,35 @@ public class FetchChunkEnumeration implements Enumeration<InputStream> {
4042
private final ObjectKey objectKey;
4143
private final SegmentManifest manifest;
4244
private final BytesRange range;
45+
private final int endOfFile;
46+
private final Optional<ObjectKey> maybeNextSegmentKey;
47+
private final Optional<Supplier<SegmentManifest>> maybeNextSegmentManifestSupplier;
4348
final int startChunkId;
4449
final int lastChunkId;
4550
private final ChunkIndex chunkIndex;
4651
int currentChunkId;
4752
public boolean closed;
4853

4954
/**
50-
*
5155
* @param chunkManager provides chunk input to fetch from
52-
* @param objectKey required by chunkManager
53-
* @param manifest provides to index to build response from
54-
* @param range original offset range start/end position
56+
* @param objectKey required by chunkManager
57+
* @param manifest provides to index to build response from
58+
* @param range original offset range start/end position
5559
*/
5660
public FetchChunkEnumeration(final ChunkManager chunkManager,
5761
final ObjectKey objectKey,
5862
final SegmentManifest manifest,
59-
final BytesRange range) {
63+
final BytesRange range,
64+
final int endOfFile,
65+
final Optional<ObjectKey> maybeNextSegmentKey,
66+
final Optional<Supplier<SegmentManifest>> maybeNextSegmentManifestSupplier) {
6067
this.chunkManager = Objects.requireNonNull(chunkManager, "chunkManager cannot be null");
6168
this.objectKey = Objects.requireNonNull(objectKey, "objectKey cannot be null");
6269
this.manifest = Objects.requireNonNull(manifest, "manifest cannot be null");
6370
this.range = Objects.requireNonNull(range, "range cannot be null");
71+
this.endOfFile = endOfFile;
72+
this.maybeNextSegmentKey = maybeNextSegmentKey;
73+
this.maybeNextSegmentManifestSupplier = maybeNextSegmentManifestSupplier;
6474

6575
this.chunkIndex = manifest.chunkIndex();
6676

@@ -138,7 +148,14 @@ public InputStream nextElement() {
138148

139149
private InputStream getChunkContent(final int chunkId) {
140150
try {
141-
return chunkManager.getChunk(objectKey, manifest, chunkId);
151+
return chunkManager.getChunk(
152+
objectKey,
153+
manifest,
154+
chunkId,
155+
endOfFile,
156+
maybeNextSegmentKey,
157+
maybeNextSegmentManifestSupplier
158+
);
142159
} catch (final KeyNotFoundException e) {
143160
throw new KeyNotFoundRuntimeException(e);
144161
} catch (final StorageBackendException | IOException e) {

‎core/src/test/java/io/aiven/kafka/tieredstorage/chunkmanager/DefaultChunkManagerTest.java‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import javax.crypto.Cipher;
2020

2121
import java.io.ByteArrayInputStream;
22+
import java.util.Optional;
2223

2324
import org.apache.kafka.server.log.remote.storage.RemoteStorageManager.IndexType;
2425

@@ -67,7 +68,8 @@ void testGetChunk() throws Exception {
6768
when(storage.fetch(OBJECT_KEY, chunkIndex.chunks().get(0).range()))
6869
.thenReturn(new ByteArrayInputStream("0123456789".getBytes()));
6970

70-
assertThat(chunkManager.getChunk(OBJECT_KEY, manifest, 0)).hasContent("0123456789");
71+
assertThat(chunkManager.getChunk(OBJECT_KEY, manifest, 0,
72+
-1, Optional.empty(), Optional.empty())).hasContent("0123456789");
7173
verify(storage).fetch(OBJECT_KEY, chunkIndex.chunks().get(0).range());
7274
}
7375

@@ -91,7 +93,8 @@ void testGetChunkWithEncryption() throws Exception {
9193
final var manifest = new SegmentManifestV1(chunkIndex, SEGMENT_INDEXES, false, encryption, null);
9294
final ChunkManager chunkManager = new DefaultChunkManager(storage, aesEncryptionProvider);
9395

94-
assertThat(chunkManager.getChunk(OBJECT_KEY, manifest, 0)).hasBinaryContent(TEST_CHUNK_CONTENT);
96+
assertThat(chunkManager.getChunk(OBJECT_KEY, manifest, 0,
97+
-1, Optional.empty(), Optional.empty())).hasBinaryContent(TEST_CHUNK_CONTENT);
9598
verify(storage).fetch(OBJECT_KEY, chunkIndex.chunks().get(0).range());
9699
}
97100

@@ -111,7 +114,8 @@ void testGetChunkWithCompression() throws Exception {
111114
final var manifest = new SegmentManifestV1(chunkIndex, SEGMENT_INDEXES, true, null, null);
112115
final ChunkManager chunkManager = new DefaultChunkManager(storage, null);
113116

114-
assertThat(chunkManager.getChunk(OBJECT_KEY, manifest, 0)).hasBinaryContent(TEST_CHUNK_CONTENT);
117+
assertThat(chunkManager.getChunk(OBJECT_KEY, manifest, 0,
118+
-1, Optional.empty(), Optional.empty())).hasBinaryContent(TEST_CHUNK_CONTENT);
115119
verify(storage).fetch(OBJECT_KEY, chunkIndex.chunks().get(0).range());
116120
}
117121
}

‎core/src/test/java/io/aiven/kafka/tieredstorage/chunkmanager/cache/ChunkCacheMetricsTest.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.lang.management.ManagementFactory;
2424
import java.nio.file.Path;
2525
import java.util.Map;
26+
import java.util.Optional;
2627
import java.util.stream.Stream;
2728

2829
import io.aiven.kafka.tieredstorage.chunkmanager.ChunkManager;
@@ -93,7 +94,7 @@ void setUp() {
9394
void shouldRecordMetrics(final Class<ChunkCache<?>> chunkCacheClass, final Map<String, ?> config)
9495
throws Exception {
9596
// Given a chunk cache implementation
96-
when(chunkManager.getChunk(any(), any(), anyInt()))
97+
when(chunkManager.getChunk(any(), any(), anyInt(), anyInt(), any(), any()))
9798
.thenReturn(new ByteArrayInputStream("test".getBytes()));
9899

99100
final var chunkCache = chunkCacheClass.getDeclaredConstructor(ChunkManager.class).newInstance(chunkManager);
@@ -102,13 +103,13 @@ void shouldRecordMetrics(final Class<ChunkCache<?>> chunkCacheClass, final Map<S
102103
final var objectName = new ObjectName("aiven.kafka.server.tieredstorage.cache:type=chunk-cache");
103104

104105
// When getting a existing chunk from cache
105-
chunkCache.getChunk(OBJECT_KEY_PATH, segmentManifest, 0);
106+
chunkCache.getChunk(OBJECT_KEY_PATH, segmentManifest, 0, -1, Optional.empty(), Optional.empty());
106107

107108
// check cache size increases after first miss
108109
assertThat(MBEAN_SERVER.getAttribute(objectName, "cache-size-total"))
109110
.isEqualTo(1.0);
110111

111-
chunkCache.getChunk(OBJECT_KEY_PATH, segmentManifest, 0);
112+
chunkCache.getChunk(OBJECT_KEY_PATH, segmentManifest, 0, -1, Optional.empty(), Optional.empty());
112113

113114
// Then the following metrics should be available
114115
assertThat(MBEAN_SERVER.getAttribute(objectName, "cache-hits-total"))

0 commit comments

Comments
 (0)