forked from dgraph-io/dgraph
-
Notifications
You must be signed in to change notification settings - Fork 0
/
watermark.go
129 lines (113 loc) · 3.81 KB
/
watermark.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
package x
import (
"container/heap"
"sync/atomic"
"time"
"golang.org/x/net/trace"
)
type uint64Heap []uint64
func (u uint64Heap) Len() int { return len(u) }
func (u uint64Heap) Less(i int, j int) bool { return u[i] < u[j] }
func (u uint64Heap) Swap(i int, j int) { u[i], u[j] = u[j], u[i] }
func (u *uint64Heap) Push(x interface{}) { *u = append(*u, x.(uint64)) }
func (u *uint64Heap) Pop() interface{} {
old := *u
n := len(old)
x := old[n-1]
*u = old[0 : n-1]
return x
}
// RaftValue contains the raft group and the raft proposal id.
// This is attached to the context, so the information could be passed
// down to the many posting lists, involved in mutations.
type RaftValue struct {
Group uint32
Index uint64
}
// Mark contains raft proposal id and a done boolean. It is used to
// update the WaterMark struct about the status of a proposal.
type Mark struct {
Index uint64
Done bool // Set to true if the pending mutation is done.
Deadline time.Time // Automatically set to true after deadline passes.
}
// WaterMark is used to keep track of the maximum done index. The right way to use
// this is to send a Mark with Done set to false, as soon as an index is known.
// WaterMark will store the index in a min-heap. It would only advance, if the minimum
// entry in the heap has been successfully done.
//
// Some time later, when this index task is completed, send another Mark, this time
// with Done set to true. It would mark the index as done, and so the min-heap can
// now advance and update the maximum done water mark.
type WaterMark struct {
Name string
Ch chan Mark
doneUntil uint64
elog trace.EventLog
}
// Init initializes a WaterMark struct. MUST be called before using it.
func (w *WaterMark) Init() {
w.Ch = make(chan Mark, 10000)
w.elog = trace.NewEventLog("Watermark", w.Name)
go w.process()
}
// DoneUntil returns the maximum index until which all tasks are done.
func (w *WaterMark) DoneUntil() uint64 {
return atomic.LoadUint64(&w.doneUntil)
}
// process is used to process the Mark channel. This is not thread-safe,
// so only run one goroutine for process. One is sufficient, because
// all goroutine ops use purely memory and cpu.
func (w *WaterMark) process() {
var indices uint64Heap
// pending maps raft proposal index to the number of pending mutations for this proposal.
pending := make(map[uint64]int)
heap.Init(&indices)
var loop uint64
for mark := range w.Ch {
// If not already done, then set. Otherwise, don't undo a done entry.
prev, present := pending[mark.Index]
if !present {
heap.Push(&indices, mark.Index)
}
delta := 1
if mark.Done {
delta = -1
}
pending[mark.Index] = prev + delta
if !mark.Deadline.IsZero() {
AssertTruef(delta == 1, "Invalid Mark: %+v", mark)
go func(m Mark) {
time.Sleep(m.Deadline.Sub(time.Now()))
w.Ch <- Mark{Index: m.Index, Done: true}
}(mark)
}
loop++
if len(indices) > 0 && loop%10000 == 0 {
min := indices[0]
w.elog.Printf("WaterMark %s: Done entry %4d. Size: %4d Watermark: %-4d Looking for: %-4d. Value: %d\n",
w.Name, mark.Index, len(indices), w.DoneUntil(), min, pending[min])
}
// Update mark by going through all indices in order; and checking if they have
// been done. Stop at the first index, which isn't done.
doneUntil := w.DoneUntil()
AssertTruef(doneUntil < mark.Index,
"Watermark %s: %d should be below current mark: %d", w.Name, doneUntil, mark.Index)
until := doneUntil
loops := 0
for len(indices) > 0 {
min := indices[0]
if done := pending[min]; done != 0 {
break
}
heap.Pop(&indices)
delete(pending, min)
until = min
loops++
}
if until != doneUntil {
AssertTrue(atomic.CompareAndSwapUint64(&w.doneUntil, doneUntil, until))
w.elog.Printf("%s: Done until %d. Loops: %d\n", w.Name, until, loops)
}
}
}