Skip to content

Commit cbc71d9

Browse files
committed
dag: basic working update-able walker
1 parent 77c445a commit cbc71d9

2 files changed

Lines changed: 310 additions & 0 deletions

File tree

‎dag/walk.go‎

Lines changed: 267 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,267 @@
1+
package dag
2+
3+
import (
4+
"fmt"
5+
"log"
6+
"sync"
7+
"time"
8+
9+
"github.com/hashicorp/go-multierror"
10+
)
11+
12+
// walker performs a graph walk
13+
type walker struct {
14+
Callback WalkFunc
15+
16+
vertices *Set
17+
edges *Set
18+
vertexMap map[Vertex]*walkerVertex
19+
20+
wait sync.WaitGroup
21+
changeLock sync.Mutex
22+
23+
errMap map[Vertex]error
24+
errLock sync.Mutex
25+
}
26+
27+
type walkerVertex struct {
28+
sync.Mutex
29+
30+
DoneCh chan struct{}
31+
CancelCh chan struct{}
32+
DepsCh chan struct{}
33+
DepsUpdateCh chan chan struct{}
34+
35+
deps map[Vertex]chan struct{}
36+
depsCancelCh chan struct{}
37+
}
38+
39+
// Wait waits for the completion of the walk and returns any errors (
40+
// in the form of a multierror) that occurred. Update should be called
41+
// to populate the walk with vertices and edges.
42+
func (w *walker) Wait() error {
43+
// Wait for completion
44+
w.wait.Wait()
45+
46+
// Grab the error lock
47+
w.errLock.Lock()
48+
defer w.errLock.Unlock()
49+
50+
// Build the error
51+
var result error
52+
for v, err := range w.errMap {
53+
result = multierror.Append(result, fmt.Errorf(
54+
"%s: %s", VertexName(v), err))
55+
}
56+
57+
return result
58+
}
59+
60+
// Update updates the currently executing walk with the given vertices
61+
// and edges. It does not block until completion.
62+
//
63+
// Update can be called in parallel to Walk.
64+
func (w *walker) Update(v, e *Set) {
65+
// Grab the change lock so no more updates happen but also so that
66+
// no new vertices are executed during this time since we may be
67+
// removing them.
68+
w.changeLock.Lock()
69+
defer w.changeLock.Unlock()
70+
71+
// Initialize fields
72+
if w.vertexMap == nil {
73+
w.vertexMap = make(map[Vertex]*walkerVertex)
74+
}
75+
76+
// Calculate all our sets
77+
newEdges := e.Difference(w.edges)
78+
newVerts := v.Difference(w.vertices)
79+
oldVerts := w.vertices.Difference(v)
80+
81+
// Add the new vertices
82+
for _, raw := range newVerts.List() {
83+
v := raw.(Vertex)
84+
85+
// Add to the waitgroup so our walk is not done until everything finishes
86+
w.wait.Add(1)
87+
88+
// Initialize the vertex info
89+
info := &walkerVertex{
90+
DoneCh: make(chan struct{}),
91+
CancelCh: make(chan struct{}),
92+
DepsCh: make(chan struct{}),
93+
DepsUpdateCh: make(chan chan struct{}, 5),
94+
deps: make(map[Vertex]chan struct{}),
95+
}
96+
97+
// Close the deps channel immediately so it passes
98+
close(info.DepsCh)
99+
100+
// Add it to the map and kick off the walk
101+
w.vertexMap[v] = info
102+
}
103+
104+
// Remove the old vertices
105+
for _, raw := range oldVerts.List() {
106+
v := raw.(Vertex)
107+
108+
// Get the vertex info so we can cancel it
109+
info, ok := w.vertexMap[v]
110+
if !ok {
111+
// This vertex for some reason was never in our map. This
112+
// shouldn't be possible.
113+
continue
114+
}
115+
116+
// Cancel the vertex
117+
close(info.CancelCh)
118+
119+
// Delete it out of the map
120+
delete(w.vertexMap, v)
121+
}
122+
123+
// Add the new edges
124+
var changedDeps Set
125+
for _, raw := range newEdges.List() {
126+
edge := raw.(Edge)
127+
128+
// waiter is the vertex that is "waiting" on this edge
129+
waiter := edge.Target()
130+
131+
// dep is the dependency we're waiting on
132+
dep := edge.Source()
133+
134+
// Get the info for the waiter
135+
waiterInfo, ok := w.vertexMap[waiter]
136+
if !ok {
137+
// Vertex doesn't exist... shouldn't be possible but ignore.
138+
continue
139+
}
140+
141+
// Get the info for the dep
142+
depInfo, ok := w.vertexMap[dep]
143+
if !ok {
144+
// Vertex doesn't exist... shouldn't be possible but ignore.
145+
continue
146+
}
147+
148+
// Add the dependency to our waiter
149+
waiterInfo.deps[dep] = depInfo.DoneCh
150+
151+
// Record that the deps changed for this waiter
152+
changedDeps.Add(waiter)
153+
}
154+
155+
for _, raw := range changedDeps.List() {
156+
v := raw.(Vertex)
157+
info, ok := w.vertexMap[v]
158+
if !ok {
159+
// Vertex doesn't exist... shouldn't be possible but ignore.
160+
continue
161+
}
162+
163+
// Create a new done channel
164+
doneCh := make(chan struct{})
165+
166+
// Create the channel we close for cancellation
167+
cancelCh := make(chan struct{})
168+
info.depsCancelCh = cancelCh
169+
170+
// Build a new deps copy
171+
deps := make(map[Vertex]<-chan struct{})
172+
for k, v := range info.deps {
173+
deps[k] = v
174+
}
175+
176+
// Update the update channel
177+
info.DepsUpdateCh <- doneCh
178+
179+
// Start the waiter
180+
go w.waitDeps(v, deps, doneCh, cancelCh)
181+
}
182+
183+
// Kickstart all the vertices
184+
for _, raw := range newVerts.List() {
185+
v := raw.(Vertex)
186+
go w.walkVertex(v, w.vertexMap[v])
187+
}
188+
}
189+
190+
// walkVertex walks a single vertex, waiting for any dependencies before
191+
// executing the callback.
192+
func (w *walker) walkVertex(v Vertex, info *walkerVertex) {
193+
// When we're done executing, lower the waitgroup count
194+
defer w.wait.Done()
195+
196+
// When we're done, always close our done channel
197+
defer close(info.DoneCh)
198+
199+
// Wait for our dependencies
200+
depsCh := info.DepsCh
201+
for {
202+
select {
203+
case <-info.CancelCh:
204+
// Cancel
205+
return
206+
207+
case <-depsCh:
208+
// Deps complete!
209+
depsCh = nil
210+
211+
case depsCh = <-info.DepsUpdateCh:
212+
// New deps, reloop
213+
}
214+
215+
if depsCh == nil {
216+
// One final check if we have an update
217+
select {
218+
case depsCh = <-info.DepsUpdateCh:
219+
default:
220+
}
221+
222+
if depsCh == nil {
223+
break
224+
}
225+
}
226+
}
227+
228+
// Call our callback
229+
if err := w.Callback(v); err != nil {
230+
w.errLock.Lock()
231+
defer w.errLock.Unlock()
232+
233+
if w.errMap == nil {
234+
w.errMap = make(map[Vertex]error)
235+
}
236+
w.errMap[v] = err
237+
}
238+
}
239+
240+
func (w *walker) waitDeps(
241+
v Vertex,
242+
deps map[Vertex]<-chan struct{},
243+
doneCh chan<- struct{},
244+
cancelCh <-chan struct{}) {
245+
// Whenever we return, mark ourselves as complete
246+
defer close(doneCh)
247+
248+
// For each dependency given to us, wait for it to complete
249+
for dep, depCh := range deps {
250+
DepSatisfied:
251+
for {
252+
select {
253+
case <-depCh:
254+
// Dependency satisfied!
255+
break DepSatisfied
256+
257+
case <-cancelCh:
258+
// Wait cancelled
259+
return
260+
261+
case <-time.After(time.Second * 5):
262+
log.Printf("[DEBUG] vertex %q, waiting for: %q",
263+
VertexName(v), VertexName(dep))
264+
}
265+
}
266+
}
267+
}

