Get rid of the package-level Connections func

probe/endpoint/procspy/background_reader_linux.go

@@ -21,36 +21,43 @@ const (
 type backgroundReader struct {
 	walker       process.Walker
 	mtx          sync.Mutex
+	running      bool
+	pleaseStop   bool
 	walkingBuf   *bytes.Buffer
 	readyBuf     *bytes.Buffer
 	readySockets map[uint64]*Proc
 }
 
-// HACK: Pretty ugly singleton interface (particularly the part of passing
-// the walker to StartBackgroundReader() and ignoring it in in Connections() )
-// experimenting with this for now.
-var singleton *backgroundReader
-
-func getBackgroundReader() (*backgroundReader, error) {
-	var err error
-	if singleton == nil {
-		err = fmt.Errorf("background reader hasn't yet been started")
-	}
-	return singleton, err
-}
-
-// StartBackgroundReader starts a ratelimited background goroutine to
-// read the expensive files from proc.
-func StartBackgroundReader(walker process.Walker) {
-	if singleton != nil {
-		return
-	}
-	singleton = &backgroundReader{
+// starts a rate-limited background goroutine to read the expensive files from
+// proc.
+func newBackgroundReader(walker process.Walker) *backgroundReader {
+	br := &backgroundReader{
 		walker:     walker,
 		walkingBuf: bytes.NewBuffer(make([]byte, 0, 5000)),
 		readyBuf:   bytes.NewBuffer(make([]byte, 0, 5000)),
 	}
-	go singleton.loop()
+	return br
+}
+
+func (br *backgroundReader) start() error {
+	br.mtx.Lock()
+	defer br.mtx.Unlock()
+	if br.running {
+		return fmt.Errorf("background reader already running")
+	}
+	br.running = true
+	go br.loop()
+	return nil
+}
+
+func (br *backgroundReader) stop() error {
+	br.mtx.Lock()
+	defer br.mtx.Unlock()
+	if !br.running {
+		return fmt.Errorf("background reader already not running")
+	}
+	br.pleaseStop = true
+	return nil
 }
 
 func (br *backgroundReader) loop() {
@@ -87,8 +94,17 @@ func (br *backgroundReader) loop() {
 
 		ticker = time.Tick(rateLimit)
 
-		// Swap buffers
 		br.mtx.Lock()
+
+		// Should we stop?
+		if br.pleaseStop {
+			br.pleaseStop = false
+			br.running = false
+			br.mtx.Unlock()
+			return
+		}
+
+		// Swap buffers
 		br.readyBuf, br.walkingBuf = br.walkingBuf, br.readyBuf
 		br.readySockets = sockets
 		br.mtx.Unlock()

probe/endpoint/procspy/fixture.go

@@ -1,13 +1,5 @@
 package procspy
 
-import (
-	"github.com/weaveworks/scope/probe/process"
-)
-
-// SetFixtures declares constant Connection and ConnectionProcs which will
-// always be returned by the package-level Connections and Processes
-// functions. It's designed to be used in tests.
-
 type fixedConnIter []Connection
 
 func (f *fixedConnIter) Next() *Connection {
@@ -21,10 +13,13 @@ func (f *fixedConnIter) Next() *Connection {
 	return &car
 }
 
-// SetFixtures is used in test scenarios to have known output.
-func SetFixtures(c []Connection) {
-	cbConnections = func(bool, process.Walker) (ConnIter, error) {
-		f := fixedConnIter(c)
-		return &f, nil
-	}
+// fixedScanner implements ConnectionScanner and uses constant Connection and
+// ConnectionProcs.  It's designed to be used in tests.
+type fixedScanner []Connection
+
+func (s fixedScanner) Connections() (ConnIter, error) {
+	iter := fixedConnIter(s)
+	return &iter, nil
 }
+
+func (s fixedScanner) Stop() {}

probe/endpoint/procspy/spy.go

@@ -5,8 +5,6 @@ package procspy
 
 import (
 	"net"
-
-	"github.com/weaveworks/scope/probe/process"
 )
 
 const (
@@ -35,11 +33,13 @@ type ConnIter interface {
 	Next() *Connection
 }
 
-// Connections returns all established (TCP) connections.  If processes is
-// false we'll just list all TCP connections, and there is no need to be root.
-// If processes is true it'll additionally try to lookup the process owning the
-// connection, filling in the Proc field. You will need to run this as root to
-// find all processes.
-func Connections(processes bool, walker process.Walker) (ConnIter, error) {
-	return cbConnections(processes, walker)
+type ConnectionScanner interface {
+	// Connections returns all established (TCP) connections.  If processes is
+	// false we'll just list all TCP connections, and there is no need to be root.
+	// If processes is true it'll additionally try to lookup the process owning the
+	// connection, filling in the Proc field. You will need to run this as root to
+	// find all processes.
+	Connections(processes bool) (ConnIter, error)
+	// Stops the scanning
+	Stop()
 }

probe/endpoint/procspy/spy_darwin.go

@@ -13,10 +13,16 @@ const (
 	lsofBinary    = "lsof"
 )
 
+func NewConnectionScanner(_ process.Walker) ConnectionScanner {
+	return &darwinScanner{}
+}
+
+type darwinScanner struct{}
+
 // Connections returns all established (TCP) connections. No need to be root
 // to run this. If processes is true it also tries to fill in the process
 // fields of the connection. You need to be root to find all processes.
-var cbConnections = func(processes bool, walker process.Walker) (ConnIter, error) {
+func (s *darwinScanner) Connections(processes bool) (ConnIter, error) {
 	out, err := exec.Command(
 		netstatBinary,
 		"-n", // no number resolving
@@ -62,3 +68,6 @@ var cbConnections = func(processes bool, walker process.Walker) (ConnIter, error
 	f := fixedConnIter(connections)
 	return &f, nil
 }
+
+// Nothing to stop since there's nothing running in the background
+func (s *darwinScanner) Stop() {}

probe/endpoint/procspy/spy_linux.go

@@ -32,19 +32,24 @@ func (c *pnConnIter) Next() *Connection {
 	return n
 }
 
-// cbConnections sets Connections()
-var cbConnections = func(processes bool, _ process.Walker) (ConnIter, error) {
+func NewConnectionScanner(walker process.Walker) ConnectionScanner {
+	br := newBackgroundReader(walker)
+	br.start()
+	return &linuxScanner{br}
+}
+
+type linuxScanner struct {
+	br *backgroundReader
+}
+
+func (s *linuxScanner) Connections(processes bool) (ConnIter, error) {
 	// buffer for contents of /proc/<pid>/net/tcp
 	buf := bufPool.Get().(*bytes.Buffer)
 	buf.Reset()
 
 	var procs map[uint64]*Proc
 	if processes {
-		br, err := getBackgroundReader()
-		if err != nil {
-			return nil, err
-		}
-		procs = br.getWalkedProcPid(buf)
+		procs = s.br.getWalkedProcPid(buf)
 	}
 
 	if buf.Len() == 0 {
@@ -58,3 +63,7 @@ var cbConnections = func(processes bool, _ process.Walker) (ConnIter, error) {
 		procs: procs,
 	}, nil
 }
+
+func (s *linuxScanner) Stop() {
+	s.br.stop()
+}

probe/endpoint/reporter.go

@@ -26,7 +26,7 @@ type Reporter struct {
 	includeProcesses bool
 	includeNAT       bool
 	flowWalker       flowWalker // interface
-	procWalker       process.Walker
+	scanner          procspy.ConnectionScanner
 	natMapper        natMapper
 	reverseResolver  *reverseResolver
 }
@@ -49,7 +49,6 @@ var SpyDuration = prometheus.NewSummaryVec(
 // is stored in the Endpoint topology. It optionally enriches that topology
 // with process (PID) information.
 func NewReporter(hostID, hostName string, includeProcesses bool, useConntrack bool, procWalker process.Walker) *Reporter {
-	procspy.StartBackgroundReader(procWalker)
 	return &Reporter{
 		hostID:           hostID,
 		hostName:         hostName,
@@ -57,7 +56,7 @@ func NewReporter(hostID, hostName string, includeProcesses bool, useConntrack bo
 		flowWalker:       newConntrackFlowWalker(useConntrack),
 		natMapper:        makeNATMapper(newConntrackFlowWalker(useConntrack, "--any-nat")),
 		reverseResolver:  newReverseResolver(),
-		procWalker:       procWalker,
+		scanner:          procspy.NewConnectionScanner(procWalker),
 	}
 }
 
@@ -69,6 +68,7 @@ func (r *Reporter) Stop() {
 	r.flowWalker.stop()
 	r.natMapper.stop()
 	r.reverseResolver.stop()
+	r.scanner.Stop()
 }
 
 // Report implements Reporter.
@@ -81,7 +81,7 @@ func (r *Reporter) Report() (report.Report, error) {
 	rpt := report.MakeReport()
 
 	{
-		conns, err := procspy.Connections(r.includeProcesses, r.procWalker)
+		conns, err := r.scanner.Connections(r.includeProcesses)
 		if err != nil {
 			return rpt, err
 		}