Reorganize code to be usable as module

grabber.go

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package main
+package realtimetraffic
 
 import (
 	"fmt"
@@ -53,7 +53,7 @@ var grabkeys = []string{
 	"tx_window_errors",
 }
 
-type grabber struct {
+type Grabber struct {
 	path    string
 	iface   string
 	quit    chan bool
@@ -61,16 +61,15 @@ type grabber struct {
 	count   uint
 }
 
-func newGrabber(iface string) *grabber {
-	return &grabber{
+func NewGrabber(iface string) *Grabber {
+	return &Grabber{
 		iface: iface,
 		path:  fmt.Sprintf("/sys/class/net/%s/statistics", iface),
 	}
 }
 
-func (g *grabber) grab() {
-
-	data := &interfacedata{}
+func (g *Grabber) grab() *Interfacedata {
+	data := &Interfacedata{}
 	statistics := map[string]interface{}{}
 	value := make([]byte, 100)
 
@@ -79,14 +78,14 @@ func (g *grabber) grab() {
 	for _, key := range grabkeys {
 		file, err = os.Open(path.Join(g.path, key))
 		if err == nil {
-			count, err := file.Read(value)
-			if err == nil {
-				statistics[key], err = strconv.ParseInt(strings.TrimSpace(string(value[:count])), 10, 64)
-				if err != nil {
-					log.Println("Failed to process data", err, value[:count], count)
+			count, ferr := file.Read(value)
+			if ferr == nil {
+				statistics[key], ferr = strconv.ParseInt(strings.TrimSpace(string(value[:count])), 10, 64)
+				if ferr != nil {
+					log.Println("Failed to process data", ferr, value[:count], count)
 				}
 			} else {
-				log.Println("Failed to read data", err)
+				log.Println("Failed to read data", ferr)
 			}
 			file.Close()
 		} else {
@@ -95,12 +94,11 @@ func (g *grabber) grab() {
 	}
 
 	data.set(g.iface, statistics)
-	h.broadcast <- data
 
+	return data
 }
 
-func (g *grabber) start() {
-
+func (g *Grabber) Start(ch chan<- *Interfacedata) {
 	g.count++
 	if g.running {
 		return
@@ -115,18 +113,16 @@ func (g *grabber) start() {
 		for {
 			select {
 			case <-ticker.C:
-				g.grab()
+				ch <- g.grab()
 			case <-g.quit:
 				ticker.Stop()
 				return
 			}
 		}
 	}()
-
 }
 
-func (g *grabber) stop() {
-
+func (g *Grabber) Stop() {
 	if !g.running {
 		return
 	}
@@ -137,5 +133,4 @@ func (g *grabber) stop() {
 		close(g.quit)
 		fmt.Printf("Grabbing stopped: %s\n", g.iface)
 	}
-
 }

interfacedata.go

@@ -15,25 +15,33 @@
  * limitations under the License.
  */
 
-package main
+package realtimetraffic
 
 import (
 	"encoding/json"
 )
 
-type interfacedata struct {
-	data  interface{}
+type Interfacedata struct {
+	data  map[string]interface{}
 	iface string
 }
 
-func (i *interfacedata) set(iface string, data interface{}) {
+func (i *Interfacedata) set(iface string, data interface{}) {
+	i.iface = iface
 	i.data = map[string]interface{}{
 		iface:  data,
 		"name": iface,
 	}
-	i.iface = iface
 }
 
-func (i *interfacedata) encode() ([]byte, error) {
+func (i *Interfacedata) Data() map[string]interface{} {
+	return i.data
+}
+
+func (i *Interfacedata) JSON() ([]byte, error) {
 	return json.Marshal(i.data)
 }
+
+func (i *Interfacedata) Name() string {
+	return i.iface
+}

realtimetrafficd/hub.go

@@ -17,25 +17,28 @@
 
 package main
 
+import (
+	"github.com/longsleep/realtimetraffic"
+)
+
 type hub struct {
-	grabbers    map[string]*grabber
+	grabbers    map[string]*realtimetraffic.Grabber
 	connections map[*connection]bool
-	broadcast   chan *interfacedata
+	broadcast   chan *realtimetraffic.Interfacedata
 	register    chan *connection
 	unregister  chan *connection
 }
 
 var h = hub{
-	broadcast:   make(chan *interfacedata),
+	broadcast:   make(chan *realtimetraffic.Interfacedata),
 	register:    make(chan *connection),
 	unregister:  make(chan *connection),
 	connections: make(map[*connection]bool),
-	grabbers:    make(map[string]*grabber),
+	grabbers:    make(map[string]*realtimetraffic.Grabber),
 }
 
 func (h *hub) run() {
-
-	var eg *grabber
+	var eg *realtimetraffic.Grabber
 	var ok bool
 
 	for {
@@ -43,20 +46,20 @@ func (h *hub) run() {
 		case c := <-h.register:
 			h.connections[c] = true
 			if eg, ok = h.grabbers[c.iface]; !ok {
-				eg = newGrabber(c.iface)
+				eg = realtimetraffic.NewGrabber(c.iface)
 				h.grabbers[c.iface] = eg
 			}
-			eg.start()
+			eg.Start(h.broadcast)
 		case c := <-h.unregister:
 			delete(h.connections, c)
 			close(c.send)
 			if eg, ok = h.grabbers[c.iface]; ok {
-				eg.stop()
+				eg.Stop()
 			}
 		case d := <-h.broadcast:
 			for c := range h.connections {
-				if c.iface == d.iface {
-					if m, err := d.encode(); err == nil {
+				if c.iface == d.Name() {
+					if m, err := d.JSON(); err == nil {
 						select {
 						case c.send <- m:
 						default: