fx-mem-v2.go

package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
)

const PAGE_SIZE = 4096
const PRESENT_BIT = 1 << 63
const FRAME_MASK = 0x007fffffffffffff

func getAllPids() ([]int, error) {
	proc, err := os.Open("/proc")

	if err != nil {
		return nil, err
	}

	defer proc.Close()

	proc_entries, err := proc.Readdirnames(0)

	if err != nil {
		return nil, err
	}

	var pids []int = make([]int, 0)

	for _, name := range proc_entries {
		pid, err := strconv.Atoi(name)

		if err != nil {
			continue
		}

		pids = append(pids, pid)
	}

	return pids, nil
}

func getFirefoxPids(pids []int) ([]int, error) {
	firefoxPids := make([]int, 0)

	for _, pid := range pids {
		exePath := path.Join("/proc", strconv.Itoa(pid), "exe")
		linkTarget, err := os.Readlink(exePath)

		if err != nil {
			continue
		}

		_, exeFilename := path.Split(linkTarget)

		if exeFilename == "firefox" || exeFilename == "firefox (deleted)" {
			firefoxPids = append(firefoxPids, pid)
		}
	}

	return firefoxPids, nil
}

func getPhysicalPages(pid int) ([]uint64, error) {
	fieldsRegexp, err := regexp.Compile("\\s+")
	if err != nil {
		panic("Can't compile basic regexp!")
	}

	mapsFile, err := os.Open(path.Join("/proc", strconv.Itoa(pid), "maps"))

	if err != nil {
		return nil, err
	}

	defer mapsFile.Close()

	pagemapFile, err := os.Open(path.Join("/proc", strconv.Itoa(pid), "pagemap"))

	if err != nil {
		return nil, err
	}

	defer pagemapFile.Close()

	mapScanner := bufio.NewScanner(mapsFile)

	pages := make([]uint64, 0)

	for mapScanner.Scan() {
		line := mapScanner.Text()

		fields := fieldsRegexp.Split(line, -1)
		addressParts := strings.Split(fields[0], "-")
		addressStart, _ := strconv.ParseUint(addressParts[0], 16, 64)
		addressEnd, _ := strconv.ParseUint(addressParts[1], 16, 64)

		addrOffset := int64(addressStart>>12) * 8 // XXX hard-coded values
		_, err := pagemapFile.Seek(addrOffset, io.SeekStart)

		if err != nil {
			panic("couldn't seek in pagemap")
		}

		buf := make([]byte, 8*(addressEnd-addressStart)/PAGE_SIZE)
		_, err = pagemapFile.Read(buf)

		if err != nil {
			panic("couldn't read from pagemap")
		}

		reader := bytes.NewReader(buf)

		for addr := addressStart; addr < addressEnd; addr += PAGE_SIZE {
			var status uint64
			binary.Read(reader, binary.LittleEndian, &status)

			if (status & PRESENT_BIT) != 0 {
				pages = append(pages, status&FRAME_MASK)
			}
		}
	}

	return pages, nil
}

func contains(values []int, target int) bool {
	for _, value := range values {
		if value == target {
			return true
		}
	}

	return false
}

func main() {
	// XXX you need to run as root to read maps for all processes

	pids, err := getAllPids()

	if err != nil {
		fmt.Printf("unable to get pids: %v\n", err)
		return
	}

	firefoxPids, err := getFirefoxPids(pids)

	if err != nil {
		fmt.Printf("unable to get firefox pids: %v\n", err)
		return
	}

	if len(firefoxPids) == 0 {
		fmt.Printf("Firefox doesn't appear to be running...")
		return
	}

	nonFirefoxPages := make(map[uint64]bool, 0)

	for _, pid := range pids {
		if contains(firefoxPids, pid) {
			continue
		}

		pages, err := getPhysicalPages(pid)

		if err != nil {
			fmt.Printf("Can't read /proc/%v/maps: %v", pid, err)
			return
		}

		for _, page := range pages {
			nonFirefoxPages[page] = true
		}
	}

	firefoxOnlyPages := make(map[uint64]bool, 0)

	for _, pid := range firefoxPids {
		pages, err := getPhysicalPages(pid)

		if err != nil {
			fmt.Printf("Can't read /proc/%v/maps: %v", pid, err)
			return
		}

		for _, page := range pages {
			if !nonFirefoxPages[page] {
				firefoxOnlyPages[page] = true
			}
		}
	}

	fmt.Printf("%v\n", len(firefoxOnlyPages)*PAGE_SIZE)
}