Improve process code in probe

Makefile

@@ -23,7 +23,7 @@ $(SCOPE_EXPORT): $(APP_EXE) $(PROBE_EXE) docker/*
 
 $(APP_EXE): app/*.go render/*.go report/*.go xfer/*.go
 
-$(PROBE_EXE): probe/*.go probe/tag/*.go probe/docker/*.go report/*.go xfer/*.go
+$(PROBE_EXE): probe/*.go probe/tag/*.go probe/docker/*.go probe/process/*.go report/*.go xfer/*.go
 
 $(APP_EXE) $(PROBE_EXE):
 	go get -tags netgo ./$(@D)

probe/docker/reporter.go

@@ -27,11 +27,11 @@ func NewReporter(registry Registry, scope string) *Reporter {
 }
 
 // Report generates a Report containing Container and ContainerImage topologies
-func (r *Reporter) Report() report.Report {
+func (r *Reporter) Report() (report.Report, error) {
 	result := report.MakeReport()
 	result.Container.Merge(r.containerTopology())
 	result.ContainerImage.Merge(r.containerImageTopology())
-	return result
+	return result, nil
 }
 
 func (r *Reporter) containerTopology() report.Topology {

probe/docker/reporter_test.go

@@ -72,7 +72,7 @@ func TestReporter(t *testing.T) {
 	}
 
 	reporter := docker.NewReporter(mockRegistryInstance, "")
-	have := reporter.Report()
+	have, _ := reporter.Report()
 	if !reflect.DeepEqual(want, have) {
 		t.Errorf("%s", test.Diff(want, have))
 	}

probe/docker/tagger.go

@@ -3,7 +3,7 @@ package docker
 import (
 	"strconv"
 
-	"github.com/weaveworks/scope/probe/tag"
+	"github.com/weaveworks/scope/probe/process"
 	"github.com/weaveworks/scope/report"
 )
 
@@ -15,7 +15,7 @@ const (
 
 // These vars are exported for testing.
 var (
-	NewPIDTreeStub = tag.NewPIDTree
+	NewProcessTreeStub = process.NewTree
 )
 
 // Tagger is a tagger that tags Docker container information to process
@@ -35,15 +35,15 @@ func NewTagger(registry Registry, procRoot string) *Tagger {
 
 // Tag implements Tagger.
 func (t *Tagger) Tag(r report.Report) (report.Report, error) {
-	pidTree, err := NewPIDTreeStub(t.procRoot)
+	tree, err := NewProcessTreeStub(t.procRoot)
 	if err != nil {
 		return report.MakeReport(), err
 	}
-	t.tag(pidTree, &r.Process)
+	t.tag(tree, &r.Process)
 	return r, nil
 }
 
-func (t *Tagger) tag(pidTree tag.PIDTree, topology *report.Topology) {
+func (t *Tagger) tag(tree process.Tree, topology *report.Topology) {
 	for nodeID, nodeMetadata := range topology.NodeMetadatas {
 		pidStr, ok := nodeMetadata["pid"]
 		if !ok {
@@ -67,7 +67,7 @@ func (t *Tagger) tag(pidTree tag.PIDTree, topology *report.Topology) {
 					break
 				}
 
-				candidate, err = pidTree.GetParent(candidate)
+				candidate, err = tree.GetParent(candidate)
 				if err != nil {
 					break
 				}

probe/docker/tagger_test.go

@@ -6,33 +6,29 @@ import (
 	"testing"
 
 	"github.com/weaveworks/scope/probe/docker"
-	"github.com/weaveworks/scope/probe/tag"
+	"github.com/weaveworks/scope/probe/process"
 	"github.com/weaveworks/scope/report"
 	"github.com/weaveworks/scope/test"
 )
 
-type mockPIDTree struct {
+type mockProcessTree struct {
 	parents map[int]int
 }
 
-func (m *mockPIDTree) GetParent(pid int) (int, error) {
+func (m *mockProcessTree) GetParent(pid int) (int, error) {
 	parent, ok := m.parents[pid]
 	if !ok {
 		return -1, fmt.Errorf("Not found %d", pid)
 	}
 	return parent, nil
 }
 
-func (m *mockPIDTree) ProcessTopology(hostID string) report.Topology {
-	panic("")
-}
-
 func TestTagger(t *testing.T) {
-	oldPIDTree := docker.NewPIDTreeStub
-	defer func() { docker.NewPIDTreeStub = oldPIDTree }()
+	oldProcessTree := docker.NewProcessTreeStub
+	defer func() { docker.NewProcessTreeStub = oldProcessTree }()
 
-	docker.NewPIDTreeStub = func(procRoot string) (tag.PIDTree, error) {
-		return &mockPIDTree{map[int]int{2: 1}}, nil
+	docker.NewProcessTreeStub = func(procRoot string) (process.Tree, error) {
+		return &mockProcessTree{map[int]int{2: 1}}, nil
 	}
 
 	var (

probe/main.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/weaveworks/procspy"
 	"github.com/weaveworks/scope/probe/docker"
+	"github.com/weaveworks/scope/probe/process"
 	"github.com/weaveworks/scope/probe/tag"
 	"github.com/weaveworks/scope/report"
 	"github.com/weaveworks/scope/xfer"
@@ -105,6 +106,11 @@ func main() {
 		taggers = append(taggers, weaveTagger)
 	}
 
+	// TODO provide an alternate implementation for Darwin.
+	if runtime.GOOS == linux {
+		reporters = append(reporters, process.NewReporter(*procRoot, hostID))
+	}
+
 	log.Printf("listening on %s", *listen)
 
 	quit := make(chan struct{})
@@ -129,18 +135,12 @@ func main() {
 				// Do this every tick so it gets tagged by the OriginHostTagger
 				r.Host = hostTopology(hostID, hostName)
 
-				// TODO abstract PIDTree to a process provider, and provide an
-				// alternate implementation for Darwin.
-				if runtime.GOOS == linux {
-					if pidTree, err := tag.NewPIDTree(*procRoot); err == nil {
-						r.Process.Merge(pidTree.ProcessTopology(hostID))
-					} else {
-						log.Printf("PIDTree: %v", err)
-					}
-				}
-
 				for _, reporter := range reporters {
-					r.Merge(reporter.Report())
+					newReport, err := reporter.Report()
+					if err != nil {
+						log.Printf("error generating report: %v", err)
+					}
+					r.Merge(newReport)
 				}
 
 				if weaveTagger != nil {

probe/process/reporter.go

@@ -0,0 +1,61 @@
+package process
+
+import (
+	"strconv"
+
+	"github.com/weaveworks/scope/probe/tag"
+	"github.com/weaveworks/scope/report"
+)
+
+// We use these keys in node metadata
+const (
+	PID     = "pid"
+	Comm    = "comm"
+	PPID    = "ppid"
+	Cmdline = "cmdline"
+	Threads = "threads"
+)
+
+// Reporter generate Reports containing the Process topology
+type reporter struct {
+	procRoot string
+	scope    string
+}
+
+// NewReporter makes a new Reporter
+func NewReporter(procRoot, scope string) tag.Reporter {
+	return &reporter{
+		procRoot: procRoot,
+		scope:    scope,
+	}
+}
+
+// Report generates a Report containing the Process topology
+func (r *reporter) Report() (report.Report, error) {
+	result := report.MakeReport()
+	processes, err := r.processTopology()
+	if err != nil {
+		return result, err
+	}
+	result.Process.Merge(processes)
+	return result, nil
+}
+
+func (r *reporter) processTopology() (report.Topology, error) {
+	t := report.NewTopology()
+	err := Walk(r.procRoot, func(p *Process) {
+		pidstr := strconv.Itoa(p.PID)
+		nodeID := report.MakeProcessNodeID(r.scope, pidstr)
+		t.NodeMetadatas[nodeID] = report.NodeMetadata{
+			PID:     pidstr,
+			Comm:    p.Comm,
+			Cmdline: p.Cmdline,
+			Threads: strconv.Itoa(p.Threads),
+		}
+		if p.PPID > 0 {
+			t.NodeMetadatas[nodeID][PPID] = strconv.Itoa(p.PPID)
+		}
+	})
+
+	return t, err
+}

probe/process/reporter_test.go

@@ -0,0 +1,68 @@
+package process_test
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/weaveworks/scope/probe/process"
+	"github.com/weaveworks/scope/report"
+	"github.com/weaveworks/scope/test"
+)
+
+func TestReporter(t *testing.T) {
+	oldWalk := process.Walk
+	defer func() { process.Walk = oldWalk }()
+
+	process.Walk = func(_ string, f func(*process.Process)) error {
+		for _, p := range []*process.Process{
+			{PID: 1, PPID: 0, Comm: "init"},
+			{PID: 2, PPID: 1, Comm: "bash"},
+			{PID: 3, PPID: 1, Comm: "apache", Threads: 2},
+			{PID: 4, PPID: 2, Comm: "ping", Cmdline: "ping foo.bar.local"},
+		} {
+			f(p)
+		}
+		return nil
+	}
+
+	reporter := process.NewReporter("", "")
+	want := report.MakeReport()
+	want.Process = report.Topology{
+		Adjacency:     report.Adjacency{},
+		EdgeMetadatas: report.EdgeMetadatas{},
+		NodeMetadatas: report.NodeMetadatas{
+			report.MakeProcessNodeID("", "1"): report.NodeMetadata{
+				process.PID:     "1",
+				process.Comm:    "init",
+				process.Cmdline: "",
+				process.Threads: "0",
+			},
+			report.MakeProcessNodeID("", "2"): report.NodeMetadata{
+				process.PID:     "2",
+				process.Comm:    "bash",
+				process.PPID:    "1",
+				process.Cmdline: "",
+				process.Threads: "0",
+			},
+			report.MakeProcessNodeID("", "3"): report.NodeMetadata{
+				process.PID:     "3",
+				process.Comm:    "apache",
+				process.PPID:    "1",
+				process.Cmdline: "",
+				process.Threads: "2",
+			},
+			report.MakeProcessNodeID("", "4"): report.NodeMetadata{
+				process.PID:     "4",
+				process.Comm:    "ping",
+				process.PPID:    "2",
+				process.Cmdline: "ping foo.bar.local",
+				process.Threads: "0",
+			},
+		},
+	}
+
+	have, err := reporter.Report()
+	if err != nil || !reflect.DeepEqual(want, have) {
+		t.Errorf("%s (%v)", test.Diff(want, have), err)
+	}
+}

probe/process/tree.go

@@ -0,0 +1,34 @@
+package process
+
+import (
+	"fmt"
+)
+
+// Tree represents all processes on the machine.
+type Tree interface {
+	GetParent(pid int) (int, error)
+}
+
+type tree struct {
+	processes map[int]*Process
+}
+
+// NewTree returns a new Tree that can be polled.
+func NewTree(procRoot string) (Tree, error) {
+	pt := tree{processes: map[int]*Process{}}
+	err := Walk(procRoot, func(p *Process) {
+		pt.processes[p.PID] = p
+	})
+
+	return &pt, err
+}
+
+// GetParent returns the pid of the parent process for a given pid
+func (pt *tree) GetParent(pid int) (int, error) {
+	proc, ok := pt.processes[pid]
+	if !ok {
+		return -1, fmt.Errorf("PID %d not found", pid)
+	}
+
+	return proc.PPID, nil
+}

probe/process/tree_test.go

@@ -0,0 +1,37 @@
+package process_test
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/weaveworks/scope/probe/process"
+)
+
+func TestTree(t *testing.T) {
+	oldWalk := process.Walk
+	defer func() { process.Walk = oldWalk }()
+
+	process.Walk = func(_ string, f func(*process.Process)) error {
+		for _, p := range []*process.Process{
+			{PID: 1, PPID: 0},
+			{PID: 2, PPID: 1},
+			{PID: 3, PPID: 1},
+			{PID: 4, PPID: 2},
+		} {
+			f(p)
+		}
+		return nil
+	}
+
+	tree, err := process.NewTree("foo")
+	if err != nil {
+		t.Fatalf("newProcessTree error: %v", err)
+	}
+
+	for pid, want := range map[int]int{2: 1, 3: 1, 4: 2} {
+		have, err := tree.GetParent(pid)
+		if err != nil || !reflect.DeepEqual(want, have) {
+			t.Errorf("%d: want %#v, have %#v (%v)", pid, want, have, err)
+		}
+	}
+}