Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
import static org.apache.paimon.utils.Preconditions.checkNotNull;

/** Cache manager to cache bytes to paged {@link MemorySegment}s. */
public class CacheManager {
public class CacheManager implements AutoCloseable {

private static final Logger LOG = LoggerFactory.getLogger(CacheManager.class);

Expand All @@ -43,8 +43,13 @@ public class CacheManager {

private final Cache dataCache;
private final Cache indexCache;
private final boolean offHeap;

public CacheManager(MemorySize maxMemorySize, double highPriorityPoolRatio) {
this(maxMemorySize, highPriorityPoolRatio, false);
}

private CacheManager(MemorySize maxMemorySize, double highPriorityPoolRatio, boolean offHeap) {
Preconditions.checkArgument(
highPriorityPoolRatio >= 0 && highPriorityPoolRatio < 1,
"The high priority pool ratio should in the range [0, 1).");
Expand All @@ -58,12 +63,19 @@ public CacheManager(MemorySize maxMemorySize, double highPriorityPoolRatio) {
} else {
this.indexCache = CacheBuilder.newBuilder().maximumWeight(indexCacheSize).build();
}
this.offHeap = offHeap;
LOG.info(
"Initialize cache manager with data cache of {} and index cache of {}.",
"Initialize {} cache manager with data cache of {} and index cache of {}.",
offHeap ? "off-heap" : "heap",
dataCacheSize,
indexCacheSize);
}

public static CacheManager createOffHeap(
MemorySize maxMemorySize, double highPriorityPoolRatio) {
return new CacheManager(maxMemorySize, highPriorityPoolRatio, true);
}

@VisibleForTesting
public Cache dataCache() {
return dataCache;
Expand All @@ -82,7 +94,7 @@ public MemorySegment getPage(CacheKey key, CacheReader reader, CacheCallback cal
k -> {
try {
return new Cache.CacheValue(
MemorySegment.wrap(reader.read(key)), callback);
toMemorySegment(reader.read(key)), callback);
} catch (IOException e) {
throw new RuntimeException(e);
}
Expand All @@ -98,6 +110,23 @@ public void invalidPage(CacheKey key) {
}
}

private MemorySegment toMemorySegment(byte[] bytes) {
if (!offHeap) {
return MemorySegment.wrap(bytes);
}
MemorySegment segment = MemorySegment.allocateOffHeapMemory(bytes.length);
segment.put(0, bytes);
return segment;
}

@Override
public void close() {
dataCache.invalidateAll();
if (indexCache != dataCache) {
indexCache.invalidateAll();
}
}

/** The container for the segment. */
public static class SegmentContainer {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,4 +66,21 @@ void testCaffeineCache() throws Exception {
}
}
}

@Test
void testOffHeapCache() throws Exception {
File file = new File(tempDir.toFile(), "test.off-heap");
assertThat(file.createNewFile()).isTrue();
CacheKey key = CacheKey.forPageIndex(new RandomAccessFile(file, "r"), 0, 0);

try (CacheManager cacheManager = CacheManager.createOffHeap(MemorySize.ofBytes(10), 0)) {
MemorySegment segment =
cacheManager.getPage(key, ignored -> new byte[] {1, 2, 3}, ignored -> {});

assertThat(segment.isOffHeap()).isTrue();
byte[] bytes = new byte[3];
segment.get(0, bytes);
assertThat(bytes).containsExactly(1, 2, 3);
}
}
}
Loading