diff --git a/src/main/java/com/jozufozu/flywheel/api/struct/StructWriter.java b/src/main/java/com/jozufozu/flywheel/api/struct/StructWriter.java index beb89d025..5b1e0bd33 100644 --- a/src/main/java/com/jozufozu/flywheel/api/struct/StructWriter.java +++ b/src/main/java/com/jozufozu/flywheel/api/struct/StructWriter.java @@ -3,11 +3,11 @@ package com.jozufozu.flywheel.api.struct; import com.jozufozu.flywheel.api.instancer.InstancedPart; /** - * StructWriters can quickly consume many instances of S and write them to some backing buffer. + * StructWriters can quickly consume many instances of S and write them to some memory address. */ public interface StructWriter { /** - * Write the given struct to given memory address. + * Write the given struct to the given memory address. */ void write(long ptr, S struct); } diff --git a/src/main/java/com/jozufozu/flywheel/backend/memory/FlwMemoryTracker.java b/src/main/java/com/jozufozu/flywheel/backend/memory/FlwMemoryTracker.java index c398eb223..80d1c486b 100644 --- a/src/main/java/com/jozufozu/flywheel/backend/memory/FlwMemoryTracker.java +++ b/src/main/java/com/jozufozu/flywheel/backend/memory/FlwMemoryTracker.java @@ -21,50 +21,40 @@ public class FlwMemoryTracker { public static MemoryBlock mallocBlock(long size) { MemoryBlock block = new MemoryBlockImpl(malloc(size), size); - cpuMemory += block.size(); + _allocCPUMemory(block.size()); return block; } public static MemoryBlock mallocBlockTracked(long size) { - TrackedMemoryBlockImpl block = new TrackedMemoryBlockImpl(malloc(size), size); - block.cleanable = CLEANER.register(block, () -> { - if (!block.isFreed()) { - block.free(); - } - }); - cpuMemory += block.size(); + MemoryBlock block = new TrackedMemoryBlockImpl(malloc(size), size, CLEANER); + _allocCPUMemory(block.size()); return block; } @Deprecated public static ByteBuffer mallocBuffer(int size) { ByteBuffer buffer = MemoryUtil.memByteBuffer(malloc(size), size); - cpuMemory += buffer.capacity(); + _allocCPUMemory(buffer.capacity()); return buffer; } public static long calloc(long num, long size) { long ptr = MemoryUtil.nmemCalloc(num, size); if (ptr == MemoryUtil.NULL) { - throw new OutOfMemoryError("Failed to allocate " + num + " blocks of size " + size + " bytes"); + throw new OutOfMemoryError("Failed to allocate " + num + " elements of size " + size + " bytes"); } return ptr; } public static MemoryBlock callocBlock(long num, long size) { MemoryBlock block = new MemoryBlockImpl(calloc(num, size), num * size); - cpuMemory += block.size(); + _allocCPUMemory(block.size()); return block; } public static MemoryBlock callocBlockTracked(long num, long size) { - TrackedMemoryBlockImpl block = new TrackedMemoryBlockImpl(calloc(num, size), num * size); - block.cleanable = CLEANER.register(block, () -> { - if (!block.isFreed()) { - block.free(); - } - }); - cpuMemory += block.size(); + MemoryBlock block = new TrackedMemoryBlockImpl(calloc(num, size), num * size, CLEANER); + _allocCPUMemory(block.size()); return block; } @@ -78,25 +68,23 @@ public class FlwMemoryTracker { public static MemoryBlock reallocBlock(MemoryBlock block, long size) { MemoryBlock newBlock = new MemoryBlockImpl(realloc(block.ptr(), size), size); - cpuMemory += -block.size() + newBlock.size(); + _freeCPUMemory(block.size()); + _allocCPUMemory(newBlock.size()); return newBlock; } public static MemoryBlock reallocBlockTracked(MemoryBlock block, long size) { - TrackedMemoryBlockImpl newBlock = new TrackedMemoryBlockImpl(realloc(block.ptr(), size), size); - newBlock.cleanable = CLEANER.register(newBlock, () -> { - if (!newBlock.isFreed()) { - newBlock.free(); - } - }); - cpuMemory += -block.size() + newBlock.size(); + MemoryBlock newBlock = new TrackedMemoryBlockImpl(realloc(block.ptr(), size), size, CLEANER); + _freeCPUMemory(block.size()); + _allocCPUMemory(newBlock.size()); return newBlock; } @Deprecated public static ByteBuffer reallocBuffer(ByteBuffer buffer, int size) { ByteBuffer newBuffer = MemoryUtil.memByteBuffer(realloc(MemoryUtil.memAddress(buffer), size), size); - cpuMemory += -buffer.capacity() + newBuffer.capacity(); + _freeCPUMemory(buffer.capacity()); + _allocCPUMemory(newBuffer.capacity()); return newBuffer; } @@ -106,13 +94,13 @@ public class FlwMemoryTracker { public static void freeBlock(MemoryBlock block) { free(block.ptr()); - cpuMemory -= block.size(); + _freeCPUMemory(block.size()); } @Deprecated public static void freeBuffer(ByteBuffer buffer) { free(MemoryUtil.memAddress(buffer)); - cpuMemory -= buffer.capacity(); + _freeCPUMemory(buffer.capacity()); } public static void _allocCPUMemory(long size) { diff --git a/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlock.java b/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlock.java index c78b8c193..9b6ad2cea 100644 --- a/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlock.java +++ b/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlock.java @@ -13,9 +13,7 @@ public sealed interface MemoryBlock permits MemoryBlockImpl { void copyTo(long ptr, long bytes); - default void copyTo(long ptr) { - copyTo(ptr, size()); - } + void copyTo(long ptr); void clear(); diff --git a/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlockImpl.java b/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlockImpl.java index 1b6d984a6..a3fb15a60 100644 --- a/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlockImpl.java +++ b/src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlockImpl.java @@ -40,6 +40,11 @@ sealed class MemoryBlockImpl implements MemoryBlock permits TrackedMemoryBlockIm MemoryUtil.memCopy(this.ptr, ptr, bytes); } + @Override + public void copyTo(long ptr) { + copyTo(ptr, size); + } + @Override public void clear() { MemoryUtil.memSet(ptr, 0, size); diff --git a/src/main/java/com/jozufozu/flywheel/backend/memory/TrackedMemoryBlockImpl.java b/src/main/java/com/jozufozu/flywheel/backend/memory/TrackedMemoryBlockImpl.java index 40b1a15eb..a37ba3590 100644 --- a/src/main/java/com/jozufozu/flywheel/backend/memory/TrackedMemoryBlockImpl.java +++ b/src/main/java/com/jozufozu/flywheel/backend/memory/TrackedMemoryBlockImpl.java @@ -3,10 +3,13 @@ package com.jozufozu.flywheel.backend.memory; import java.lang.ref.Cleaner; final class TrackedMemoryBlockImpl extends MemoryBlockImpl { - Cleaner.Cleanable cleanable; + final CleaningAction cleaningAction; + final Cleaner.Cleanable cleanable; - TrackedMemoryBlockImpl(long ptr, long size) { + TrackedMemoryBlockImpl(long ptr, long size, Cleaner cleaner) { super(ptr, size); + cleaningAction = new CleaningAction(ptr, size); + cleanable = cleaner.register(this, cleaningAction); } @Override @@ -14,25 +17,49 @@ final class TrackedMemoryBlockImpl extends MemoryBlockImpl { return true; } + void freeInner() { + freed = true; + cleaningAction.freed = true; + cleanable.clean(); + } + @Override public MemoryBlock realloc(long size) { MemoryBlock block = FlwMemoryTracker.reallocBlock(this, size); - freed = true; - cleanable.clean(); + freeInner(); return block; } @Override public MemoryBlock reallocTracked(long size) { MemoryBlock block = FlwMemoryTracker.reallocBlockTracked(this, size); - freed = true; - cleanable.clean(); + freeInner(); return block; } @Override public void free() { - cleanable.clean(); - freed = true; + FlwMemoryTracker.freeBlock(this); + freeInner(); + } + + static class CleaningAction implements Runnable { + final long ptr; + final long size; + + boolean freed; + + CleaningAction(long ptr, long size) { + this.ptr = ptr; + this.size = size; + } + + @Override + public void run() { + if (!freed) { + FlwMemoryTracker.free(ptr); + FlwMemoryTracker._freeCPUMemory(size); + } + } } }