udp-receiver async - fix data corruption (with buffer pools)

.chloggen/pkg-stanza-input-udp-async-fix-data-corruption.yaml

@@ -0,0 +1,27 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: bug_fix
+
+# The name of the component, or a single word describing the area of concern, (e.g. filelogreceiver)
+component: pkg/stanza
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: Fix data-corruption/race-condition issue in udp async (reuse of buffer); use buffer pool isntead.
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [27613]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext: 
+
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: []

pkg/stanza/operator/input/udp/udp.go

@@ -138,6 +138,12 @@ func (c Config) Build(logger *zap.SugaredLogger) (operator.Operator, error) {
 
 	if c.AsyncConfig != nil {
 		udpInput.messageQueue = make(chan messageAndAddress, c.AsyncConfig.MaxQueueLength)
+		udpInput.readBufferPool = sync.Pool{
+			New: func() interface{} {
+				buffer := make([]byte, MaxUDPSize)
+				return &buffer
+			},
+		}
 	}
 	return udpInput, nil
 }
@@ -159,13 +165,15 @@ type Input struct {
 	splitFunc bufio.SplitFunc
 	resolver  *helper.IPResolver
 
-	messageQueue chan messageAndAddress
-	stopOnce     sync.Once
+	messageQueue   chan messageAndAddress
+	readBufferPool sync.Pool
+	stopOnce       sync.Once
 }
 
 type messageAndAddress struct {
-	Message    []byte
-	RemoteAddr net.Addr
+	Message       *[]byte
+	RemoteAddr    net.Addr
+	MessageLength int
 }
 
 // Start will start listening for messages on a socket.
@@ -206,9 +214,12 @@ func (u *Input) readAndProcessMessages(ctx context.Context) {
 	defer u.wg.Done()
 
 	dec := decode.New(u.encoding)
-	buf := make([]byte, 0, MaxUDPSize)
+	readBuffer := make([]byte, MaxUDPSize)
+	scannerBuffer := make([]byte, 0, MaxUDPSize)
 	for {
-		message, remoteAddr, err := u.readMessage()
+		message, remoteAddr, bufferLength, err := u.readMessage(readBuffer)
+		message = u.removeTrailingCharactersAndNULsFromBuffer(message, bufferLength)
+
 		if err != nil {
 			select {
 			case <-ctx.Done():
@@ -219,19 +230,19 @@ func (u *Input) readAndProcessMessages(ctx context.Context) {
 			break
 		}
 
-		u.processMessage(ctx, message, remoteAddr, dec, buf)
+		u.processMessage(ctx, message, remoteAddr, dec, scannerBuffer)
 	}
 }
 
-func (u *Input) processMessage(ctx context.Context, message []byte, remoteAddr net.Addr, dec *decode.Decoder, buf []byte) {
+func (u *Input) processMessage(ctx context.Context, message []byte, remoteAddr net.Addr, dec *decode.Decoder, scannerBuffer []byte) {
 	if u.OneLogPerPacket {
 		log := truncateMaxLog(message)
 		u.handleMessage(ctx, remoteAddr, dec, log)
 		return
 	}
 
 	scanner := bufio.NewScanner(bytes.NewReader(message))
-	scanner.Buffer(buf, MaxUDPSize)
+	scanner.Buffer(scannerBuffer, MaxUDPSize)
 
 	scanner.Split(u.splitFunc)
 
@@ -247,8 +258,10 @@ func (u *Input) readMessagesAsync(ctx context.Context) {
 	defer u.wg.Done()
 
 	for {
-		message, remoteAddr, err := u.readMessage()
+		readBuffer := u.readBufferPool.Get().(*[]byte) // Can't reuse the same buffer since same references would be written multiple times to the messageQueue (and cause data override of previous entries)
+		message, remoteAddr, bufferLength, err := u.readMessage(*readBuffer)
 		if err != nil {
+			u.readBufferPool.Put(readBuffer)
 			select {
 			case <-ctx.Done():
 				return
@@ -259,8 +272,9 @@ func (u *Input) readMessagesAsync(ctx context.Context) {
 		}
 
 		messageAndAddr := messageAndAddress{
-			Message:    message,
-			RemoteAddr: remoteAddr,
+			Message:       &message,
+			MessageLength: bufferLength,
+			RemoteAddr:    remoteAddr,
 		}
 
 		// Send the message to the message queue for processing
@@ -272,7 +286,7 @@ func (u *Input) processMessagesAsync(ctx context.Context) {
 	defer u.wg.Done()
 
 	dec := decode.New(u.encoding)
-	buf := make([]byte, 0, MaxUDPSize)
+	scannerBuffer := make([]byte, 0, MaxUDPSize)
 
 	for {
 		// Read a message from the message queue.
@@ -281,7 +295,9 @@ func (u *Input) processMessagesAsync(ctx context.Context) {
 			return // Channel closed, exit the goroutine.
 		}
 
-		u.processMessage(ctx, messageAndAddr.Message, messageAndAddr.RemoteAddr, dec, buf)
+		trimmedMessage := u.removeTrailingCharactersAndNULsFromBuffer(*messageAndAddr.Message, messageAndAddr.MessageLength)
+		u.processMessage(ctx, trimmedMessage, messageAndAddr.RemoteAddr, dec, scannerBuffer)
+		u.readBufferPool.Put(messageAndAddr.Message)
 	}
 }
 
@@ -331,17 +347,22 @@ func (u *Input) handleMessage(ctx context.Context, remoteAddr net.Addr, dec *dec
 }
 
 // readMessage will read log messages from the connection.
-func (u *Input) readMessage() ([]byte, net.Addr, error) {
-	n, addr, err := u.connection.ReadFrom(u.buffer)
+func (u *Input) readMessage(buffer []byte) ([]byte, net.Addr, int, error) {
+	n, addr, err := u.connection.ReadFrom(buffer)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, 0, err
 	}
 
+	return buffer, addr, n, nil
+}
+
+// This will remove trailing characters and NULs from the buffer
+func (u *Input) removeTrailingCharactersAndNULsFromBuffer(buffer []byte, n int) []byte {
 	// Remove trailing characters and NULs
-	for ; (n > 0) && (u.buffer[n-1] < 32); n-- { // nolint
+	for ; (n > 0) && (buffer[n-1] < 32); n-- { // nolint
 	}
 
-	return u.buffer[:n], addr, nil
+	return buffer[:n]
 }
 
 // Stop will stop listening for udp messages.