forked from paulstuart/influxsnmp
-
Notifications
You must be signed in to change notification settings - Fork 0
/
influx.go
143 lines (127 loc) · 3.01 KB
/
influx.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
package main
import (
"fmt"
"time"
client "github.com/influxdata/influxdb/client/v2"
"github.com/pkg/errors"
)
var (
retry = time.Second * 30
)
// Sender is a function that accepts the components of a datapoint
type Sender func(string, map[string]string, map[string]interface{}, time.Time) error
const (
// DefaultBatchSize is the default number points to batch before sending
DefaultBatchSize = 8192
// DefaultQueueSize is the default size the write queue
DefaultQueueSize = 65535
// DefaultFlush is the default of how often to send accumulated datapoints (in seconds)
DefaultFlush = 10
)
// dbCheck ensures the given database exists
func dbCheck(conn client.Client, database string) error {
if len(database) == 0 {
return fmt.Errorf("no database specified")
}
q := client.Query{Command: "show databases"}
resp, err := conn.Query(q)
if err != nil {
return err
}
for _, r := range resp.Results {
for _, s := range r.Series {
for _, v := range s.Values {
for _, d := range v {
if d.(string) == database {
return nil
}
}
}
}
}
return fmt.Errorf("database %s does not exist", database)
}
// NewSender returns a function that will accept datapoints to send to influxdb
func NewSender(
config interface{},
batch client.BatchPointsConfig,
batchSize int,
queueSize int,
flush int,
errFunc func(error),
) (Sender, error) {
if batchSize <= 0 {
batchSize = DefaultBatchSize
}
if queueSize <= 0 {
queueSize = DefaultQueueSize
}
if flush <= 0 {
flush = DefaultFlush
}
var conn client.Client
var err error
switch conf := config.(type) {
case client.HTTPConfig:
conn, err = client.NewHTTPClient(conf)
if err != nil {
return nil, errors.Wrap(err, "error creating HTTPClient")
}
_, _, err = conn.Ping(conf.Timeout)
if err != nil {
return nil, fmt.Errorf("cannot ping influxdb server: %s", conf.Addr)
}
if err := dbCheck(conn, batch.Database); err != nil {
return nil, errors.Wrapf(err, "check for database %s failed", batch.Database)
}
case client.UDPConfig:
conn, err = client.NewUDPClient(conf)
if err != nil {
return nil, errors.Wrap(err, "error creating UDPClient")
}
}
pts := make(chan *client.Point, queueSize)
bp, err := client.NewBatchPoints(batch)
if err != nil {
return nil, errors.Wrap(err, "batchpoints error")
}
go func() {
delay := time.Duration(flush) * time.Second
tick := time.Tick(delay)
count := 0
for {
select {
case p := <-pts:
bp.AddPoint(p)
count++
if count < batchSize {
continue
}
case <-tick:
if len(bp.Points()) == 0 {
continue
}
}
for {
if err := conn.Write(bp); err != nil {
if errFunc != nil {
errFunc(err)
}
time.Sleep(retry)
continue
}
bp, _ = client.NewBatchPoints(batch)
count = 0
break
}
}
}()
return func(key string, tags map[string]string, fields map[string]interface{}, ts time.Time) error {
pt, err := client.NewPoint(key, tags, fields, ts)
if err != nil {
return err
}
pts <- pt
return nil
}, nil
}