Move SendFlows to separate struct for more obvious usage

Author: aaronfuj

config.go

@@ -198,7 +198,7 @@ func (c *Config) start(client *api.Client, dev *api.Device, errors chan<- error)
 		return nil, fmt.Errorf("agg setup error: %s", err)
 	}
 
-	sender := newSender(c.flow, c.timeout)
+	sender := newSender(c.flow)
 	sender.Errors = errors
 	sender.sample = c.sample
 	sender.Metrics = senderMetrics

direct_sender.go

@@ -0,0 +1,97 @@
+package libkflow
+
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"net/url"
+
+	capnp "zombiezen.com/go/capnproto2"
+
+	"github.com/kentik/libkflow/api"
+	"github.com/kentik/libkflow/flow"
+)
+
+// messagePrefix is an 80-byte prefix for the message header when sending kflow to the Kentik API. This is a deprecated
+// header, but the bytes must remain for backwards compatibility with the Kentik API.
+var messagePrefix = [80]byte{}
+
+// A DirectFlowSender transmits flows to Kentik.
+type DirectFlowSender struct {
+	client *api.Client
+
+	url      string
+	deviceID uint32
+}
+
+// NewDirectSender creates a new DirectFlowSender.
+func NewDirectSender(client *api.Client, device *api.Device, url *url.URL) *DirectFlowSender {
+	return &DirectFlowSender{
+		client: client,
+
+		url:      createURLString(*url, device.ClientID()),
+		deviceID: uint32(device.ID),
+	}
+}
+
+// SendFlows sends the flows to the Kentik API, returning the number of bytes sent as the payload. The device ID on
+// the flows is set to the device ID of the sender, regardless of what it was previously set to. This is to ensure all
+// data matches the expectations of the downstream URL/API.
+//
+// This will directly send the slice of flows without any additional downsampling or rate limiting. This does not
+// contribute to the underlying Send call.
+func (s *DirectFlowSender) SendFlows(flows []flow.Flow) (int64, error) {
+	if len(flows) == 0 {
+		return 0, nil
+	}
+	// ensure all flows have the device ID set; otherwise it may not be properly queried
+	for i := range flows {
+		flows[i].DeviceId = s.deviceID
+	}
+
+	// ensure the sample rate is matching the kentik api expectations
+	flow.NormalizeSampleRate(flows, 0)
+
+	// serialize the data
+	_, segment, err := capnp.NewMessage(capnp.SingleSegment(nil))
+	if err != nil {
+		return 0, fmt.Errorf("failed to create capn proto segment: %w", err)
+	}
+	message, err := flow.ToCapnProtoMessage(flows, segment)
+	if err != nil {
+		return 0, fmt.Errorf("failed to convert flows to capn proto: %w", err)
+	}
+
+	// write the data with additional gzip compression
+	buf := &bytes.Buffer{}
+	z := gzip.NewWriter(buf)
+	_, err = z.Write(messagePrefix[:])
+	if err != nil {
+		return 0, fmt.Errorf("failed to write empty message header: %w", err)
+	}
+	err = capnp.NewPackedEncoder(z).Encode(message)
+	if err != nil {
+		return 0, fmt.Errorf("failed to encode packed capn proto message: %w", err)
+	}
+	err = z.Close()
+	if err != nil {
+		return 0, fmt.Errorf("failed to close gzip writer: %w", err)
+	}
+
+	// send the compressed and packed message to the Kentik API
+	payloadLength := int64(len(buf.Bytes()))
+	err = s.client.SendFlow(s.url, buf)
+	if err != nil {
+		return 0, err
+	}
+	return payloadLength, nil
+}
+
+// createURLString creates the full URL string to use when sending data to the Kentik API.
+func createURLString(u url.URL, clientID string) string {
+	q := u.Query()
+	q.Set("sid", "0")
+	q.Set("sender_id", clientID)
+	u.RawQuery = q.Encode()
+	return u.String()
+}

direct_sender_test.go

