改进nwipe工具：解决AES-CTR PRNG多线程并行化性能骤降问题

我正在改进nwipe工具，实现基于AES-128计数器模式（AES-CTR）的PRNG，生成高质量随机数用于安全擦除HDD和SSD。原单核心实现性能可达200MB/s，但用Pthreads并行化后性能骤降至约15MB/s，求问题原因及高效多核利用方案。

单核心实现代码 #

int nwipe_aes_ctr_prng_read(NWIPE_PRNG_READ_SIGNATURE) {
    u8* restrict bufpos = buffer;
    size_t words = count / SIZE_OF_AES_CTR_PRNG;

    for(size_t ii = 0; ii < words; ++ii) {
        aes_ctr_prng_genrand_uint128_to_buf((aes_ctr_state_t*) *state, bufpos);
        bufpos += 16; // Move to the next block
    }

    // Handle remaining bytes if count is not a multiple of SIZE_OF_AES_CTR_PRNG
    const size_t remain = count % SIZE_OF_AES_CTR_PRNG;
    if(remain > 0) {
        unsigned char temp_output[16]; // Temporary buffer for the last block
        aes_ctr_prng_genrand_uint128_to_buf((aes_ctr_state_t*) *state, temp_output);
        memcpy(bufpos, temp_output, remain);
    }

    return 0; // Success
}

并行化尝试代码 #

typedef struct {
    aes_ctr_state_t* state;
    u8* buffer;
    size_t start;
    size_t end;
} prng_thread_arg_t;

void* nwipe_aes_ctr_prng_read_thread(void* arg) {
    prng_thread_arg_t* thread_arg = (prng_thread_arg_t*)arg;
    aes_ctr_state_t* state = thread_arg->state;
    u8* buffer = thread_arg->buffer + thread_arg->start;
    size_t words = (thread_arg->end - thread_arg->start) / SIZE_OF_AES_CTR_PRNG;

    for(size_t ii = 0; ii < words; ++ii) {
        aes_ctr_prng_genrand_uint128_to_buf(state, buffer);
        buffer += SIZE_OF_AES_CTR_PRNG;
    }

    return NULL;
}

int nwipe_aes_ctr_prng_read(NWIPE_PRNG_READ_SIGNATURE) {
    int num_threads = 8; // Adjustable based on requirements
    pthread_t threads[num_threads];
    prng_thread_arg_t thread_args[num_threads];

    size_t total_words = count / SIZE_OF_AES_CTR_PRNG;
    size_t words_per_thread = total_words / num_threads;

    for(int i = 0; i < num_threads; i++) {
        size_t start = i * words_per_thread * SIZE_OF_AES_CTR_PRNG;
        size_t end = (i + 1) * words_per_thread * SIZE_OF_AES_CTR_PRNG;

        if(i == num_threads - 1) {
            end = total_words * SIZE_OF_AES_CTR_PRNG; // Correct end calculation
        }

        thread_args[i].state = (aes_ctr_state_t*)*state;
        thread_args[i].buffer = buffer;
        thread_args[i].start = start;
        thread_args[i].end = end;

        pthread_create(&threads[i], NULL, nwipe_aes_ctr_prng_read_thread, &thread_args[i]);
    }

    for(int i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }

    // Remaining bytes handling omitted for brevity
    return 0;
}

通用随机数生成函数 #

void aes_ctr_prng_genrand_uint128_to_buf(aes_ctr_state_t* state, unsigned char* bufpos) {
    CRYPTO_ctr128_encrypt(bufpos, bufpos, 16, &state->aes_key, state->ivec, state->ecount, &state->num, (block128_f) AES_encrypt);
    next_state(state);
}

• 共享状态的缓存一致性风暴：所有线程共用同一个aes_ctr_state_t，每次调用aes_ctr_prng_genrand_uint128_to_buf都会修改ivec、num等字段。多个线程同时读写这些变量会触发缓存行失效，每个核心都要重新从主内存加载状态，导致大量内存等待开销，完全抵消多核并行的收益。
• 违背CTR模式并行特性：AES-CTR的每个加密块完全独立，不需要依赖前一块的输出，但当前并行实现强制所有线程共享计数器状态，相当于串行执行计数器递增，完全浪费了CTR模式的并行潜力。

正确的思路是为每个线程分配独立的CTR状态，预先计算好每个线程负责的计数器起始值，让线程独立生成对应块的随机数，彻底避免共享状态竞争。

