square-mem

#!/usr/bin/python3
import os
import time
import math
import optparse
import subprocess

import curses
import squarify

from helpers_categorizecmd import categorize_cmd

kilo = 1024
mega = kilo*kilo
giga = mega*kilo




### Helpers

def kmg(amount,kilo=1000, append='',thresh=15, nextup=0.9, rstrip0=True, extradigits=0, i_for_1024=True):
    """ For more easily skimmable sizes

        e.g.
             kmg(3429873278462) == '3.4T'
             kmg(342987327)     == '343M'
             kmg(34298)         == '34K'

             '%sB'%kmg(2342342324)                           == '2.3GB'
             '%sB'%kmg(2342342324, kilo=1024)                == '2.2GiB'
             '%sB'%kmg(2342342324, kilo=1024, extradigits=1) == '2.18GiB'
             '%sB'%kmg(19342342324, kilo=1024)                == '18GiB'
             '%sB'%kmg(19342342324, kilo=1024, extradigits=1) == '18GiB'  (because of rstrip0)

        Decimal/SI kilos by default, so useful beyond bytes.
        Specify kilo=1024 if you want binary kilos. By default this also adds the i.

        thresh is the controls where we take one digit away, e.g. for 1.3GB but 16GB.
        Default is at 15 which is entirely arbitrary. 
        Disable using None.

        nextup makes us switch to the next higher up earlier, e.g. 700GB but 0.96TB
        Disable using None.
 
        extradigits=1 (or maybe more) to unconditionally see a less-rounded number
        (though note rstrip can still apply)

        rstrip0     whether to take off '.0' if present (defaults to true)
        append      is mostly meant for optional space between number and unit.
    """
    ret = None

    mega = kilo*kilo
    giga = mega*kilo
    tera = giga*kilo
    peta = tera*kilo
    exa   = peta*kilo
    zetta = exa*kilo
    yotta = zetta*kilo

    if nextup==None:
        nextup = 1.0
    if thresh==None:
        thresh = 1000
    nextup = float(nextup)

    # Yes, could be handled a bunch more more compactly (and used to be)

    showdigits=0
    if abs(amount) < nextup*kilo: # less than a kilo; omits multiplier and i
        showval = amount
    else:
        for csize, mchar in ( (peta, 'P'),
                              (tera, 'T'),
                              (giga, 'G'),
                              (mega, 'M'),
                              (kilo, 'K'),
                              #(exa,  'E'),# exa, zetta, yotta is shown as peta amounts. Too large to comprehend anyway.
                              #(zeta, 'Z'),
                              #(yotta,'Y'),
           ):
            if abs(amount) > nextup*csize:
                showval = amount/float(csize)
                if showval<thresh:
                    showdigits = 1 + extradigits
                else:
                    showdigits = 0 + extradigits
                append += mchar
                if i_for_1024 and kilo==1024:
                    append += 'i'
                break

    ret = ("%%.%df"%(showdigits))%showval

    if rstrip0:
        if '.' in ret:
            ret=ret.rstrip('0').rstrip('.')

    ret+=append

    return ret

def parse_kmg(str, kilo=1024):
    " E.g. '1MB' -> 1048576.   Quick and dirty implementation, could stand cleaning "
    #if type(str) is int:
    #    return str
    try:
        ns = str.rstrip('kmgtbKMGTBiI')
        ret = float(ns)
        sl = str.lower()
        if 'k' in sl:
            ret *= kilo
        if 'm' in sl:
            ret *= kilo*kilo
        if 'g' in sl:
            ret *= kilo*kilo*kilo
        if 't' in sl:
            ret *= kilo*kilo*kilo*kilo
        ret=int(ret)
        return ret
    except Exception as e:        
        print( "Didn't understand value %r"%str )
        print( e )
        raise


def fetch_procmem():
    ' mostly used for totals, which smem and ps do not report ' 
    f = open('/proc/meminfo')
    lines = f.readlines()
    f.close()
    mi = {}
    for line in lines:
        if ':' in line:
            var, val = line.strip().split(':',1)
            val = val.strip()
            if 'kB' in val:
                val = 1024*int(val.replace('kB','').strip()) # hacky hacky :)
                mi[var]=val
    return mi


def fetch_smem():
    ' returns a dict:  pid -> (cmd, memsize) '
    pid_cmd_mem = {}
    cmd = ["/usr/bin/smem", '-n','-H', '-c','pid pss swap name']
    p = subprocess.Popen(cmd,shell=False, stdout=subprocess.PIPE, encoding='utf8')
    out,_ = p.communicate()
    #print out
    for line in out.splitlines():
        ll = line.strip().split(None,3)
        pid, pss, swap, cmd = ll
        pss = int(pss)*1024
        swap = int(swap)*1024
        # clean up cmd we display (if necessary)
        if ':' in cmd:
            cmd = cmd.split(':')[0]
        if '/' in cmd:
            cmd = cmd.split('/')[-1]
        if ' ' in cmd:
            cmd = cmd.split(' ')[0]
        pid_cmd_mem[pid] = (cmd,pss)
    return pid_cmd_mem


