【Java】记录一次 JMH 性能测试与调优

有了前两次“安全”方面的处理经验，部门决定接下来“安全”方面的工作都全部交给我来处理（被赶鸭子上架）。于是...公司的一个关于“加解密工具”安全审计工作被安排到我这边了。

在配合专业团队进行为期一周的代码审计后发现“加解密工具”存在结构性漏洞，后续需要对代码进行修复（本人参与这部分工作）。由于代码是用 Java 编写的，因此修复完成后选择采用 JMH 对代码性能进行测试。

JMH 是什么？

JMH（Java Microbenchmark Harness）是一个专为 Java 代码设计的微基准测试框架。它是由OpenJDK 和 Oracle 官方开发并维护的工具，主要用于在方法级别对代码片段进行高性能和低延迟的基准测试。JMH 的设计目的是为了帮助开发者理解和优化 Java 程序中的关键性能瓶颈。

而且考虑到本次只是针对一部分应用代码进行性能测试（不涉及到网络请求），因此就没有选用 JMeter 来做测试了。具体的使用也是非常简单。首先肯定要在 pom.xml 中添加对应的依赖和插件，如下图：

...

<dependency>
  <groupId>org.openjdk.jmh</groupId>
  <artifactId>jmh-core</artifactId>
  <version>${jmh.version}</version>
</dependency>
<dependency>
  <groupId>org.openjdk.jmh</groupId>
  <artifactId>jmh-generator-annprocess</artifactId>
  <version>${jmh.version}</version>
</dependency>

...

<build>
  <plugins>
    ...
    <plugin>
      <groupId>pw.krejci</groupId>
      <artifactId>jmh-maven-plugin</artifactId>
      <version>0.2.2</version>
      <executions>
        <execution>
          <goals>
            <goal>benchmark</goal>
          </goals>
        </execution>
      </executions>
    </plugin>
  </plugins>
</build>

注意了，mvnrepository.com 获取的依赖事例一般都是带有 <scope>test</scope> 标签的，这个在你打包后用 main 方法运行时就会报错说找不到 main 方法，因此建议统一将 <scope> 标签去掉。

接下来就可以编写测试代码了，如下图：

// 采用吞吐量压测模式
@BenchmarkMode(Mode.Throughput)
// 输出的时间单位是秒（因为压测时间较长）
@OutputTimeUnit(TimeUnit.SECONDS)
// 叉分次数，只运行一次fork，减少初始化开销
@Fork(1)
// 线程数（旧款的 Intel Mac CPU i5 四核）
@Threads(10)
// 5次预热迭代，每次5秒， 由于每个迭代都需要做一次 setupPerIteration ，因此预热时间足够了
@Warmup(iterations = 5, time = 5, timeUnit = TimeUnit.SECONDS)
// 单次测量时间为10分钟，以捕捉长时间运行的性能特征
@Measurement(iterations = 1, time = 600, timeUnit = TimeUnit.SECONDS)
// 自动为每个线程创建独立的状态，在本例中每个线程都会有自己的 testData 实例
@State(Scope.Thread)
public class AhEUtilTest {

    // 定义数据大小
    private static final int DATA_SIZE = 1024 * 1024;
    // 定义安全随机数生成器，用于填充测试数据
    private static final SecureRandom SECURE_RANDOM = new SecureRandom();
    // 测试数据数组
    private byte[] testData;
    // 记录加密后数据
    private String encryptedData;

    // 在每次fork开始前调用此方法来初始化状态
    @Setup
    public void setup() {
        testData = new byte[DATA_SIZE];
        // 使用安全随机数生成器填充测试数据
        SECURE_RANDOM.nextBytes(testData);
    }

    // 加密压测
    @Benchmark
    public void benchmarkEncryption() throws Exception {
        // AhEUtil 是加解密工具（已封装），这里调用 encrypt 加密方法
        AhEUtil.encrypt(new String(testData, StandardCharsets.UTF_8));
    }

    // 每次迭代前都调用一次（耗时不算入测试结果）
    @Setup(Level.Iteration)
    public void setupPerIteration() throws Exception {
        // 调用一次加密将加密后的动态数据保存到变量供解密时试用
        testData = new byte[DATA_SIZE];
        SECURE_RANDOM.nextBytes(testData);
        encryptedData = AhEUtil.encrypt(new String(testData, StandardCharsets.UTF_8));
    }

    // 解密压测
    @Benchmark
    public void benchmarkDecryption() throws Exception {
        // 将 setupPerIteration 加密后的数据进行解密，调用 decrypt 方法
        String decryptedData = AhEUtil.decrypt(encryptedData);

        // 验证是否解密成功这里加上字符串的一致性判断
        if (!decryptedData.equals(new String(testData, StandardCharsets.UTF_8))) {
            throw new IllegalStateException("Decrypted data does not match original data.");
        }
    }

