@@ -4,20 +4,26 @@ import (
44 "bytes"
55 "compress/gzip"
66 "context"
7+ "fmt"
78 "net/url"
89 "os"
910 "sync"
1011 "time"
1112
13+ "github.com/tinylib/msgp/msgp"
14+ capnp "zombiezen.com/go/capnproto2"
15+
1216 "github.com/kentik/libkflow/agg"
1317 "github.com/kentik/libkflow/api"
1418 "github.com/kentik/libkflow/flow"
1519 "github.com/kentik/libkflow/log"
1620 "github.com/kentik/libkflow/metrics"
17- "github.com/tinylib/msgp/msgp"
18- capnp "zombiezen.com/go/capnproto2"
1921)
2022
23+ // messagePrefix is an 80-byte prefix for the message header when sending kflow to the Kentik API. This is a deprecated
24+ // header, but the bytes must remain for backwards compatibility with the Kentik API.
25+ var messagePrefix = [80 ]byte {}
26+
2127// A Sender aggregates and transmits flow information to Kentik.
2228type Sender struct {
2329 agg * agg.Agg
@@ -56,6 +62,80 @@ func (s *Sender) Send(flow *flow.Flow) {
5662 s .agg .Add (flow )
5763}
5864
65+ // SendFlows sends the flows to the Kentik API, returning the number of bytes sent as the payload. The device ID on
66+ // the flows is set to the device ID of the sender, regardless of what it was previously set to. This is to ensure all
67+ // data matches the expectations of the downstream URL/API.
68+ //
69+ // This will directly send the slice of flows without any additional downsampling or rate limiting. This does not
70+ // contribute to the underlying Send call.
71+ func (s * Sender ) SendFlows (flows []flow.Flow ) (int64 , error ) {
72+ s .workers .Add (1 )
73+ defer s .workers .Done ()
74+
75+ if s .Device == nil {
76+ return 0 , fmt .Errorf ("device not initialized" )
77+ }
78+ if len (flows ) == 0 {
79+ return 0 , nil
80+ }
81+ decoratedURL , err := s .createURLString ()
82+ if err != nil {
83+ return 0 , fmt .Errorf ("failed to create URL string: %w" , err )
84+ }
85+
86+ if s .Metrics != nil {
87+ s .Metrics .TotalFlowsIn .Mark (int64 (len (flows )))
88+ }
89+
90+ // ensure all flows have the device ID set; otherwise it may not be properly queried
91+ for i := range flows {
92+ flows [i ].DeviceId = uint32 (s .Device .ID )
93+ }
94+
95+ // ensure the sample rate is matching the kentik api expectations
96+ flow .NormalizeSampleRate (flows , 0 )
97+
98+ // serialize the data
99+ _ , segment , err := capnp .NewMessage (capnp .SingleSegment (nil ))
100+ if err != nil {
101+ return 0 , fmt .Errorf ("failed to create capn proto segment: %w" , err )
102+ }
103+ message , err := flow .ToCapnProtoMessage (flows , segment )
104+ if err != nil {
105+ return 0 , fmt .Errorf ("failed to convert flows to capn proto: %w" , err )
106+ }
107+
108+ // write the data with additional gzip compression
109+ buf := & bytes.Buffer {}
110+ z := gzip .NewWriter (buf )
111+ _ , err = z .Write (messagePrefix [:])
112+ if err != nil {
113+ return 0 , fmt .Errorf ("failed to write empty message header: %w" , err )
114+ }
115+ err = capnp .NewPackedEncoder (z ).Encode (message )
116+ if err != nil {
117+ return 0 , fmt .Errorf ("failed to encode packed capn proto message: %w" , err )
118+ }
119+ err = z .Close ()
120+ if err != nil {
121+ return 0 , fmt .Errorf ("failed to close gzip writer: %w" , err )
122+ }
123+
124+ // send the compressed and packed message to the Kentik API
125+ payloadLength := int64 (len (buf .Bytes ()))
126+ err = s .client .SendFlow (decoratedURL , buf )
127+ if err != nil {
128+ return 0 , err
129+ }
130+
131+ if s .Metrics != nil {
132+ s .Metrics .TotalFlowsOut .Mark (int64 (len (flows )))
133+ s .Metrics .BytesSent .Mark (payloadLength )
134+ }
135+
136+ return payloadLength , nil
137+ }
138+
59139// Stop requests a graceful shutdown of the Sender.
60140func (s * Sender ) Stop (wait time.Duration ) bool {
61141 s .agg .Stop ()
@@ -107,18 +187,18 @@ func (s *Sender) SendEncodedDNS(data []byte) {
107187}
108188
109189func (s * Sender ) start (agg * agg.Agg , client * api.Client , device * api.Device , n int ) error {
110- q := s .url .Query ()
111- q .Set ("sid" , "0" )
112- q .Set ("sender_id" , device .ClientID ())
113-
114190 s .agg = agg
115- s .url .RawQuery = q .Encode ()
116191 s .Device = device
117192 s .client = client
118193 s .workers .Add (n )
119194
195+ decoratedURL , err := s .createURLString ()
196+ if err != nil {
197+ return fmt .Errorf ("failed to create URL string: %w" , err )
198+ }
199+
120200 for i := 0 ; i < n ; i ++ {
121- go s .dispatch ()
201+ go s .dispatch (decoratedURL )
122202 }
123203 go s .monitor ()
124204 go s .update ()
@@ -128,16 +208,16 @@ func (s *Sender) start(agg *agg.Agg, client *api.Client, device *api.Device, n i
128208 return nil
129209}
130210
131- func (s * Sender ) dispatch () {
211+ // dispatch runs a loop to send aggregated flow from the [agg.Agg] to the Kentik API.
212+ func (s * Sender ) dispatch (url string ) {
132213 buf := & bytes.Buffer {}
133214 cid := [80 ]byte {}
134- url := s .url .String ()
135215 z := gzip .NewWriter (buf )
136216
137217 for msg := range s .agg .Output () {
138218 log .Debugf ("dispatching aggregated flow" )
139219 z .Reset (buf )
140- z .Write (cid [:])
220+ _ , _ = z .Write (cid [:])
141221
142222 err := capnp .NewPackedEncoder (z ).Encode (msg )
143223 if err != nil {
@@ -247,3 +327,21 @@ func (s *Sender) error(err error) {
247327 default :
248328 }
249329}
330+
331+ // createURLString creates the full URL to use when sending data to the Kentik API.
332+ func (s * Sender ) createURLString () (string , error ) {
333+ if s .Device == nil {
334+ return "" , fmt .Errorf ("device not initialized" )
335+ }
336+ if s .url == nil {
337+ return "" , fmt .Errorf ("url not initialized" )
338+ }
339+
340+ // Create a new URL to avoid modifying the original, backed by the config
341+ u := * s .url
342+ q := u .Query ()
343+ q .Set ("sid" , "0" )
344+ q .Set ("sender_id" , s .Device .ClientID ())
345+ u .RawQuery = q .Encode ()
346+ return u .String (), nil
347+ }
0 commit comments