@@ -0,0 +1,46 @@
+package libkflow
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/kentik/libkflow/api/test"
+	"github.com/kentik/libkflow/flow"
+)
+
+func TestDirectFlowSender_SendFlows(t *testing.T) {
+	client, server, device, err := test.NewClientServer()
+	require.NoError(t, err)
+	url := server.URL(test.FLOW)
+
+	directSender := NewDirectSender(client, device, url)
+
+	expected1 := flow.Flow{
+		DeviceId:  uint32(device.ID),
+		SrcAs:     rand.Uint32(),
+		DstAs:     rand.Uint32(),
+		SampleAdj: true,
+	}
+	expected2 := flow.Flow{
+		DeviceId:  uint32(device.ID),
+		SrcAs:     rand.Uint32(),
+		DstAs:     rand.Uint32(),
+		SampleAdj: true,
+	}
+
+	flows := []flow.Flow{expected1, expected2}
+	n, err := directSender.SendFlows(flows)
+	require.NoError(t, err)
+
+	msgs, err := receive(server)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, len(flows), msgs.Len())
+	assert.Greater(t, n, int64(0))
+	assert.Equal(t, flowToCHF(expected1, t).String(), msgs.At(0).String())
+	assert.Equal(t, flowToCHF(expected2, t).String(), msgs.At(1).String())
+}

send.go

@@ -20,16 +20,11 @@ import (
 	"github.com/kentik/libkflow/metrics"
 )
 
