Tool for experimenting with MADV_FREE

madv_free.go

// You may want to first disable transparent huge pages:
//
//    echo never | sudo tee /sys/kernel/mm/transparent_hugepage/enabled

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"unsafe"
)

/*
#include <stdlib.h>
#include <unistd.h>
#include <sys/mman.h>
#include <stdint.h>

static int inCore(void *base, uint64_t length, uint64_t pages) {
	int count = 0;
	unsigned char *vec = malloc(pages);
	if (vec == NULL)
		return -1;
	if (mincore(base, length, vec) < 0)
		return -1;
	for (int i = 0; i < pages; i++)
		if (vec[i] != 0)
			count++;
	free(vec);
	return count;
}
*/
import "C"

var pageSize = syscall.Getpagesize()

func main() {
	makeClean := flag.Bool("clean", false, "test clean file pages instead of dirty file pages")
	useDontneed := flag.Bool("dontneed", false, "use MADV_DONTNEED instead of MADV_FREE")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: %s [flags] anon-MiB file-MiB\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(2)
	}
	flag.Parse()
	if flag.NArg() != 2 {
		flag.Usage()
	}
	anonMB, err := strconv.Atoi(flag.Arg(0))
	if err != nil {
		flag.Usage()
	}
	fileMB, err := strconv.Atoi(flag.Arg(1))
	if err != nil {
		flag.Usage()
	}

	// Map anonymous memory.
	m, err := syscall.Mmap(-1, 0, anonMB<<20, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE|syscall.MAP_ANON)
	if err != nil {
		log.Fatal(err)
	}
	printStats("After anon mmap:", m, nil)

	// Fault in anonymous memory.
	for i := 0; i < len(m); i += pageSize {
		m[i] = 42
	}
	printStats("After anon fault:", m, nil)

	if *useDontneed {
		// MADV_DONTNEED the memory
		err = syscall.Madvise(m, syscall.MADV_DONTNEED)
		if err != nil {
			log.Fatal(err)
		}
		printStats("After MADV_DONTNEED:", m, nil)
	} else {
		// MADV_FREE the memory.
		err = syscall.Madvise(m, C.MADV_FREE)
		if err != nil {
			log.Fatal(err)
		}
		printStats("After MADV_FREE:", m, nil)
	}

	// Create a file to map.
	f, err := ioutil.TempFile("", "madv")
	if err != nil {
		log.Fatal(err)
	}
	os.Remove(f.Name())
	err = f.Truncate(int64(fileMB) << 20)
	if err != nil {
		log.Fatal(err)
	}

	// Map file memory.
	fm, err := syscall.Mmap(int(f.Fd()), 0, fileMB<<20, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE)
	if err != nil {
		log.Fatal(err)
	}
	f.Close()
	printStats("After file mmap:", m, fm)

	// Fault in file memory in 10 pieces, reporting stats at each
	// point.
	const nparts = 10
	fpages := len(fm) / pageSize
	for part := 0; part < nparts; part++ {
		start, end := fpages/nparts*part*pageSize, fpages/nparts*(1+part)*pageSize
		for i := start; i < end; i += pageSize {
			fm[i] = 42
		}
		if *makeClean {
			x, err := C.msync(unsafe.Pointer(&fm[start]), C.size_t(end-start), C.MS_SYNC)
			if x < 0 {
				log.Fatal("msync:", err)
			}
		}
		printStats(fmt.Sprintf("After fault %d MiB:", end/(1<<20)), m, fm)
	}

	runtime.KeepAlive(m)
	runtime.KeepAlive(fm)
}

func printStats(ident string, m, fm []byte) {
	fmt.Print(ident, " ", rss()/(1<<20), " MiB RSS, ", inCore(m)/(1<<20), " MiB anon in core")
	if fm != nil {
		fmt.Print(", ", inCore(fm)/(1<<20), " MiB file in core")
	}
	fmt.Println()
}

func rss() uintptr {
	data, err := ioutil.ReadFile("/proc/self/stat")
	if err != nil {
		log.Fatal(err)
	}
	fs := strings.Fields(string(data))
	rss, err := strconv.ParseInt(fs[23], 10, 64)
	if err != nil {
		log.Fatal(err)
	}
	return uintptr(rss) * uintptr(pageSize)
}

func inCore(b []byte) uintptr {
	n, err := C.inCore(unsafe.Pointer(&b[0]), C.uint64_t(len(b)), C.uint64_t(len(b)/pageSize))
	if n < 0 {
		log.Fatal(err)
	}
	return uintptr(n) * uintptr(pageSize)
}