1. 修改线程参数结构 #

每个线程拥有独立的CTR状态副本：

typedef struct {
    aes_ctr_state_t state; // 线程专属状态，避免共享
    u8* buffer;
    size_t start;
    size_t end;
} prng_thread_arg_t;

2. 初始化线程独立状态 #

主线程中基于原始状态为每个线程计算起始计数器值：

#include <unistd.h> // 用于sysconf

int nwipe_aes_ctr_prng_read(NWIPE_PRNG_READ_SIGNATURE) {
    // 获取系统可用核心数，替代硬编码的8
    int num_threads = sysconf(_SC_NPROCESSORS_ONLN);
    pthread_t threads[num_threads];
    prng_thread_arg_t thread_args[num_threads];

    size_t total_blocks = count / SIZE_OF_AES_CTR_PRNG;
    size_t blocks_per_thread = total_blocks / num_threads;
    aes_ctr_state_t* base_state = (aes_ctr_state_t*)*state;

    for(int i = 0; i < num_threads; i++) {
        // 复制原始密钥和IV，生成线程独立状态
        memcpy(&thread_args[i].state, base_state, sizeof(aes_ctr_state_t));
        
        // 计算当前线程负责的起始块偏移，更新计数器
        increment_ctr_by(&thread_args[i].state.ivec, i * blocks_per_thread);

        thread_args[i].buffer = buffer;
        thread_args[i].start = i * blocks_per_thread * SIZE_OF_AES_CTR_PRNG;
        // 最后一个线程处理剩余所有块
        thread_args[i].end = (i == num_threads - 1) ? 
            total_blocks * SIZE_OF_AES_CTR_PRNG : 
            (i + 1) * blocks_per_thread * SIZE_OF_AES_CTR_PRNG;

        pthread_create(&threads[i], NULL, nwipe_aes_ctr_prng_read_thread, &thread_args[i]);
    }

    // 等待所有线程完成
    for(int i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }

    // 处理剩余字节
    const size_t remain = count % SIZE_OF_AES_CTR_PRNG;
    if(remain > 0) {
        unsigned char temp_output[16];
        aes_ctr_state_t temp_state = *base_state;
        // 从总块数后的计数器开始生成
        increment_ctr_by(&temp_state.ivec, total_blocks);
        aes_ctr_prng_genrand_uint128_to_buf(&temp_state, temp_output);
        memcpy(buffer + total_blocks * SIZE_OF_AES_CTR_PRNG, temp_output, remain);
        // 更新原始状态，保证后续调用的连续性
        increment_ctr_by(&base_state->ivec, total_blocks + (remain > 0 ? 1 : 0));
    }

    return 0;
}

3. 实现计数器递增函数 #

处理AES CTR的大端计数器递增：

void increment_ctr_by(unsigned char* ctr, size_t increment) {
    size_t carry = increment;
    // 从最后一个字节开始递增（大端格式）
    for(int i = 15; i >= 0 && carry > 0; i--) {
        size_t val = (unsigned char)ctr[i] + carry;
        ctr[i] = val & 0xFF;
        carry = val >> 8;
    }
}

4. 修改线程执行函数 #

使用线程专属状态生成随机数：

void* nwipe_aes_ctr_prng_read_thread(void* arg) {
    prng_thread_arg_t* thread_arg = (prng_thread_arg_t*)arg;
    aes_ctr_state_t* state = &thread_arg->state;
    u8* buffer = thread_arg->buffer + thread_arg->start;
    size_t blocks = (thread_arg->end - thread_arg->start) / SIZE_OF_AES_CTR_PRNG;

    for(size_t ii = 0; ii < blocks; ++ii) {
        aes_ctr_prng_genrand_uint128_to_buf(state, buffer);
        buffer += SIZE_OF_AES_CTR_PRNG;
    }

    return NULL;
}

额外优化建议 #

• 启用AES-NI硬件加速：确保OpenSSL编译时启用了AES-NI支持，硬件加速能让单核心AES性能提升数倍，并行后收益更显著。
• 使用大页内存：如果擦除缓冲区较大，配置大页内存减少TLB miss，提升内存访问效率。
• 批量生成优化：可以让每个线程一次性生成多个块，减少函数调用开销，但CTR模式本身已足够高效，此优化优先级较低。

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