用于测试记录bsdiff内存占用的python脚本

record_bsdiff_memory.py

#!/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division
from threading import Thread
import subprocess
import psutil
import time
import os

import global_flag

SOFTWARE_DIR = 'software'
PATCH_DIR = 'patch'


def format_size(memory):

    flag = False
    if memory < 0:
        flag = True

    memory = abs(memory)

    if memory > 1000000000:
        m_raw = memory
        m_show = round(memory / 1000000000, 2)
        m_unit = str(m_show) + 'G'
    elif memory > 1000000:
        m_raw = memory
        m_show = round(memory / 1000000, 2)
        m_unit = str(m_show) + 'M'
    elif memory > 1000:
        m_raw = memory
        m_show = round(memory / 1000, 2)
        m_unit = str(m_show) + 'K'
    else:
        m_raw = memory
        m_show = round(memory, 2)
        m_unit = str(m_show)

    if flag:
        m_unit = '-' + m_unit

    return (m_raw, m_show, m_unit)


def get_used_memory():

    info = psutil.virtual_memory()
    used = info[3]

    return used


def get_max_memory(normal_memory, old, new, patch, old_size, new_size):

    # nmr: normal memory raw ...
    nmr, nms, nmu = format_size(normal_memory)
    max_memory = 0

    #print("start log max memory")
    while True:

        used_memory = get_used_memory()

        if used_memory > max_memory:
            max_memory = used_memory

        time.sleep(0.5)

        if global_flag.get():
            break

    # mmr: max memory raw
    mmr, _, mmu = format_size(max_memory)
    max_used, _, max_used_unit = format_size(mmr - nmr)

    expect = expect_max_memory_used(old_size, new_size)
    _, _, expect_unit = format_size(expect)
    _, _, o1 = format_size(max_used - expect)

    print("noraml: {}, max: {}, bsdiff_max_used: {}".format(nmu, mmu, max_used_unit))
    print("expect_max_memory_used: {}, O(1): {}".format(expect_unit, o1))
    return



def start_get_max_memory_with_thread(meta_info):

    normal_memory = meta_info['normal_memory']
    old = meta_info['old']
    new = meta_info['new']
    patch = meta_info['patch']
    old_size = meta_info['old_size']
    new_size = meta_info['new_size']

    thread = Thread(target=get_max_memory, args=(normal_memory, old, new, patch, old_size, new_size))
    thread.daemon = True
    thread.start()

    return thread



def expect_max_memory_used(old_size, new_size):

    return max(old_size * 17, old_size * 9 + new_size)



def start_bsdiff(meta_info):

    #print("start bsdiff")

    old_path = meta_info['old_path']
    new_path = meta_info['new_path']
    patch_path = meta_info['patch_path']

    p = subprocess.Popen(["bsdiff", old_path, new_path, patch_path], stdout=subprocess.PIPE)
    output, err = p.communicate()

    # when bsdiff done, change the flag
    global_flag.change()



def main():

    file_list = os.listdir(SOFTWARE_DIR)
    #print(file_list)
    file_list = file_list[::-1]

    for i in range(len(file_list) - 1):

        print("------------ start({}/{}) ------------".format(i+1, len(file_list)))

        normal_memory = get_used_memory()
        old_name = file_list[i]
        new_name = file_list[i+1]
        patch_name = ''.join(old_name.split('.')[:-1]) + "_upto_" + ''.join(new_name.split('.')[:-1])
        print(old_name, new_name, patch_name)

        old_path = os.path.join(SOFTWARE_DIR, old_name)
        new_path = os.path.join(SOFTWARE_DIR, new_name)
        patch_path = os.path.join(PATCH_DIR, patch_name)

        # software size
        old_size = os.path.getsize(old_path)
        new_size = os.path.getsize(new_path)
        _, _, old_size_unit = format_size(old_size)
        _, _, new_size_unit = format_size(new_size)

        meta_info = {
            "normal_memory": normal_memory,
            "old": old_name,
            "new": new_name,
            "patch": patch_name,
            "old_path": old_path,
            "new_path": new_path,
            "patch_path": patch_path,
            "old_size": old_size,
            "new_size": new_size
        }

        # start logging memory used
        thread = start_get_max_memory_with_thread(meta_info)

        # wait memory function launch
        time.sleep(1)

        # start bsdiff and timeit
        t1 = time.time()
        start_bsdiff(meta_info)
        t2 = time.time()

        thread.join()
        global_flag.change()


        print("old size: {}, new size: {}".format(old_size_unit, new_size_unit))
        print("bsdiff used time: {}s".format(round(int(t2-t1), 2)))

        time.sleep(5)



if __name__ == "__main__":

    main()

global_flag 包在这里： https://gist.github.com/sincerefly/4fa11ec77302557c2626af6a84bbfe32