-
Notifications
You must be signed in to change notification settings - Fork 492
/
moving_avg.go
176 lines (155 loc) · 3.96 KB
/
moving_avg.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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
package main
import (
"bytes"
"encoding/gob"
"errors"
"log"
"os"
"github.com/influxdata/kapacitor/udf"
"github.com/influxdata/kapacitor/udf/agent"
)
// An Agent.Handler that computes a moving average of the data it receives.
type avgHandler struct {
field string
as string
size int
state map[string]*avgState
agent *agent.Agent
}
// The state required to compute the moving average.
type avgState struct {
Size int
Window []float64
Avg float64
}
// Update the moving average with the next data point.
func (a *avgState) update(value float64) float64 {
l := len(a.Window)
if a.Size == l {
a.Avg += value/float64(l) - a.Window[0]/float64(l)
a.Window = a.Window[1:]
} else {
a.Avg = (value + float64(l)*a.Avg) / float64(l+1)
}
a.Window = append(a.Window, value)
return a.Avg
}
func newMovingAvgHandler(a *agent.Agent) *avgHandler {
return &avgHandler{
state: make(map[string]*avgState),
as: "avg",
agent: a,
}
}
// Return the InfoResponse. Describing the properties of this UDF agent.
func (a *avgHandler) Info() (*udf.InfoResponse, error) {
info := &udf.InfoResponse{
Wants: udf.EdgeType_STREAM,
Provides: udf.EdgeType_STREAM,
Options: map[string]*udf.OptionInfo{
"field": {ValueTypes: []udf.ValueType{udf.ValueType_STRING}},
"size": {ValueTypes: []udf.ValueType{udf.ValueType_INT}},
"as": {ValueTypes: []udf.ValueType{udf.ValueType_STRING}},
},
}
return info, nil
}
// Initialze the handler based of the provided options.
func (a *avgHandler) Init(r *udf.InitRequest) (*udf.InitResponse, error) {
init := &udf.InitResponse{
Success: true,
Error: "",
}
for _, opt := range r.Options {
switch opt.Name {
case "field":
a.field = opt.Values[0].Value.(*udf.OptionValue_StringValue).StringValue
case "size":
a.size = int(opt.Values[0].Value.(*udf.OptionValue_IntValue).IntValue)
case "as":
a.as = opt.Values[0].Value.(*udf.OptionValue_StringValue).StringValue
}
}
if a.field == "" {
init.Success = false
init.Error += " must supply field"
}
if a.size == 0 {
init.Success = false
init.Error += " must supply window size"
}
if a.as == "" {
init.Success = false
init.Error += " invalid as name provided"
}
return init, nil
}
// Create a snapshot of the running state of the process.
func (a *avgHandler) Snaphost() (*udf.SnapshotResponse, error) {
var buf bytes.Buffer
enc := gob.NewEncoder(&buf)
enc.Encode(a.state)
return &udf.SnapshotResponse{
Snapshot: buf.Bytes(),
}, nil
}
// Restore a previous snapshot.
func (a *avgHandler) Restore(req *udf.RestoreRequest) (*udf.RestoreResponse, error) {
buf := bytes.NewReader(req.Snapshot)
dec := gob.NewDecoder(buf)
err := dec.Decode(&a.state)
msg := ""
if err != nil {
msg = err.Error()
}
return &udf.RestoreResponse{
Success: err == nil,
Error: msg,
}, nil
}
// This handler does not do batching
func (a *avgHandler) BeginBatch(*udf.BeginBatch) error {
return errors.New("batching not supported")
}
// Receive a point and compute the average.
// Send a response with the average value.
func (a *avgHandler) Point(p *udf.Point) error {
// Update the moving average.
value := p.FieldsDouble[a.field]
state := a.state[p.Group]
if state == nil {
state = &avgState{Size: a.size}
a.state[p.Group] = state
}
avg := state.update(value)
// Re-use the existing point so we keep the same tags etc.
p.FieldsDouble = map[string]float64{a.as: avg}
p.FieldsInt = nil
p.FieldsString = nil
// Send point with average value.
a.agent.Responses <- &udf.Response{
Message: &udf.Response_Point{
Point: p,
},
}
return nil
}
// This handler does not do batching
func (a *avgHandler) EndBatch(*udf.EndBatch) error {
return errors.New("batching not supported")
}
// Stop the handler gracefully.
func (a *avgHandler) Stop() {
close(a.agent.Responses)
}
func main() {
a := agent.New(os.Stdin, os.Stdout)
h := newMovingAvgHandler(a)
a.Handler = h
log.Println("Starting agent")
a.Start()
err := a.Wait()
if err != nil {
log.Fatal(err)
}
}