Stream the body in the trace-agent dogstatsd proxy

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.