star/hcache/main.go

//go:build linux

package hcache

/*
 * Copyright 2014-2015 Albert P. Tobey <atobey@datastax.com> @AlTobey
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * pcstat.go - page cache stat
 *
 * uses the mincore(2) syscall to find out which pages (almost always 4k)
 * of a file are currently cached in memory
 *
 */

import (
	"b612.me/apps/b612/hcache/pkg/utils"
	"bufio"
	"fmt"
	"github.com/spf13/cobra"
	"log"
	"os"
	"path"
	"sort"
	"strings"

	pcstat "github.com/tobert/pcstat/pkg"
)

var (
	pidFlag, topFlag                            int
	terseFlag, nohdrFlag, jsonFlag, unicodeFlag bool
	plainFlag, ppsFlag, histoFlag, bnameFlag    bool
)

func init() {
	// TODO: error on useless/broken combinations
	Cmd.Flags().IntVarP(&pidFlag, "pid", "p", 0, "show all open maps for the given pid")
	Cmd.Flags().IntVarP(&topFlag, "top", "t", 0, "show top x cached files in descending order")
	Cmd.Flags().BoolVarP(&terseFlag, "terse", "T", false, "show terse output")
	Cmd.Flags().BoolVarP(&nohdrFlag, "nohdr", "n", false, "omit the header from terse & text output")
	Cmd.Flags().BoolVarP(&jsonFlag, "json", "j", false, "return data in JSON format")
	Cmd.Flags().BoolVarP(&unicodeFlag, "unicode", "u", false, "return data with unicode box characters")
	Cmd.Flags().BoolVarP(&plainFlag, "plain", "P", false, "return data with no box characters")
	Cmd.Flags().BoolVarP(&ppsFlag, "pps", "S", false, "include the per-page status in JSON output")
	Cmd.Flags().BoolVarP(&histoFlag, "histo", "H", false, "print a simple histogram instead of raw data")
	Cmd.Flags().BoolVarP(&bnameFlag, "bname", "B", false, "convert paths to basename to narrow the output")
}

var Cmd = &cobra.Command{
	Use:   "hcache",
	Short: "Page cache status",
	Long:  `hcache - Page cache status`,
	Run: func(cmd *cobra.Command, args []string) {
		run(args)
	},
}

func uniqueSlice(slice *[]string) {
	found := make(map[string]bool)
	total := 0
	for i, val := range *slice {
		if _, ok := found[val]; !ok {
			found[val] = true
			(*slice)[total] = (*slice)[i]
			total++
		}
	}

	*slice = (*slice)[:total]
}

func getStatsFromFiles(files []string) PcStatusList {

	stats := make(PcStatusList, 0, len(files))
	for _, fname := range files {
		status, err := pcstat.GetPcStatus(fname)
		if err != nil {
			log.Printf("skipping %q: %v", fname, err)
			continue
		}

		// convert long paths to their basename with the -bname flag
		// this overwrites the original filename in pcs but it doesn't matter since
		// it's not used to access the file again -- and should not be!
		if bnameFlag {
			status.Name = path.Base(fname)
		}

		stats = append(stats, status)
	}
	return stats
}

func formatStats(stats PcStatusList) {
	if jsonFlag {
		stats.FormatJson(!ppsFlag)
	} else if terseFlag {
		stats.FormatTerse()
	} else if histoFlag {
		stats.FormatHistogram()
	} else if unicodeFlag {
		stats.FormatUnicode()
	} else if plainFlag {
		stats.FormatPlain()
	} else {
		stats.FormatText()
	}
}

func top(top int) {
	p, err := utils.Processes()
	if err != nil {
		log.Fatalf("err: %s", err)
	}

	if len(p) <= 0 {
		log.Fatal("Cannot find any process.")
	}

	results := make([]utils.Process, 0, 50)

	for _, p1 := range p {
		if p1.RSS() != 0 {
			results = append(results, p1)
		}
	}

	var files []string

	for _, process := range results {
		pcstat.SwitchMountNs(process.Pid())
		maps := getPidMaps(process.Pid())
		files = append(files, maps...)
	}

	uniqueSlice(&files)

	stats := getStatsFromFiles(files)

	sort.Sort(PcStatusList(stats))
	topStats := stats[:top]
	formatStats(topStats)
}

func run(files []string) {

	if topFlag != 0 {
		top(topFlag)
		os.Exit(0)
	}

	if pidFlag != 0 {
		pcstat.SwitchMountNs(pidFlag)
		maps := getPidMaps(pidFlag)
		files = append(files, maps...)
	}

	// all non-flag arguments are considered to be filenames
	// this works well with shell globbing
	// file order is preserved throughout this program
	if len(files) == 0 {
		os.Exit(1)
	}

	stats := getStatsFromFiles(files)
	sort.Sort(PcStatusList(stats))
	formatStats(stats)
}

func getPidMaps(pid int) []string {
	fname := fmt.Sprintf("/proc/%d/maps", pid)

	f, err := os.Open(fname)
	if err != nil {
		log.Fatalf("could not open '%s' for read: %v", fname, err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)

	// use a map to help avoid duplicates
	maps := make(map[string]bool)

	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) == 6 && strings.HasPrefix(parts[5], "/") {
			// found something that looks like a file
			maps[parts[5]] = true
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("reading '%s' failed: %s", fname, err)
	}

	// convert back to a list
	out := make([]string, 0, len(maps))
	for key := range maps {
		out = append(out, key)
	}

	return out
}