Stream the body in the trace-agent dogstatsd proxy (#54226)

### What does this PR do?

Replaces the eager `bytes.Split` in the trace-agent dogstatsd proxy handler
(`/dogstatsd/v1/proxy`, `/dogstatsd/v2/proxy`) with a `bufio.Scanner` that walks
the request body one payload at a time, and:

- skips empty payloads, which carry nothing for dogstatsd to parse
- caps the payloads relayed per request (`maxDogstatsdProxyPayloads`, 100k),
  returning 413 beyond that

### Motivation

The handler read the whole body and then split it on newlines before relaying
each payload over UDP. `bytes.Split` pre-materializes one 24-byte slice header
per newline byte, so a body made mostly of newlines used several times its own
size in memory — at the default 25MB body limit that is hundreds of MB of
headers, and the memory is not returned promptly afterwards. Separately, the
number of UDP writes one request could drive was bounded only by the body size.

Scanning keeps the extra memory proportional to the longest line, and the cap
keeps the writes a single request can cause to a fixed cost.

### Describe how you validated your changes

`TestDogStatsDReverseProxyPayloadRelay` is a table test binding a real UDP
socket and asserting both the status code and what actually reaches dogstatsd:

| subtest | body | expects |
|---|---|---|
| `newlines only` | 100k newlines | 200, nothing relayed |
| `unreadable body` | one line then a failing reader | 500, first line relayed |
| `body over size limit` | limit cuts mid second line | 500, first line relayed |
| `more payloads than the proxy relays` | cap + 1 lines | 413, first payload relayed |

The pre-existing end-to-end UDP and UDS proxy tests still pass, and the whole
package passes with `-test.short=false` (545 tests). `dda inv linter.go
--targets=./pkg/trace/api` is clean.

### Additional Notes

The `bufio.Scanner` default 64KiB token limit is above the maximum UDP payload
(65507 bytes), so it cannot reject a payload that would otherwise have been
deliverable.

Co-authored-by: andrew.glaude <andrew.glaude@datadoghq.com>

Author: ajgajg1134

pkg/trace/api/dogstatsd.go

@@ -6,8 +6,7 @@
 package api
 
 import (
-	"bytes"
-	"io"
+	"bufio"
 	"net"
 	"net/http"
 	"strconv"
@@ -16,10 +15,25 @@ import (
 	"github.com/DataDog/datadog-agent/pkg/trace/log"
 )
 
+// maxDogstatsdProxyLines bounds the number of lines read from the body of a
+// single proxied request. Without it, both the number of UDP writes and the
+// number of scanner iterations one request causes are only bounded by the
+// request size, so a body of two-byte lines can drive one write per two bytes
+// sent, and a body of newlines one iteration per byte sent. Empty lines relay
+// nothing but still count, so that the work a request can ask for stays a fixed
+// cost. The limit stays well above any realistic batch, which is a handful of
+// metrics per request.
+const maxDogstatsdProxyLines = 100_000
+
 // dogstatsdProxyHandler returns a new HTTP handler which will proxy requests to
 // the DogStatsD endpoint in the Core Agent over UDP. Communication between the
 // proxy and the agent does not support UDS (see #13628), and so does not guarantee delivery of
 // all statsd payloads.
+//
+// The request body is relayed as it is read, so a body that turns out to be
+// unreadable, over the size limit, or over maxDogstatsdProxyLines is
+// reported as an error only after its earlier payloads have been sent. A client
+// that retries such a request may therefore submit those payloads twice.
 func (r *HTTPReceiver) dogstatsdProxyHandler() http.Handler {
 	if !r.conf.StatsdEnabled {
 		log.Info("DogstatsD disabled in the Agent configuration. The DogstatsD proxy endpoint will be non-functional.")
@@ -41,25 +55,40 @@ func (r *HTTPReceiver) dogstatsdProxyHandler() http.Handler {
 	}
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		req.Body = apiutil.NewLimitedReader(req.Body, r.conf.MaxRequestBytes)
-		body, err := io.ReadAll(req.Body)
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
-		payloads := bytes.Split(body, []byte("\n"))
-
 		conn, err := net.DialUDP("udp", nil, addr)
 		if err != nil {
 			log.Errorf("Error connecting to %s endpoint at %q: %v", "udp", addr, err)
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
 		defer conn.Close()
-		for _, p := range payloads {
-			if _, err := conn.Write(p); err != nil {
+
+		// Scan the body one line at a time rather than splitting it up front:
+		// splitting materializes a slice header per newline byte, so a body made
+		// of newlines multiplies its own size in memory several times over. The
+		// scanner keeps the extra memory proportional to the longest line.
+		scanner := bufio.NewScanner(req.Body)
+		lines := 0
+		for scanner.Scan() {
+			lines++
+			if lines > maxDogstatsdProxyLines {
+				log.Errorf("Dogstatsd proxy request contains more than %d lines, dropping the rest.", maxDogstatsdProxyLines)
+				http.Error(w, "too many dogstatsd payloads in request", http.StatusRequestEntityTooLarge)
+				return
+			}
+			payload := scanner.Bytes()
+			if len(payload) == 0 {
+				// Nothing for DogStatsD to parse; don't spend a syscall on it.
+				continue
+			}
+			if _, err := conn.Write(payload); err != nil {
 				http.Error(w, err.Error(), http.StatusInternalServerError)
 				return
 			}
 		}
+		if err := scanner.Err(); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
 	})
 }

pkg/trace/api/dogstatsd_test.go

@@ -7,13 +7,17 @@ package api
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
+	"strings"
+	"sync"
 	"testing"
+	"testing/iotest"
 	"time"
 
 	"github.com/stretchr/testify/require"
@@ -78,6 +82,116 @@ func TestDogStatsDReverseProxy(t *testing.T) {
 	})
 }
 
+// TestDogStatsDReverseProxyPayloadRelay covers how the proxy turns a request
+// body into DogStatsD payloads: newlines on their own carry no payload, and a
+// body that is unreadable, over the size limit, or over the payload cap is
+// reported as an error only once the payloads read before it have been relayed.
+func TestDogStatsDReverseProxyPayloadRelay(t *testing.T) {
+	const line = "users.online:1|c"
+	testCases := []struct {
+		name string
+		// maxRequestBytes overrides the default body size limit when non-zero.
+		maxRequestBytes int64
+		body            io.Reader
+		errCode         int
+		// wantFirstPayload is the first payload expected to reach DogStatsD,
+		// or empty when nothing at all should be relayed.
+		wantFirstPayload string
+	}{
+		{
+			name:    "newlines only",
+			body:    bytes.NewReader(bytes.Repeat([]byte("\n"), maxDogstatsdProxyLines)),
+			errCode: http.StatusOK,
+		},
+		{
+			// Empty lines relay nothing, but still count towards the limit.
+			name:    "more newlines than the proxy scans",
+			body:    bytes.NewReader(bytes.Repeat([]byte("\n"), maxDogstatsdProxyLines+1)),
+			errCode: http.StatusRequestEntityTooLarge,
+		},
+		{
+			name:             "unreadable body",
+			body:             io.MultiReader(strings.NewReader(line+"\n"), iotest.ErrReader(errors.New("read failed"))),
+			errCode:          http.StatusInternalServerError,
+			wantFirstPayload: line,
+		},
+		{
+			// The limit cuts the body mid way through its second line.
+			name:             "body over size limit",
+			maxRequestBytes:  int64(len(line)) + 5,
+			body:             strings.NewReader(line + "\n" + line),
+			errCode:          http.StatusInternalServerError,
+			wantFirstPayload: line,
+		},
+		{
+			name:             "more payloads than the proxy relays",
+			body:             strings.NewReader(strings.Repeat("x\n", maxDogstatsdProxyLines+1)),
+			errCode:          http.StatusRequestEntityTooLarge,
+			wantFirstPayload: "x",
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Bind a UDP socket so we can observe what the proxy relays.
+			conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
+			require.NoError(t, err)
+			defer conn.Close()
+
+			cfg := config.New()
+			cfg.StatsdHost = "127.0.0.1"
+			cfg.StatsdPort = conn.LocalAddr().(*net.UDPAddr).Port
+			if tc.maxRequestBytes != 0 {
+				cfg.MaxRequestBytes = tc.maxRequestBytes
+			}
+			receiver := newTestReceiverFromConfig(cfg)
+			proxy := receiver.dogstatsdProxyHandler()
+			require.NotNil(t, proxy)
+
+			// Drain the socket while the proxy runs, so that a case relaying
+			// many payloads cannot fill the receive buffer and have the rest
+			// dropped; the first payload is kept for the assertion below.
+			var mu sync.Mutex
+			var first []byte
+			var count int
+			go func() {
+				buf := make([]byte, 1024)
+				for {
+					n, _, err := conn.ReadFrom(buf)
+					if err != nil { // the socket was closed, the subtest is over
+						return
+					}
+					mu.Lock()
+					if count == 0 {
+						first = append([]byte(nil), buf[:n]...)
+					}
+					count++
+					mu.Unlock()
+				}
+			}()
+			// Counting rather than inspecting first, so that an empty payload
+			// is not mistaken for no payload at all.
+			relayed := func() bool {
+				mu.Lock()
+				defer mu.Unlock()
+				return count > 0
+			}
+
+			rec := httptest.NewRecorder()
+			proxy.ServeHTTP(rec, httptest.NewRequest("POST", "/", tc.body))
+			require.Equal(t, tc.errCode, rec.Code)
+
+			if tc.wantFirstPayload == "" {
+				require.Never(t, relayed, 100*time.Millisecond, 10*time.Millisecond, "expected no payload to be relayed")
+				return
+			}
+			require.Eventually(t, relayed, 5*time.Second, 10*time.Millisecond, "expected a payload to be relayed")
+			mu.Lock()
+			defer mu.Unlock()
+			require.Equal(t, tc.wantFirstPayload, string(first))
+		})
+	}
+}
+
 func testDogStatsDReverseProxyEndToEndUDP(t *testing.T, cfg *config.AgentConfig) {
 	hosts := []string{"localhost", "127.0.0.1", "::1"}
 	for _, host := range hosts {

releasenotes/notes/dogstatsd-proxy-line-scanner-e0a0a2c599995297.yaml

@@ -0,0 +1,9 @@
+---
+fixes:
+  - |
+    Fixed a bug in the trace-agent DogStatsD proxy endpoints
+    (``/dogstatsd/v1/proxy`` and ``/dogstatsd/v2/proxy``) where a request body
+    was split into all of its newline-separated payloads at once, so a body
+    containing many newlines used several times its own size in memory. The
+    body is now scanned one payload at a time, empty payloads are skipped, and
+    the number of payloads relayed per request is capped.