-// messagePrefix is an 80-byte prefix for the message header when sending kflow to the Kentik API. This is a deprecated
-// header, but the bytes must remain for backwards compatibility with the Kentik API.
-var messagePrefix = [80]byte{}
-
 // A Sender aggregates and transmits flow information to Kentik.
 type Sender struct {
 	agg               *agg.Agg
 	exit              chan struct{}
 	url               *url.URL
-	timeout           time.Duration
 	client            *api.Client
 	sample            int
 	ticker            *time.Ticker
@@ -43,12 +38,11 @@ type Sender struct {
 	Metrics           *metrics.Metrics
 }
 
-func newSender(url *url.URL, timeout time.Duration) *Sender {
+func newSender(url *url.URL) *Sender {
 	tickerCtx, cancelFunc := context.WithCancel(context.Background())
 	return &Sender{
 		exit:             make(chan struct{}),
 		url:              url,
-		timeout:          timeout,
 		ticker:           time.NewTicker(20 * time.Minute),
 		tickerCtx:        tickerCtx,
 		tickerCancelFunc: cancelFunc,
@@ -62,80 +56,6 @@ func (s *Sender) Send(flow *flow.Flow) {
 	s.agg.Add(flow)
 }
 
-// SendFlows sends the flows to the Kentik API, returning the number of bytes sent as the payload. The device ID on
-// the flows is set to the device ID of the sender, regardless of what it was previously set to. This is to ensure all
-// data matches the expectations of the downstream URL/API.
-//
-// This will directly send the slice of flows without any additional downsampling or rate limiting. This does not
-// contribute to the underlying Send call.
-func (s *Sender) SendFlows(flows []flow.Flow) (int64, error) {
-	s.workers.Add(1)
-	defer s.workers.Done()
-
-	if s.Device == nil {
-		return 0, fmt.Errorf("device not initialized")
-	}
-	if len(flows) == 0 {
-		return 0, nil
-	}
-	decoratedURL, err := s.createURLString()
-	if err != nil {
-		return 0, fmt.Errorf("failed to create URL string: %w", err)
-	}
-
-	if s.Metrics != nil {
-		s.Metrics.TotalFlowsIn.Mark(int64(len(flows)))
-	}
-
-	// ensure all flows have the device ID set; otherwise it may not be properly queried
-	for i := range flows {
-		flows[i].DeviceId = uint32(s.Device.ID)
-	}
-
-	// ensure the sample rate is matching the kentik api expectations
-	flow.NormalizeSampleRate(flows, 0)
-
-	// serialize the data
-	_, segment, err := capnp.NewMessage(capnp.SingleSegment(nil))
-	if err != nil {
-		return 0, fmt.Errorf("failed to create capn proto segment: %w", err)
-	}
-	message, err := flow.ToCapnProtoMessage(flows, segment)
-	if err != nil {
-		return 0, fmt.Errorf("failed to convert flows to capn proto: %w", err)
-	}
-
-	// write the data with additional gzip compression
-	buf := &bytes.Buffer{}
-	z := gzip.NewWriter(buf)
-	_, err = z.Write(messagePrefix[:])
-	if err != nil {
-		return 0, fmt.Errorf("failed to write empty message header: %w", err)
-	}
-	err = capnp.NewPackedEncoder(z).Encode(message)
-	if err != nil {
-		return 0, fmt.Errorf("failed to encode packed capn proto message: %w", err)
-	}
-	err = z.Close()
-	if err != nil {
-		return 0, fmt.Errorf("failed to close gzip writer: %w", err)
-	}
-
-	// send the compressed and packed message to the Kentik API
-	payloadLength := int64(len(buf.Bytes()))
-	err = s.client.SendFlow(decoratedURL, buf)
-	if err != nil {
-		return 0, err
-	}
-
-	if s.Metrics != nil {
-		s.Metrics.TotalFlowsOut.Mark(int64(len(flows)))
-		s.Metrics.BytesSent.Mark(payloadLength)
-	}
-
-	return payloadLength, nil
-}
-
 // Stop requests a graceful shutdown of the Sender.
 func (s *Sender) Stop(wait time.Duration) bool {
 	s.agg.Stop()
@@ -336,12 +256,5 @@ func (s *Sender) createURLString() (string, error) {
 	if s.url == nil {
 		return "", fmt.Errorf("url not initialized")
 	}
-
-	// Create a new URL to avoid modifying the original, backed by the config
-	u := *s.url
-	q := u.Query()
-	q.Set("sid", "0")
-	q.Set("sender_id", s.Device.ClientID())
-	u.RawQuery = q.Encode()
-	return u.String(), nil
+	return createURLString(*s.url, s.Device.ClientID()), nil
 }

send_test.go

@@ -42,38 +42,6 @@ func TestSender(t *testing.T) {
 	assert.Equal(flowToCHF(expected, t).String(), msgs.At(0).String())
 }
 
-func TestSender_SendFlows(t *testing.T) {
-	sender, server, assert := setup(t)
-
-	expected1 := flow.Flow{
-		DeviceId:  uint32(sender.Device.ID),
-		SrcAs:     rand.Uint32(),
-		DstAs:     rand.Uint32(),
-		SampleAdj: true,
-	}
-	expected2 := flow.Flow{
-		DeviceId:  uint32(sender.Device.ID),
-		SrcAs:     rand.Uint32(),
-		DstAs:     rand.Uint32(),
-		SampleAdj: true,
-	}
-
-	flows := []flow.Flow{expected1, expected2}
-	n, err := sender.SendFlows(flows)
-	assert.NoError(err)
-
-	msgs, err := receive(server)
-	if err != nil {
-		t.Fatal(err)
-	}
-	assert.Equal(int64(len(flows)), sender.Metrics.TotalFlowsIn.Count())
-	assert.Equal(int64(len(flows)), sender.Metrics.TotalFlowsOut.Count())
-	assert.Equal(n, sender.Metrics.BytesSent.Count())
-	assert.Equal(len(flows), msgs.Len())
-	assert.Equal(flowToCHF(expected1, t).String(), msgs.At(0).String())
-	assert.Equal(flowToCHF(expected2, t).String(), msgs.At(1).String())
-}
-
 func TestSenderStop(t *testing.T) {
 	sender, _, assert := setup(t)
 	stopped := sender.Stop(100 * time.Millisecond)
@@ -140,7 +108,7 @@ func setup(t testing.TB) (*Sender, *test.Server, *assert.Assertions) {
 	server.Log.SetOutput(io.Discard)
 
 	url := server.URL(test.FLOW)
-	sender := newSender(url, 1*time.Second)
+	sender := newSender(url)
 	sender.Metrics = metrics
 	sender.start(agg, client, device, 1)
 
@@ -226,7 +194,7 @@ func TestSender_createURLString(t *testing.T) {
 	agg, err := agg.NewAgg(10*time.Millisecond, 100, m)
 	require.NoError(t, err)
 
-	sender := newSender(u, 1*time.Second)
+	sender := newSender(u)
 	err = sender.start(agg, client, device, 1)
 	require.NoError(t, err)