memory.go

package main

import (
        "bufio"
        "bytes"
        "fmt"
        "io/ioutil"
        "os"
        "os/exec"
        "path/filepath"
        "regexp"
        "strconv"
)

func calcPss(pid int) (int, error) {
        path := filepath.Join("/", "proc", strconv.Itoa(pid), "smaps")
        f, err := os.Open(path)
        if err != nil {
                return 0, err
        }
        defer f.Close()
        scanner := bufio.NewScanner(f)
        pssReg := regexp.MustCompile(`^Pss:\s+(\d+) (.{2})\s?$`)
        sizeInBytes := 0
        for scanner.Scan() {
                line := scanner.Text()
                m := pssReg.FindStringSubmatch(line)
                if len(m) == 0 {
                        continue
                }
                if m[2] != "kB" {
                        return 0, fmt.Errorf("%s is not recognized unit", m[2])
                }
                s, err := strconv.Atoi(m[1])
                if err != nil {
                        return 0, err
                }
                sizeInBytes += s * 1024
        }
        if err := scanner.Err(); err != nil {
                return 0, err
        }
        return sizeInBytes, nil
}

func getChildren(pid int) ([]int, error) {
        out, err := exec.Command("ps", "--ppid", strconv.Itoa(pid)).Output()
        if err != nil {
                if _, ok := err.(*exec.ExitError); !ok {
                        return nil, err
                }
        }
        scanner := bufio.NewScanner(bytes.NewBuffer(out))
        psReg := regexp.MustCompile(`^(\d+)\s`)
        children := []int{}
        for scanner.Scan() {
                line := scanner.Text()
                m := psReg.FindStringSubmatch(line)
                if len(m) == 0 {
                        continue
                }
                cpid, err := strconv.Atoi(m[1])
                if err != nil {
                        return nil, err
                }
                children = append(children, cpid)
        }
        if err := scanner.Err(); err != nil {
                return nil, err
        }
        return children, nil
}

type Process struct {
        pid      int
        cmdLine  string
        pss      int
        children []*Process
}

func NewProcess(pid int) (*Process, error) {
        pss, err := calcPss(pid)
        if err != nil {
                return nil, err
        }
        cmdLine, err := ioutil.ReadFile(filepath.Join("/", "proc", strconv.Itoa(pid), "cmdline"))
        if err != nil {
                return nil, err
        }
        process := &Process{
                pid:     pid,
                cmdLine: string(cmdLine),
                pss:     pss,
        }
        childrenPids, err := getChildren(pid)
        if err != nil {
                return nil, err
        }
        for _, pid := range childrenPids {
                p, err := NewProcess(pid)
                if err != nil {
                        return nil, err
                }
                process.children = append(process.children, p)
        }
        return process, nil
}

func (p *Process) TotalPss() int{
        pss := p.pss
        for _, child := range p.children {
                pss += child.TotalPss()
        }
        return pss
}

func main() {
        pid, err := strconv.Atoi(os.Args[1])
        if err != nil {
                panic(err)
        }
        p, err := NewProcess(pid)
        if err != nil {
                panic(err)
        }
        pss := p.TotalPss()
        fmt.Printf("%d\n", pss)
}

test.py

import time
import subprocess
from selenium import webdriver

class Test:
    def __init__(self, name, sites, driver):
        self.__name = name
        self.__sites = sites
        self.__driver_name = driver

    def Run(self):
        if self.__driver_name == 'chrome':
            self.__driver = webdriver.Chrome()                                                                                                                                                                                                                                                                                          
        elif self.__driver_name == 'firefox':
            self.__driver = webdriver.Firefox()                                                                                                                                                                                                                                                                          
        else:
            raise ValueError('invalid driver %s' % [self.__driver])
        pid = self.__driver.service.process.pid
        for site in self.__sites:
            self.__driver.get(site);
            time.sleep(5) # Let the user actually see something!                                                                                                                                                                                                                                                                                                             
        pss = int(subprocess.check_output(['go', 'run', 'memory.go', str(pid)]))
        print '%s,%s,%d' % (self.__name, self.__driver_name, pss)
        self.__driver.quit()

class PageSet:
    def __init__(self, name, sites):
        self.__name = name
        self.__sites = sites

    @property
    def name(self):
        return self.__name

    @property
    def sites(self):
        return self.__sites

if __name__ == '__main__':
    pagesets = [
        PageSet('Facebook', ['http://facebook.com/barackobama']),
    ]
    for pageset in pagesets:
        for driver in ['chrome', 'firefox']:
            test = Test(pageset.name, pageset.sites, driver)
            test.Run()