#!/usr/bin/env python3
"""
内存压力测试脚本 - 用于演示swap使用
该脚本会分配大量内存来触发系统使用swap
"""

import time
import sys

def memory_stress_test():
    """分配内存来触发swap使用"""
    print("开始内存压力测试...")
    print("当前内存和swap使用情况:")
    
    # 显示当前内存状态
    with open('/proc/meminfo', 'r') as f:
        meminfo = f.read()
        for line in meminfo.split('\n')[:6]:
            if line.strip():
                print(f"  {line}")
    
    # 分配内存列表
    memory_blocks = []
    block_size = 50 * 1024 * 1024  # 50MB per block
    
    try:
        for i in range(20):  # 尝试分配总共约1GB
            print(f"\n分配第 {i+1} 个内存块 (50MB)...")
            # 分配50MB的内存块
            block = bytearray(block_size)
            # 写入一些数据确保内存真正分配
            for j in range(0, block_size, 1024):
                block[j] = j % 256
            memory_blocks.append(block)
            
            # 每分配5个块后检查内存状态
            if (i + 1) % 5 == 0:
                print(f"\n已分配 {(i+1)*50}MB，当前内存状态:")
                time.sleep(2)  # 等待系统调整
                
    except MemoryError:
        print("\n内存分配失败，已达到系统限制")
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    
    print(f"\n测试完成，共分配了 {len(memory_blocks)*50}MB 内存")
    print("等待5秒后释放内存...")
    time.sleep(5)
    
    # 释放内存
    memory_blocks.clear()
    print("内存已释放")
    
    print("\n测试后的内存状态:")
    with open('/proc/meminfo', 'r') as f:
        meminfo = f.read()
        for line in meminfo.split('\n')[:6]:
            if line.strip():
                print(f"  {line}")

if __name__ == "__main__":
    memory_stress_test()