1 files changed, 116 insertions, 0 deletions
diff --git a/tools/numap/numa.go b/tools/numap/numa.go
new file mode 100644
index 0000000..402ea1d
--- /dev/null
+++ b/tools/numap/numa.go
@@ -0,0 +1,116 @@
+package main
+
+import (
+	"fmt"
+	"io/ioutil"
+	"path"
+	"path/filepath"
+	"strings"
+
+	"golang.fcuny.net/numap/internal/hwids"
+	"golang.fcuny.net/numap/internal/sysfs"
+)
+
+const (
+	node_root      = "/sys/devices/system/node/node*"
+	CLASS_NVMe     = 67586
+	CLASS_ETHERNET = 131072
+	CLASS_GPU      = 197120
+)
+
+type node struct {
+	Name       string      `json:"name"`
+	Path       string      `json:"path"`
+	CpuList    string      `json:"cpulist"`
+	PCIDevices []PCIDevice `json:"pci_devices"`
+}
+
+type PCIDevice struct {
+	Vendor string `json:"vendor"`
+	Name   string `json:"name"`
+}
+
+func findNodes(hwdb hwids.PciDevices) (map[string]node, error) {
+	nodes := make(map[string]node)
+
+	files, err := filepath.Glob(node_root)
+	if err != nil {
+		return nil, fmt.Errorf("Failed to find NUMA nodes under %s: %+v", node_root, err)
+	}
+	if len(files) == 0 {
+		return nil, fmt.Errorf("Could not find NUMA node in %s", node_root)
+	}
+
+	for _, f := range files {
+		n, err := newNode(f)
+		if err != nil {
+			return make(map[string]node), err
+		}
+		nodes[n.Name] = n
+	}
+
+	r, err := mapPCIDevicesToNumaNode(hwdb)
+	if err != nil {
+		panic(err)
+	}
+	for k, v := range r {
+		nodeName := fmt.Sprintf("node%d", k)
+		n := nodes[nodeName]
+		n.PCIDevices = v
+		nodes[nodeName] = n
+	}
+	return nodes, nil
+}
+
+func mapPCIDevicesToNumaNode(hwdb hwids.PciDevices) (map[int][]PCIDevice, error) {
+	devices := sysfs.ScanPCIDevices()
+	r := map[int][]PCIDevice{}
+
+	for _, d := range devices {
+		if d.Class == CLASS_NVMe || d.Class == CLASS_ETHERNET || d.Class == CLASS_GPU {
+			_, ok := hwdb[uint16(d.Vendor)]
+			if ok {
+				desc := hwdb[uint16(d.Vendor)]
+				var vendor, name string
+				for _, m := range desc {
+					if uint64(m.Device) == d.Device && uint64(m.Vendor) == d.Vendor {
+						vendor = m.VendorName
+						name = m.DeviceName
+						break
+					}
+				}
+				pciDevice := PCIDevice{
+					Vendor: vendor,
+					Name:   name,
+				}
+				r[d.NumaNode] = append(r[d.NumaNode], pciDevice)
+			}
+		}
+	}
+	return r, nil
+}
+
+func newNode(p string) (node, error) {
+	_, name := path.Split(p)
+
+	cpulist, err := cpuList(p)
+	if err != nil {
+		return node{}, err
+	}
+
+	return node{
+		Name:       name,
+		Path:       p,
+		CpuList:    cpulist,
+		PCIDevices: []PCIDevice{},
+	}, nil
+}
+
+func cpuList(p string) (string, error) {
+	lpath := filepath.Join(p, "cpulist")
+	c, err := ioutil.ReadFile(lpath)
+	if err != nil {
+		return "", fmt.Errorf("Failed to open %s: %+v", lpath, err)
+	}
+	return strings.TrimRight(string(c), "\n"), nil
+}