-
Notifications
You must be signed in to change notification settings - Fork 20
/
consumer.go
150 lines (128 loc) · 3.69 KB
/
consumer.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
package consumer
import (
"encoding/json"
"os"
"os/signal"
"strconv"
"strings"
"sync"
"github.com/Shopify/sarama"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
log "github.com/sirupsen/logrus"
)
var (
flowReceiveBytesTotal = promauto.NewCounterVec(
prometheus.CounterOpts{
Name: "flow_receive_bytes_total",
Help: "Bytes received.",
},
[]string{"source_as", "source_as_name", "destination_as", "destination_as_name", "hostname"},
)
flowTransmitBytesTotal = promauto.NewCounterVec(
prometheus.CounterOpts{
Name: "flow_transmit_bytes_total",
Help: "Bytes transferred.",
},
[]string{"source_as", "source_as_name", "destination_as", "destination_as_name", "hostname"},
)
)
type flow struct {
SourceAS int `json:"as_src"`
DestinationAS int `json:"as_dst"`
SourceIP string `json:"ip_dst"`
DestinationIP string `json:"ip_src"`
Bytes int `json:"bytes"`
Hostname string `json:"label"`
}
// Consume ...
func Consume(brokers string, topic string, partitions string, asn int, asns map[int]string) {
log.Info("Starting Kafka consumer")
c, err := sarama.NewConsumer(strings.Split(brokers, ","), nil)
if err != nil {
log.Fatalf("Failed to start consumer: %s", err)
}
partitionList, err := getPartitions(c, topic, partitions)
if err != nil {
log.Fatalf("Failed to start consumer: %s", err)
}
var (
messages = make(chan *sarama.ConsumerMessage, 256)
closing = make(chan struct{})
wg sync.WaitGroup
)
go func() {
signals := make(chan os.Signal, 1)
signal.Notify(signals, os.Kill, os.Interrupt)
<-signals
close(closing)
}()
for _, partition := range partitionList {
partitionConsumer, err := c.ConsumePartition(topic, partition, sarama.OffsetNewest)
if err != nil {
log.Fatalf("Failed to start consumer for partition %d: %s", partition, err)
}
go func(pc sarama.PartitionConsumer) {
<-closing
pc.AsyncClose()
}(partitionConsumer)
wg.Add(1)
go func(partitionConsumer sarama.PartitionConsumer) {
defer wg.Done()
for message := range partitionConsumer.Messages() {
messages <- message
}
}(partitionConsumer)
}
go func() {
for message := range messages {
logFlow(*message, asns, asn)
}
}()
wg.Wait()
close(messages)
if err := c.Close(); err != nil {
log.Warnf("Failed to close consumer: %s", err)
}
}
func getPartitions(c sarama.Consumer, topic string, partitions string) ([]int32, error) {
if partitions == "all" {
return c.Partitions(topic)
}
tmp := strings.Split(partitions, ",")
var partitionList []int32
for i := range tmp {
val, err := strconv.ParseInt(tmp[i], 10, 32)
if err != nil {
return nil, err
}
partitionList = append(partitionList, int32(val))
}
return partitionList, nil
}
func logFlow(message sarama.ConsumerMessage, asns map[int]string, asn int) {
var f flow
json.Unmarshal([]byte(message.Value), &f)
if f.SourceAS == asn {
flowTransmitBytesTotal.With(
prometheus.Labels{
"source_as": strconv.Itoa(f.SourceAS),
"source_as_name": asns[f.SourceAS],
"destination_as": strconv.Itoa(f.DestinationAS),
"destination_as_name": asns[f.DestinationAS],
"hostname": f.Hostname,
},
).Add(float64(f.Bytes))
} else if f.DestinationAS == asn {
flowReceiveBytesTotal.With(
prometheus.Labels{
"source_as": strconv.Itoa(f.SourceAS),
"source_as_name": asns[f.SourceAS],
"destination_as": strconv.Itoa(f.DestinationAS),
"destination_as_name": asns[f.DestinationAS],
"hostname": f.Hostname,
},
).Add(float64(f.Bytes))
}
log.WithFields(log.Fields{"offset": message.Offset}).Debug("Consumed message offset")
}