memory.py

#!/usr/bin/env python3
#
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""Utilities for measuring memory utilization for a process."""
import time
import sys

from subprocess import run, CalledProcessError, PIPE
from threading import Thread

MEMORY_COP_TIMEOUT = 5

def usage():
    print('Usage: %s <pid> <guest memory size>' % sys.argv[0])
    sys.exit(1)

def main():
    if len(sys.argv) < 3:
        usage()
    pid, _memory = sys.argv[1:3]
    memory = int(_memory)
    
    pmap_cmd = 'pmap -xq {}'.format(pid)
    while True:
        mem_total = 0

        try:
            pmap_out = run(
                pmap_cmd,
                shell=True,
                check=True,
                stdout=PIPE
            ).stdout.decode('utf-8').split('\n')
        except CalledProcessError:
            break

        for line in pmap_out:
            tokens = line.split()
            if not tokens:
                break
            try:
                total_size = int(tokens[1])
                rss = int(tokens[2])
#                print('  RSS %s total size %s (vs %s)' % (rss, total_size, memory * 1024))
                
            except ValueError:
                # This line doesn't contain memory related information.
                continue
            if total_size >= memory * 1024:
                # This is the guest's memory region.
                # TODO Check for the address of the guest's memory instead.
                continue
            mem_total += rss
#            print('      Footprint %s' % mem_total)

        if not mem_total:
            return

        print('Cloud Hypervisor VMM memory footprint: %s MB' % (mem_total / 1024))
        
        time.sleep(MEMORY_COP_TIMEOUT)


if __name__ == '__main__':
    main()