    // main 方法调用
    public static void main(String[] args) throws RunnerException {
        Options opt = new OptionsBuilder()
                .include(AhEUtilTest.class.getSimpleName())
                .result("ahe_result.json")
                .resultFormat(ResultFormatType.JSON).build();
        new Runner(opt).run();
    }
}

上面的注释应该已经将代码都解释得非常清晰了，至于 JMH 的参数网上实在太多例子了我就不多列举了，本人都是通过其他大神的博客学习的，接下来看看压测结果。

为了将 JMH 结果可视化，我使用了 https://jmh.morethan.io/ 可视化工具，将压测后的 json 导入到这个工具里面就能够直观看到结果，如下图。

额...只有 21 ops/s，难道我又要吃瘪了？虽然我的电脑配置有点差，但也不至于吧...

为了更好地做对比，用公司提供的封装好的 AES 工具类做了一下测试（开始自我安慰），这里稍微调整了 AhEUtilTest 类并尽可能保持与本次加解密工具参数保持一致，调整代码如下图所示：

...
public class AhEUtilTest {

    ...
    private String encryptedKey;

    ...
    @Benchmark
    public void benchmarkEncryption() throws Exception {
        //AhEUtil.encrypt(new String(testData, StandardCharsets.UTF_8));

        // 用雪花算法获取密钥 id（AES 工具调用所需参数）
        encryptedKey = SnowFlakeUtil.get().getItemID(16);
        // 调用 AES 工具进行加密
        nativeAesEncryptForTest(testData,encryptedKey);
    }

    @Setup(Level.Iteration)
    public void setupPerIteration() throws Exception {
        testData = new byte[DATA_SIZE];
        SECURE_RANDOM.nextBytes(testData);
        //encryptedData = AhEUtil.encrypt(new String(testData, StandardCharsets.UTF_8));
       
        // 自动调用一次加密用于缓存密文（AES 工具调用所需参数）
        encryptedKey = SnowFlakeUtil.get().getItemID(16);
        encryptedData = nativeAesEncryptForTest(testData,encryptedKey);

    }

    @Benchmark
    public void benchmarkDecryption() throws Exception {
        //String decryptedData = AhEUtil.decrypt(encryptedData);

        // 调用 AES 解密
        String decryptedData = nativeAesDecryptForTest(encryptedData,encryptedKey);
        if (!decryptedData.equals(new String(testData, StandardCharsets.UTF_8))) {
            throw new IllegalStateException("Decrypted data does not match original data.");
        }
    }

    ...

    private String nativeAesEncryptForTest(byte[] testData,String encryptedKey) throws Exception{
        // 采用 CBC PKCS7Padding 的加密方式，密钥长度 256（跟目标算法一致）
        return AESUtil.encrypt(new String(testData, StandardCharsets.UTF_8), encryptedKey, AESUtil.CRYPTO_CBC,
                                "PKCS7Padding",
                                256, encryptedKey);
    }

    private String nativeAesDecryptForTest(String encryptedDatas,String encryptedKeys) throws Exception{
        // 调用公司 AES 工具解密
        return AESUtil.decrypt(encryptedDatas, encryptedKeys, AESUtil.CRYPTO_CBC,
                                "PKCS7Padding",
                                256, encryptedKeys);
    }
}

之所以选用 AES 做对比压测，是因为前工程师们在开发这套加解密工具时就是在 AES 算法之外再叠加处理逻辑的，也可以说是基于 AES 加解密的变种吧，因此拿 AES 做比较还算是比较合适的。压测结果如下：

果然又要吃瘪了，原生 AES 加解密加密性能能够去到 49 ops/s，解密也有 50 ops/s。这一半以上的差距，只好重新审查代码并看看有没有继续优化的可能。

经过几天的奋战，将原代码中动态密钥用到的 PBKDF2 算法替换成 HmacSHA256 算法，然后在加密因子中增加了高熵处理。简单来说就是之前每次加密都需要调用 PBKDF2 来创建动态密钥，现在换了一种玩法，将这个需要算力密集的操作通过提前做一次高熵处理给做了，然后动态密钥就由 HmacSHA256 来接管，这样可以有效减少 CPU 算力消耗（当然了这样的逻辑改动肯定要过审）。接下来又做了一次压测，结果如下：

性能直接飙升到加密 47 ops/s，解密 47 ops/s。再看看整个提升的汇总，如下图：

嗯...感觉这样才算是比较正常的，总算是修正过来了。当然了如果是有抗暴力破解和抗彩虹表攻击需求的 PBKDF2 仍然是一个可靠的选择。