def fetch_ps():
    ' returns a dict:  pid -> (cmd, memsize) '
    pid_cmd_mem = {}
    p = subprocess.Popen("ps --no-header -eo pid,rss,comm",shell=True, stdout=subprocess.PIPE, encoding='utf8')          # rss is non-swapped physical
    out,_ = p.communicate()
    pid_cmd_mem = {}
    for line in out.splitlines():
        ll = line.strip().split()
        pid = ll[0]
        rss = int(ll[1])*1024
        cmd = ' '.join(ll[2:])
        pid_cmd_mem[pid] = (cmd,rss)
    return pid_cmd_mem



def main():
    p=optparse.OptionParser()
    p.add_option('-s','--accurate-shared',  default=False,   action="store_true", dest="ashare",
                 help="Count shared memory more accurately. Slow, only really necessary when you have much of it. Uses smem, which you need to install. Falls back on not doing this, i.e. using ps.")
    p.add_option('-S','--show-swap',        default=False,   action="store_true", dest="showswap",
                 help="Add used and free swap")
    p.add_option('-g','--group-small',      default='.005',   action="store", dest="smallthresh",
                 help="Anything smaller than this is summed into a single item called (sumsmaller). Can be either a size like 200M, or a fraction interpreted that will be multiplied by MemTotal. Default is 0.005 (meaning half a percent)")
    p.add_option('-i','--interval',  default=0.5,   action="store", dest="interval",
                 help="Interval (sleep time) between measuring/showing. (note that something higher will be used when using smem)")    
    options,args = p.parse_args()
    interval     = float(options.interval)
    try:
        smallthresh = float(options.smallthresh)
        if smallthresh > 1:
            raise ValueError('not a fraction')
    except ValueError:
        smallthresh  = parse_kmg(options.smallthresh)
    
    if options.ashare:
        try:
            p = subprocess.Popen(["/usr/bin/smem",'-n','-H', '-c','pid name pss swap'],shell=False, stdout=subprocess.PIPE, encoding='utf8')
            p.wait()
        except OSError:
            options.ashare = False
            print( "WARNING: not doing accurate shared memory reporting (smem not present, falling back to ps)")
        time.sleep(2)
    else:
        print( "INFO: Using faster reporting - that doesn't show shared memory so well")
        
    if options.ashare and os.getuid()!=0:
        print( "WARNING: Using smem but not running as root / via sudo, so memory summary will exclude system processes")
        time.sleep(5)
    
    screen = curses.initscr()
    try:
        curses.noecho()
        curses.curs_set(0)

        #curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_WHITE)
        curses.start_color()
        curses.use_default_colors()
        num_inited_cols = 0

        initcolnum = min(curses.COLOR_PAIRS, curses.COLORS) # I should check whether this makes any real sesne
            
        if curses.COLORS in (8,256):
            #import syslog
            while num_inited_cols < curses.COLOR_PAIRS:
                for bg in range(8):
                    if bg in (7,6,5,3,2): # brighter colors
                        fg=0 # black
                    else: # darker colors (black, red, blue)
                        fg = 7 # white
                        # CONSIDER: also using bright white, 15

                    if fg!=bg:
                        num_inited_cols += 1
                        curses.init_pair(num_inited_cols, fg, bg)

                    if num_inited_cols > curses.COLOR_PAIRS:
                        break

            #syslog.syslog("inited %d color pairs"%num_inited_cols)

        while True:
            # deal with WINCHes
            if curses.is_term_resized(curses.COLS,curses.LINES):
                y, x = screen.getmaxyx()
                screen.clear()
                curses.resizeterm(y, x)
            
            ### get memory per process
            if options.ashare:
                try:
                    pid_cmd_mem = fetch_smem()
                except OSError: # shouldn't happen since we checked before starting, but hey
                    pid_cmd_mem = fetch_ps()
            else:
                pid_cmd_mem = fetch_ps()

                    
            ### sort that into more human-parseable categories
            pername = {}
            for pid in pid_cmd_mem:
                cmd,rss = pid_cmd_mem[pid]
                name = cmd

                if rss==0.0:
                    continue
                name = categorize_cmd(cmd)
                if name not in pername:
                    pername[name]=0
                pername[name]+=rss
                #syslog.syslog( '%s --> %s  (%s)'%(cmd,name, kmg(rss)) )
                
                
            ### augment with some other useful things
            mi = fetch_procmem()
            MT  = mi['MemTotal']
            MF  = mi['MemFree']

            pername['(unused)'] = MF

            CD  = mi['Cached']
            pername['(caches, usable)']  = CD + mi['SReclaimable'] 


            #pername['(kernel)'] = mi['Buffers'] + mi['Dirty'] + mi['Writeback'] + mi['SUnreclaim']  # misc system stuff
            # split out because sometimes SUnreclaim is large (e.g. ZFS ARC).   They may well fall into (sumsmaller) anyway
            pername['(k/buffers)'] = mi['Buffers']
            pername['(k/dirty)'] = mi['Dirty']
            pername['(k/writeback)'] = mi['Writeback']
            pername['(k/slab unreclaim)'] = mi['SUnreclaim']  # note: SReclaimable is counted towards caches
            # CONSIDER
            # - using /proc/spl/kstat/zfs/arcstats to split off ZFS, which if used will make SUnreclaim large due to its ARC
            # - using slabinfo to estimate largest abusers (Except I don't think you can tell the differene between reclaim and unreclaim)
            SH  = mi['Shmem'] # TODO: figure out whether this does exactly what I think
            #CL  = mi['CommitLimit']
            #CAS = mi['Committed_AS']
            
            # ashare means it's added towards processes. 
            if not options.ashare: # If not, though, we add some indication of it here (overestimated!)
                pername['(shared, overestimated)'] = SH

            ST  = mi['SwapTotal']
            SF  = mi['SwapFree']
            SU  = ST - SF # swap use
            
            if options.showswap:
                # CONSIDER: make this a separate bar, maybe just first line.
                pername['(swap_used)'] = SU # TODO: force RED color
                pername['(swap_free)'] = SF



            # make list, filter, sum everything under 1% into 'small misc', sort
            name_mem = pername.items()

            if smallthresh < 1: # treat as faction of MemTotal
                size_thresh = smallthresh * MT
            else:
                size_thresh = smallthresh
            temp=[]
            restsize=0
            for name, mem in name_mem:                
                if mem < size_thresh and name not in ('(unused)','(caches, usable)'):
                    restsize+=mem
                else:
                    temp.append( (name,mem) )
            temp.append( ('(sumsmaller)', restsize) )
            name_mem = temp

            name_mem.sort(key=lambda x: x[1], reverse=True)

            rsses = list(mem   for name, mem in name_mem)
            normed = squarify.normalize_sizes(rsses, curses.COLS,curses.LINES)
            rects = squarify.squarify(normed, 0, 0, curses.COLS,curses.LINES)

            full = curses.newwin(curses.LINES,curses.COLS,0,0)
            full.clear()
            full.refresh()
            
            for i, rect in enumerate(rects):
                x = rect['x']
                y = rect['y']
                w = rect['dx']
                h = rect['dy']
                #x=int(x)
                #y=int(y)
                #w=int(w)
                #h=int(h)
                #TODO: check whether this math and ceil stuff is correct, I suspect it isn't.
                x=int(math.floor(round(x)))
                y=int(math.floor(round(y)))
                w=int(math.ceil(round(w)))
                h=int(math.ceil(round(h)))
                try:
                    win = curses.newwin(h,w,y,x)
                    name = name_mem[i][0]

                    #win.border()
                    win.addstr(1,1, name[:max(0,w-1)] )
                    win.addstr(2,1, ('%sB'%kmg(name_mem[i][1],kilo=1024))[:max(0,w-1)]  )
                    #win.border('|', '|', '-', '-', '+','+','+','+')
                    try:
                        import binascii
                        checksum = abs(binascii.crc32(name))
                    except:
                        checksum = sum(list(ord(c) for c in name))
                    colnum = 1+ (checksum%num_inited_cols)
                    if name=='(caches, usable)':
                        colnum = 0 # under current color logic this forces it to black
                    if name=='(unused)':
                        colnum = 0 
                    if name=='(available)':
                        colnum = 0 
                    if name=='(swap_free)':
                        colnum = 0 
                    
                    #print `curses.color_pair(colnum)` 
                    #r,g,b = curses.color_content(colnum)
                    #win.addstr(3,1, ('col %s|%s %s %s'%(colnum,r,g,b))[:max(0,w-1)]  )
                    win.bkgd( ' ', curses.color_pair(colnum) )
                except curses.error:
                    pass # assume it's because we're tring to draw out of screen due to rounding blah I should look at 
                win.refresh()

                #for win in wins:
                #del wins

            if options.ashare:
                time.sleep(5*interval) # back off the CPU use if we're using smem
            else:
                time.sleep(interval)
                
    finally:
        curses.endwin()
        curses.echo()

if __name__ == '__main__':
    main()