Added utility class for limiting memory allocated for VMs.

This commit is contained in:
Florian Nücke
2020-10-05 11:26:27 +02:00
parent 5f1a69efc3
commit d8d568daca
3 changed files with 111 additions and 0 deletions

View File

@@ -0,0 +1,5 @@
package li.cil.oc2;
public final class Config {
public static final long maxAllocatedData = 512 * 1024 * 1024;
}

View File

@@ -0,0 +1,99 @@
package li.cil.oc2.common.sandbox;
import li.cil.oc2.Config;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.HashMap;
import java.util.UUID;
/**
* Cooperative memory allocation limit enforcement.
* <p>
* Call sites must be cooperative and only free claimed memory when actually being sure the
* allocated memory associated with the claim will be garbage collected.
*/
public final class Allocator {
private static final Logger LOGGER = LogManager.getLogger();
private static final HashMap<UUID, Allocation> ALLOCATIONS = new HashMap<>();
private static long allocated;
/**
* Creates a new handle that can be used to claim memory.
*
* @return a new handle.
*/
public static UUID createHandle() {
return UUID.randomUUID();
}
/**
* Tries to claim the specified amount of memory using the specified {@code handle}.
* <p>
* Claimed memory <em>must</em> be returned using {@link #freeMemory(UUID)} to prevent leaks.
*
* @param handle the handle to use for claiming memory.
* @param size the amount of memory to claim.
* @return {@code true} if the memory was successfully claimed; {@code false} otherwise.
*/
public static boolean claimMemory(final UUID handle, final int size) {
if (!checkArgs(handle, size)) {
return false;
}
if (size != 0) {
ALLOCATIONS.put(handle, new Allocation(size));
}
return true;
}
/**
* Frees memory that was claimed using the specified handle.
* <p>
* Using this if there was no memory claimed with ths handle or if the handle has already been
* freed does nothing.
*
* @param handle the handle to release the claimed memory for.
*/
public static void freeMemory(final UUID handle) {
final Allocation allocation = ALLOCATIONS.remove(handle);
if (allocation != null) {
allocated -= allocation.size;
}
}
/**
* Clears all remaining allocations and logs their stack traces.
*/
public static void resetAndCheckLeaks() {
if (allocated > 0) {
LOGGER.error("Not all memory was released; leaked allocation stack traces follow.");
for (final Allocation allocation : ALLOCATIONS.values()) {
LOGGER.error(allocation.stacktrace);
}
}
ALLOCATIONS.clear();
allocated = 0;
}
private static boolean checkArgs(final UUID handle, final int size) {
if (ALLOCATIONS.containsKey(handle)) {
throw new IllegalStateException("Handle is already in use. It must be freed before it can be reused.");
}
if (size < 0) {
throw new IllegalArgumentException();
}
return Config.maxAllocatedData - size >= allocated;
}
private static final class Allocation {
public final int size;
private final StackTraceElement[] stacktrace;
private Allocation(final int size) {
this.size = size;
this.stacktrace = new Throwable().getStackTrace();
}
}
}

View File

@@ -0,0 +1,7 @@
@ParametersAreNonnullByDefault
@MethodsReturnNonnullByDefault
package li.cil.oc2.common.sandbox;
import mcp.MethodsReturnNonnullByDefault;
import javax.annotation.ParametersAreNonnullByDefault;