如何限制JVM线程内存使用？基于Rhino的JS执行场景

在使用Mozilla Rhino为用户提供JavaScript自定义能力时，用户代码可能通过无限创建对象（如示例中的无限往数组push元素）导致JVM OOM。由于性能考量，不希望将用户JS放到独立JVM进程中运行，因此需要实现单线程级别的内存限制方案。

用户示例代码：

test();

function test(){
   var arr = [];
   while(true){
       arr.push("test string");
   }
}

核心思路：拦截Rhino对象创建+线程级内存跟踪 #

Rhino将JS对象编译为Java对象执行，因此可以通过扩展Rhino的核心类，结合线程本地存储（ThreadLocal）跟踪当前执行上下文的内存使用量，达到阈值时终止脚本执行。以下是具体可行的方案：

1. 自定义ContextFactory实现内存监控 #

利用Rhino的ContextFactory扩展执行上下文，添加内存计数器和定期检查逻辑：

• 通过ThreadLocal存储当前线程的已用内存量
• 设置指令观察阈值，每执行一定数量的JS指令就触发内存检查
• 超过阈值时抛出ScriptException终止脚本

示例代码：

public class MemoryLimitedContextFactory extends ContextFactory {
    private final long memoryLimitBytes;
    private final ThreadLocal<Long> threadUsedMemory = new ThreadLocal<>();

    public MemoryLimitedContextFactory(long memoryLimitBytes) {
        this.memoryLimitBytes = memoryLimitBytes;
    }

    @Override
    protected Context makeContext() {
        Context ctx = super.makeContext();
        // 初始化当前线程的内存计数器
        threadUsedMemory.set(0L);
        // 每执行1000条JS指令触发一次内存检查
        ctx.setInstructionObserverThreshold(1000);
        ctx.setInstructionObserver((context, instructionCount) -> {
            long currentUsed = threadUsedMemory.get();
            if (currentUsed > memoryLimitBytes) {
                throw new ScriptException("脚本内存使用超出限制");
            }
        });
        return ctx;
    }

    // 供对象创建逻辑调用，累加内存使用量
    public void addUsedMemory(long size) {
        threadUsedMemory.set(threadUsedMemory.get() + size);
    }
}

2. 扩展Rhino原生对象拦截内存分配 #

针对JS中常用的对象（如数组、字符串），扩展Rhino的原生实现，在对象创建/修改时估算内存并更新计数器：

• 以NativeArray为例，重写push方法，每次添加元素时估算元素内存并累加
• 对字符串、对象等其他类型，实现近似的内存估算逻辑（无需绝对精确，保证趋势监控即可）

示例代码（扩展NativeArray）：

public class MemoryLimitedNativeArray extends NativeArray {
    private final MemoryLimitedContextFactory contextFactory;

    public MemoryLimitedNativeArray(MemoryLimitedContextFactory factory, Object[] elements) {
        super(elements);
        this.contextFactory = factory;
    }

    @Override
    public Object push(Context cx, Scriptable thisObj, Object[] args, Function funObj) {
        for (Object arg : args) {
            // 估算元素的内存占用（近似值）
            long elementSize = estimateObjectSize(arg);
            contextFactory.addUsedMemory(elementSize);
        }
        return super.push(cx, thisObj, args, funObj);
    }

    private long estimateObjectSize(Object obj) {
        if (obj instanceof String) {
            // 字符串内存估算：对象头(16字节) + char数组长度*2 + 数组对象头(16字节)
            String str = (String) obj;
            return 16 + (long) str.length() * 2 + 16;
        } else if (obj instanceof NativeObject) {
            // 估算NativeObject大小：对象头 + 属性数量*单属性开销
            return 16 + ((NativeObject) obj).getIds().length * 32;
        } else {
            // 基本类型或其他对象默认估算32字节
            return 32;
        }
    }
}

3. 结合Java Instrumentation API提升内存估算精度 #

如果需要更准确的内存计算，可以使用Java的Instrumentation API：

• 通过getObjectSize方法获取单个对象的精确大小（注意：该值不包含对象引用的其他对象，需递归计算嵌套对象大小）
• 同样结合ThreadLocal累加内存使用量，达到阈值时终止脚本
• 注意：该方案需要在JVM启动时添加agent（如-javaagent:xxx.jar），部署稍复杂，但内存计算更准确

4. 补充：拦截无限循环的辅助逻辑 #

除了内存监控，还可以通过Context.setMaximumInstructionCount设置脚本允许执行的最大指令数，防止无限循环导致的内存持续增长：

ctx.setMaximumInstructionCount(1000000); // 限制最多执行100万条指令

方案选型建议 #

• 性能优先：优先使用自定义ContextFactory+原生对象扩展的方案，近似内存估算带来的性能开销极小
• 精度优先：结合Instrumentation API实现更准确的内存统计，但需承担agent部署的复杂度
• 简单快速：直接设置MaximumInstructionCount可以快速拦截无限循环类的恶意脚本，作为内存监控的补充

内容的提问来源于stack exchange，提问作者zhufeizzz