Allow configuring UDP packets buffer in receiver. (#760)

* Allow configuring UDP packets buffer in receiver.

* Fix mixed descriptions

* Update CHANGELOG.md

---------

Co-authored-by: Mateusz Wedel <[email protected]>

Author: Wheezyx

CHANGELOG.md

@@ -1,3 +1,7 @@
+40.1.0
+------
+- Add support for configuration of receiver's buffer size - `receive-buffer-size`
+
 40.0.5
 ------
 - Clean up a bunch of old stuff in the repo

README.md

@@ -98,6 +98,8 @@ This configuration mode allows the following configuration options:
   Defaults to `false`.
 - `receive-batch-size`: the number of datagrams to attempt to read.  It is more CPU efficient to read multiple, however
   it takes extra memory.  See [Memory allocation for read buffers] section below for details.  Defaults to 50.
+- `receive-buffer-size`: the size of single buffer that's used to store datagram during read. Number of buffers is defined
+   by `receive-batch-size`. Default to 65535 what is theoretical maximum UDP packet size.
 - `conn-per-reader`: attempts to create a connection for every UDP receiver.  Not supported by all OS versions. It will
   be ignored when unix sockets are used for the connection.
   Defaults to `false`.
@@ -398,7 +400,7 @@ Memory allocation for read buffers
 By default `gostatsd` will batch read multiple packets to optimise read performance. The amount of memory allocated
 for these read buffers is determined by the config options:
 
-    max-readers * receive-batch-size * 64KB (max packet size)
+    max-readers * receive-batch-size * receive-buffer-size
 
 The metric `avg_packets_in_batch` can be used to track the average number of datagrams received per batch, and the
 `--receive-batch-size` flag used to tune it.  There may be some benefit to tuning the `--max-readers` flag as well.

cmd/gostatsd/main.go

@@ -166,6 +166,7 @@ func constructServer(v *viper.Viper) (*statsd.Server, error) {
 		PercentThreshold:      pt,
 		HeartbeatEnabled:      v.GetBool(gostatsd.ParamHeartbeatEnabled),
 		ReceiveBatchSize:      v.GetInt(gostatsd.ParamReceiveBatchSize),
+		ReceiveBufferSize:     v.GetInt(gostatsd.ParamReceiveBufferSize),
 		ConnPerReader:         v.GetBool(gostatsd.ParamConnPerReader),
 		ServerMode:            v.GetString(gostatsd.ParamServerMode),
 		LogRawMetric:          v.GetBool(gostatsd.ParamLogRawMetric),

cmd/lambda-extension/main.go

@@ -87,6 +87,7 @@ func NewServer(v *viper.Viper, logger logrus.FieldLogger) *statsd.Server {
 		StatserType:       v.GetString(gostatsd.ParamStatserType),
 		HeartbeatEnabled:  v.GetBool(gostatsd.ParamHeartbeatEnabled),
 		ReceiveBatchSize:  v.GetInt(gostatsd.ParamReceiveBatchSize),
+		ReceiveBufferSize: v.GetInt(gostatsd.ParamReceiveBufferSize),
 		ConnPerReader:     v.GetBool(gostatsd.ParamConnPerReader),
 		ServerMode:        GostatsdForwarderMode,
 		LogRawMetric:      v.GetBool(gostatsd.ParamLogRawMetric),

cmd/tester/main.go

@@ -49,6 +49,7 @@ func main() {
 			MaxQueueSize:          gostatsd.DefaultMaxQueueSize,
 			PercentThreshold:      gostatsd.DefaultPercentThreshold,
 			ReceiveBatchSize:      gostatsd.DefaultReceiveBatchSize,
+			ReceiveBufferSize:     gostatsd.DefaultReceiveBufferSize,
 			Viper:                 viper.New(),
 		}
 		ctx, cancelFunc := context.WithTimeout(context.Background(), time.Duration(s.Benchmark)*time.Second)

defaults_and_params.go

@@ -79,6 +79,8 @@ const (
 	DefaultHeartbeatEnabled = false
 	// DefaultReceiveBatchSize is the number of datagrams to read in each receive batch
 	DefaultReceiveBatchSize = 50
+	// DefaultReceiveBufferSize is the number of size of a buffer for each datagram during reads
+	DefaultReceiveBufferSize = 0xffff
 	// DefaultEstimatedTags is the estimated number of expected tags on an individual metric submitted externally
 	DefaultEstimatedTags = 4
 	// DefaultConnPerReader is the default for whether to create a connection per reader
@@ -162,6 +164,8 @@ const (
 	ParamHeartbeatEnabled = "heartbeat-enabled"
 	// ParamReceiveBatchSize is the name of the parameter with the number of datagrams to read in each receive batch
 	ParamReceiveBatchSize = "receive-batch-size"
+	// ParamReceiveBufferSize is the name of the parameter with the number that defines size of buffer of each datagram during reads
+	ParamReceiveBufferSize = "receive-buffer-size"
 	// ParamConnPerReader is the name of the parameter indicating whether to create a connection per reader
 	ParamConnPerReader = "conn-per-reader"
 	// ParamBadLineRateLimitPerMinute is the name of the parameter indicating how many bad lines can be logged per minute
@@ -216,6 +220,7 @@ func AddFlags(fs *pflag.FlagSet) {
 	fs.String(ParamPercentThreshold, strings.Join(toStringSlice(DefaultPercentThreshold), " "), "Space separated list of percentiles")
 	fs.Bool(ParamHeartbeatEnabled, DefaultHeartbeatEnabled, "Enables heartbeat")
 	fs.Int(ParamReceiveBatchSize, DefaultReceiveBatchSize, "The number of datagrams to read in each receive batch")
+	fs.Int(ParamReceiveBufferSize, DefaultReceiveBufferSize, "The number that defines size of buffer of each datagram during reads")
 	fs.Bool(ParamConnPerReader, DefaultConnPerReader, "Create a separate connection per reader (requires system support for reusing addresses)")
 	fs.String(ParamServerMode, DefaultServerMode, "The server mode to run in")
 	fs.String(ParamHostname, getHost(), "overrides the hostname of the server")

internal/awslambda/extension/manager_integration_test.go

@@ -128,6 +128,7 @@ func createStatsDServer(apiAddr string, fc flush.Coordinator, log logrus.FieldLo
 		MaxParsers:                1,
 		MaxWorkers:                1,
 		ReceiveBatchSize:          50,
+		ReceiveBufferSize:         0xffff,
 		MetricsAddr:               l.LocalAddr().String(),
 		ServerMode:                "forwarder",
 		Viper:                     v,

pkg/statsd/receiver.go

@@ -16,10 +16,6 @@ import (
 	"github.com/atlassian/gostatsd/pkg/stats"
 )
 
-// ip packet size is stored in two bytes and that is how big in theory the packet can be.
-// In practice it is highly unlikely but still possible to get packets bigger than usual MTU of 1500.
-const packetSizeUDP = 0xffff
-
 // DatagramReceiver receives datagrams on its PacketConn and passes them off to be parsed
 type DatagramReceiver struct {
 	// Counter fields below must be read/written only using atomic instructions.
@@ -31,20 +27,26 @@ type DatagramReceiver struct {
 	bufPool *pool.DatagramBufferPool
 
 	receiveBatchSize int // The number of datagrams to read in each batch
-	numReaders       int
-	socketFactory    SocketFactory
+
+	// ip packet size is stored in two bytes and that is how big in theory the packet can be.
+	// In practice it is highly unlikely but still possible to get packets bigger than usual MTU of 1500.
+	// Defaults to 64kB
+	receiveBufferSize int // size of each buffer used to read datagrams
+	numReaders        int
+	socketFactory     SocketFactory
 
 	out chan<- []*Datagram // Output chan of read datagram batches
 }
 
 // NewDatagramReceiver initialises a new DatagramReceiver.
-func NewDatagramReceiver(out chan<- []*Datagram, sf SocketFactory, numReaders, receiveBatchSize int) *DatagramReceiver {
+func NewDatagramReceiver(out chan<- []*Datagram, sf SocketFactory, numReaders, receiveBatchSize int, receiveBufferSize int) *DatagramReceiver {
 	return &DatagramReceiver{
-		out:              out,
-		receiveBatchSize: receiveBatchSize,
-		numReaders:       numReaders,
-		socketFactory:    sf,
-		bufPool:          pool.NewDatagramBufferPool(packetSizeUDP),
+		out:               out,
+		receiveBatchSize:  receiveBatchSize,
+		receiveBufferSize: receiveBufferSize,
+		numReaders:        numReaders,
+		socketFactory:     sf,
+		bufPool:           pool.NewDatagramBufferPool(receiveBufferSize),
 	}
 }
 

pkg/statsd/receiver_test.go

@@ -28,7 +28,7 @@ func BenchmarkReceive(b *testing.B) {
 	//
 	// ... so this is pretty arbitrary.
 	ch := make(chan []*Datagram, 5000)
-	mr := NewDatagramReceiver(ch, nil, 0, gostatsd.DefaultReceiveBatchSize)
+	mr := NewDatagramReceiver(ch, nil, 0, gostatsd.DefaultReceiveBatchSize, gostatsd.DefaultReceiveBufferSize)
 	c, done := fakesocket.NewCountedFakePacketConn(uint64(b.N))
 
 	var wg sync.WaitGroup
@@ -74,7 +74,7 @@ func BenchmarkReceive(b *testing.B) {
 
 func TestDatagramReceiver_Receive(t *testing.T) {
 	ch := make(chan []*Datagram, 1)
-	mr := NewDatagramReceiver(ch, nil, 0, 2)
+	mr := NewDatagramReceiver(ch, nil, 0, 2, gostatsd.DefaultReceiveBufferSize)
 	c := fakesocket.NewFakePacketConn()
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -103,7 +103,7 @@ func TestDatagramReceiver_UnixSocketConnection(t *testing.T) {
 
 	// Datagram receiver listening in Unix Domain Socket
 	socketPath := os.TempDir() + "/gostatsd_receiver_test_receive_uds.sock"
-	mr := NewDatagramReceiver(ch, socketFactory(socketPath, false), 1, 2)
+	mr := NewDatagramReceiver(ch, socketFactory(socketPath, false), 1, 2, gostatsd.DefaultReceiveBatchSize)
 	ctx, cancel := context.WithCancel(context.Background())
 
 	var wg sync.WaitGroup
@@ -137,7 +137,7 @@ func TestDatagramReceiver_UnixSocketIsRemovedOnContextCancellation(t *testing.T)
 	ch := make(chan []*Datagram, 1)
 
 	socketPath := os.TempDir() + "/gostatsd_receiver_test_receive_uds.sock"
-	mr := NewDatagramReceiver(ch, socketFactory(socketPath, false), 1, 2)
+	mr := NewDatagramReceiver(ch, socketFactory(socketPath, false), 1, 2, gostatsd.DefaultReceiveBatchSize)
 	ctx, cancel := context.WithCancel(context.Background())
 
 	var wg sync.WaitGroup

pkg/statsd/statsd.go

@@ -54,6 +54,7 @@ type Server struct {
 	HeartbeatEnabled          bool
 	HeartbeatTags             gostatsd.Tags
 	ReceiveBatchSize          int
+	ReceiveBufferSize         int
 	DisabledSubTypes          gostatsd.TimerSubtypes
 	HistogramLimit            uint32
 	BadLineRateLimitPerSecond rate.Limit
@@ -196,7 +197,7 @@ func (s *Server) RunWithCustomSocket(ctx context.Context, sf SocketFactory) erro
 	}
 
 	// Create the Receiver
-	receiver := NewDatagramReceiver(datagrams, sf, s.MaxReaders, s.ReceiveBatchSize)
+	receiver := NewDatagramReceiver(datagrams, sf, s.MaxReaders, s.ReceiveBatchSize, s.ReceiveBufferSize)
 	runnables = gostatsd.MaybeAppendRunnable(runnables, receiver)
 
 	// Create the Statser