misc: several improvements to the graph module

modules/graph/edges.go

@@ -0,0 +1,158 @@
+package graph
+
+import (
+	"encoding/json"
+	"github.com/evilsocket/islazy/fs"
+	"io/ioutil"
+	"os"
+	"path"
+	"sort"
+	"sync"
+	"time"
+)
+
+const edgesIndexName = "edges.json"
+
+type EdgesTo map[string][]Edge
+
+type EdgesCallback func(string, []Edge, string) error
+
+type Edges struct {
+	sync.RWMutex
+	timestamp time.Time
+	fileName  string
+	size      int
+	from      map[string]EdgesTo
+}
+
+type edgesJSON struct {
+	Timestamp time.Time          `json:"timestamp"`
+	Size      int                `json:"size"`
+	Edges     map[string]EdgesTo `json:"edges"`
+}
+
+func LoadEdges(basePath string) (*Edges, error) {
+	edges := Edges{
+		fileName: path.Join(basePath, edgesIndexName),
+		from:     make(map[string]EdgesTo),
+	}
+
+	if fs.Exists(edges.fileName) {
+		var js edgesJSON
+
+		if raw, err := ioutil.ReadFile(edges.fileName); err != nil {
+			return nil, err
+		} else if err = json.Unmarshal(raw, &js); err != nil {
+			return nil, err
+		}
+
+		edges.timestamp = js.Timestamp
+		edges.from = js.Edges
+		edges.size = js.Size
+	}
+
+	return &edges, nil
+}
+
+func (e *Edges) flush() error {
+	e.timestamp = time.Now()
+	js := edgesJSON{
+		Timestamp: e.timestamp,
+		Size:      e.size,
+		Edges:     e.from,
+	}
+
+	if raw, err := json.Marshal(js); err != nil {
+		return err
+	} else if err = ioutil.WriteFile(e.fileName, raw, os.ModePerm); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (e *Edges) Flush() error {
+	e.RLock()
+	defer e.RUnlock()
+	return e.flush()
+}
+
+func (e *Edges) ForEachEdge(cb EdgesCallback) error {
+	e.RLock()
+	defer e.RUnlock()
+
+	for from, edgesTo := range e.from {
+		for to, edges := range edgesTo {
+			if err := cb(from, edges, to); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func (e *Edges) ForEachEdgeFrom(nodeID string, cb EdgesCallback) error {
+	e.RLock()
+	defer e.RUnlock()
+
+	if edgesTo, found := e.from[nodeID]; found {
+		for to, edges := range edgesTo {
+			if err := cb(nodeID, edges, to); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func (e *Edges) IsConnected(nodeID string) bool {
+	e.RLock()
+	defer e.RUnlock()
+
+	if edgesTo, found := e.from[nodeID]; found {
+		return len(edgesTo) > 0
+	}
+
+	return false
+}
+
+func (e *Edges) FindEdges(fromID, toID string, doSort bool) []Edge {
+	e.RLock()
+	defer e.RUnlock()
+
+	if edgesTo, foundFrom := e.from[fromID]; foundFrom {
+		if edges, foundTo := edgesTo[toID]; foundTo {
+			if doSort {
+				// sort edges from oldest to newer
+				sort.Slice(edges, func(i, j int) bool {
+					return edges[i].CreatedAt.Before(edges[j].CreatedAt)
+				})
+			}
+			return edges
+		}
+	}
+
+	return nil
+}
+
+func (e *Edges) Connect(fromID, toID string, edge Edge) error {
+	e.Lock()
+	defer e.Unlock()
+
+	if edgesTo, foundFrom := e.from[fromID]; foundFrom {
+		edges := edgesTo[toID]
+		edges = append(edges, edge)
+		e.from[fromID][toID] = edges
+	} else {
+		// create the entire path
+		e.from[fromID] = EdgesTo{
+			toID: {edge},
+		}
+	}
+
+	e.size++
+
+	return e.flush()
+}