diff --git a/dag/graph.go b/dag/graph.go index 4e4f5be28..2834f1dc7 100644 --- a/dag/graph.go +++ b/dag/graph.go @@ -337,6 +337,16 @@ func (g *Graph) SetDebugWriter(w io.Writer) { g.debug.Encode(newMarshalGraph("root", g)) } +func (g *Graph) AnnotateVertex(v Vertex, info string) { + va := newVertexAnnotation(v, info) + g.debug.Encode(va) +} + +func (g *Graph) AnnotateEdge(e Edge, info string) { + ea := newEdgeAnnotation(e, info) + g.debug.Encode(ea) +} + // VertexName returns the name of a vertex. func VertexName(raw Vertex) string { switch v := raw.(type) { diff --git a/dag/marshal.go b/dag/marshal.go index 722fb4b93..47daadcbc 100644 --- a/dag/marshal.go +++ b/dag/marshal.go @@ -467,3 +467,33 @@ func decodeGraph(r io.Reader) (*marshalGraph, error) { } return g, nil } + +// *Annotation structs allow encoding arbitrary information about the graph in +// the logs. +type vertexAnnotation struct { + Type string + Vertex *marshalVertex + Info string +} + +func newVertexAnnotation(v Vertex, info string) *vertexAnnotation { + return &vertexAnnotation{ + Type: "VertexAnnotation", + Vertex: newMarshalVertex(v), + Info: info, + } +} + +type edgeAnnotation struct { + Type string + Edge *marshalEdge + Info string +} + +func newEdgeAnnotation(e Edge, info string) *edgeAnnotation { + return &edgeAnnotation{ + Type: "EdgeAnnotation", + Edge: newMarshalEdge(e), + Info: info, + } +} diff --git a/dag/marshal_test.go b/dag/marshal_test.go index 945d12f7c..970d359ec 100644 --- a/dag/marshal_test.go +++ b/dag/marshal_test.go @@ -129,6 +129,84 @@ func TestGraphJSON_basicRecord(t *testing.T) { } } +// Verify that Vertex and Edge annotations appear in the debug output +func TestGraphJSON_annotations(t *testing.T) { + var g Graph + var buf bytes.Buffer + g.SetDebugWriter(&buf) + + g.Add(1) + g.Add(2) + g.Add(3) + g.Connect(BasicEdge(1, 2)) + + g.AnnotateVertex(2, "2") + g.AnnotateVertex(3, "3") + g.AnnotateEdge(BasicEdge(1, 2), "1|2") + + dec := json.NewDecoder(bytes.NewReader(buf.Bytes())) + + var found2, found3, foundEdge bool + for dec.More() { + var d streamDecode + + err := dec.Decode(&d) + if err != nil { + t.Fatal(err) + } + + switch d.Type { + case "VertexAnnotation": + va := &vertexAnnotation{} + err := json.Unmarshal(d.JSON, va) + if err != nil { + t.Fatal(err) + } + + switch va.Info { + case "2": + if va.Vertex.Name != "2" { + t.Fatalf("wrong vertex annotated 2: %#v", va) + } + found2 = true + case "3": + if va.Vertex.Name != "3" { + t.Fatalf("wrong vertex annotated 3: %#v", va) + } + found3 = true + default: + t.Fatalf("unexpected annotation:", va) + } + case "EdgeAnnotation": + ea := &edgeAnnotation{} + err := json.Unmarshal(d.JSON, ea) + if err != nil { + t.Fatal(err) + } + + switch ea.Info { + case "1|2": + if ea.Edge.Name != "1|2" { + t.Fatalf("incorrect edge annotation: %#v\n", ea) + } + foundEdge = true + default: + t.Fatalf("unexpected edge Info: %#v", ea) + } + } + } + + if !found2 { + t.Fatal("annotation 2 not found") + } + if !found3 { + t.Fatal("annotation 3 not found") + } + if !foundEdge { + t.Fatal("edge annotation not found") + } +} + const testGraphJSONEmptyStr = `{ "Type": "Graph", "Name": "root",