‎dag/walk_test.go‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package dag
2+
3+
import (
4+
"reflect"
5+
"sync"
6+
"testing"
7+
)
8+
9+
func TestWalker_basic(t *testing.T) {
10+
var g Graph
11+
g.Add(1)
12+
g.Add(2)
13+
g.Connect(BasicEdge(1, 2))
14+
15+
// Run it a bunch of times since it is timing dependent
16+
for i := 0; i < 50; i++ {
17+
var order []interface{}
18+
w := &walker{Callback: walkCbRecord(&order)}
19+
w.Update(g.vertices, g.edges)
20+
21+
// Wait
22+
if err := w.Wait(); err != nil {
23+
t.Fatalf("err: %s", err)
24+
}
25+
26+
// Check
27+
expected := []interface{}{1, 2}
28+
if !reflect.DeepEqual(order, expected) {
29+
t.Fatalf("bad: %#v", order)
30+
}
31+
}
32+
}
33+
34+
// walkCbRecord is a test helper callback that just records the order called.
35+
func walkCbRecord(order *[]interface{}) WalkFunc {
36+
var l sync.Mutex
37+
return func(v Vertex) error {
38+
l.Lock()
39+
defer l.Unlock()
40+
*order = append(*order, v)
41+
return nil
42+
}
43+
}

0 commit comments

Comments
 (0)