// Copyright (C) 2017 Luke Shumaker <lukeshu@sbcglobal.net>

// Command cgswap displays which cgroups are using swap space.
package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
)

// You may be wondering why this program is so parallel.  Needless
// complexity, right?  Well, because processes may be very short
// lived, we want to read each process's info as soon as possible,
// before it goes away.

type cginfo struct {
	Cgroup string
	VmSwap int
}

type cginfos []cginfo

func (l cginfos) Len() int           { return len(l) }
func (l cginfos) Less(i, j int) bool { return l[i].VmSwap < l[j].VmSwap }
func (l cginfos) Swap(i, j int)      { l[i], l[j] = l[j], l[i] }

func inArray(needle string, haystack []string) bool {
	for _, straw := range haystack {
		if needle == straw {
			return true
		}
	}
	return false
}

func getVmSwap(statusFile io.Reader) (int, bool) {
	buf := bufio.NewScanner(statusFile)
	for buf.Scan() {
		line := buf.Text()
		if strings.HasPrefix(line, "VmSwap:") {
			swap, err := strconv.Atoi(strings.TrimSpace(strings.TrimSuffix(strings.TrimPrefix(line, "VmSwap:"), "kB")))
			if err != nil {
				return 0, false
			}
			return swap, true
		}
	}
	return 0, false
}

func getCgroup(cgroupFile io.Reader) (string, bool) {
	buf := bufio.NewScanner(cgroupFile)
	for buf.Scan() {
		parts := strings.SplitN(buf.Text(), ":", 3)
		if len(parts) != 3 {
			continue
		}
		heir := parts[0]
		controllers := strings.Split(parts[1], ",")
		cgroup := parts[2]
		if heir == "0" || inArray("name=systemd", controllers) {
			return cgroup, true
		}
	}
	return "", false
}

func main() {
	dir, err := os.Open("/proc")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	fileinfos, err := dir.Readdir(-1)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	ch := make(chan cginfo)
	var infos cginfos
	cgroup2idx := make(map[string]int)
	var consumers sync.WaitGroup
	consumers.Add(1)
	go func() {
		for info := range ch {
			idx, ok := cgroup2idx[info.Cgroup]
			if ok {
				infos[idx].VmSwap += info.VmSwap
			} else {
				cgroup2idx[info.Cgroup] = len(infos)
				infos = append(infos, info)
			}
		}
		consumers.Done()
	}()

	var producers sync.WaitGroup
	for _, fileinfo := range fileinfos {
		if pid, err := strconv.Atoi(fileinfo.Name()); fileinfo.IsDir() && err == nil {
			producers.Add(1)
			go func(pid int) {
				defer producers.Done()
				statusFile, err := os.Open(fmt.Sprintf("/proc/%d/status", pid))
				if err != nil {
					return
				}
				cgroupFile, err := os.Open(fmt.Sprintf("/proc/%d/cgroup", pid))
				if err != nil {
					return
				}
				swap, ok := getVmSwap(statusFile)
				if !ok || swap == 0 {
					return
				}
				cgroup, ok := getCgroup(cgroupFile)
				if !ok {
					return
				}
				ch <- cginfo{VmSwap: swap, Cgroup: cgroup}
			}(pid)
		}
	}

	producers.Wait()
	close(ch)
	consumers.Wait()

	sort.Sort(infos)

	total := 0
	for _, info := range infos {
		total += info.VmSwap
	}
	vmswapWidth := len(strconv.Itoa(total))

	for _, info := range infos {
		fmt.Printf("%[1]*d kB  %s\n", vmswapWidth, info.VmSwap, info.Cgroup)
	}
	fmt.Printf("%[1]*d kB  %s\n", vmswapWidth, total, "total")
}