feat: Apply the init-concurrency limiter to polling and streaming

docs/configuration.md

@@ -137,6 +137,19 @@ _(9)_ The `metricsCapacity` setting controls the queue for usage metrics events,
 _(9)_ The `maxClientRequestBodySize` setting limits how much of a `REPORT` evaluation request body the Relay Proxy will read into memory before decoding the context, protecting the process from memory exhaustion caused by oversized request bodies. It applies to the `evalx` context/user endpoints for client-side, mobile, and server-side SDKs. The default value is `5MiB`. Requests whose body exceeds the limit receive an HTTP `413 Request Entity Too Large` response. Setting a non-positive value (such as `0B`) is rejected at startup; to raise or lower the limit, specify a positive value using the same units as `maxInboundPayloadSize` (for example, `10MiB`).
 
 
+### File section: `[Concurrency]`
+
+This section bounds how many SDK *initialization deliveries* — the full data-set payload the Relay Proxy serializes and sends when an SDK first connects — it will perform at once. It covers the FDv2 poll endpoints, the FDv1 PHP all-flags poll, and the server-side streaming endpoints (`/all` and `/sdk/stream`); the `evalx` evaluation endpoints and the legacy FDv1 flags-only stream are not gated. It protects the Relay Proxy from the memory and egress spikes a large burst of connecting or reconnecting SDKs can cause. It is disabled by default, so behavior is unchanged unless you set `maxConcurrent`.
+
+| Property in file | Environment var       |   Type   | Default | Description                                                                                                                                                                                                       |
+|------------------|-----------------------|:--------:|:--------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `maxConcurrent`  | `INIT_MAX_CONCURRENT` |  Number  | none    | Maximum number of initialization deliveries in flight at once. `0` or unset disables the limit. Cheap operations (an up-to-date reply, deltas, heartbeats, and single-item lookups) are never counted.            |
+| `maxQueued`      | `INIT_MAX_QUEUED`     |  Number  | `0`     | Maximum number of requests that may wait for a slot once `maxConcurrent` is reached. `0` sheds excess requests immediately rather than queueing them. Only meaningful when `maxConcurrent` is set.                 |
+| `sendTimeout`    | `INIT_SEND_TIMEOUT`   | Duration | `2m`    | Absolute cap on how long a single gated delivery may hold a slot. A ~64 KB/s throughput floor closes a client that stalls or reads slower than the floor well before this. The cap governs for a delivery large enough that even a floor-rate client would exceed it (roughly above `sendTimeout × 64 KB/s`, i.e. ~7.5 MiB at the default), so a client on a very large data set may be cut and reconnect. On expiry the connection is closed so the slot is reclaimed and the SDK reconnects. Only meaningful when `maxConcurrent` is set. |
+
+When the budget is full, a polling request is shed with an HTTP `503` and a `Retry-After` header; a streaming request, whose response has already started, has its connection closed so the SDK reconnects with backoff.
+
+
 ### File section: `[Environment "NAME"]`
 
 The Relay Proxy allows you to proxy any number of LaunchDarkly environments; there must be at least one. In a configuration file, each of these is a separate section in the format `[Environment "MyEnvName"]`, where `MyEnvName` is a unique identifier for the environment (this does not have to match the environment name on your LaunchDarkly dashboard, but it is recommended to). If you are using environment variables, you will add the `MyEnvName` identifier to the variable name prefix for each property. See examples below.

internal/middleware/concurrency.go

@@ -0,0 +1,112 @@
+package middleware
+
+import (
+	"context"
+	"math/rand/v2"
+	"net/http"
+	"strconv"
+	"time"
+
+	"github.com/launchdarkly/ld-relay/v9/internal/concurrency"
+	"github.com/launchdarkly/ld-relay/v9/internal/initwrite"
+	"github.com/launchdarkly/ld-relay/v9/internal/util"
+)
+
+type initLimiterCtxKey struct{}
+
+type initLimiterHolder struct {
+	limiter *concurrency.Limiter
+	maxHold time.Duration
+}
+
+// AcquireInitSlot draws one slot from the shared initialization-delivery limiter for the
+// current request, holding it until the returned release is called. If the budget is full it
+// writes a 503 (with a jittered Retry-After and a JSON body) and returns ok=false; the caller
+// must then write nothing further. A disabled or nil limiter always admits with a no-op
+// release, so callers can invoke this unconditionally.
+//
+// The response write itself is bounded separately, by the progress-aware writer that
+// LimitConcurrency / ProvideInitLimiter install, so a slow client cannot park the slot.
+//
+// It is called both by LimitConcurrency (which gates a whole handler, e.g. the FDv1 all-flags
+// poll) and directly by the FDv2 poll handlers, which acquire only on the full-basis branch so
+// a cheap up-to-date reply is never charged.
+func AcquireInitSlot(limiter *concurrency.Limiter, w http.ResponseWriter, r *http.Request) (release func(), ok bool) {
+	if !limiter.Enabled() {
+		return func() {}, true
+	}
+	rel, ok := limiter.Acquire(r.Context())
+	if !ok {
+		w.Header().Set("Content-Type", "application/json")
+		w.Header().Set("Retry-After", strconv.Itoa(retryAfterSeconds()))
+		w.WriteHeader(http.StatusServiceUnavailable)
+		_, _ = w.Write(util.ErrorJSONMsg("relay initialization concurrency limit reached; retry shortly"))
+		return func() {}, false
+	}
+	return rel, true
+}
+
+// LimitConcurrency wraps a handler with the shared initialization-delivery limiter for a
+// request whose whole response is a full-dataset delivery (the FDv1 all-flags poll). The slot
+// is acquired on entry and held until the handler returns, and the response is written through
+// a progress-aware writer (see initwrite) so a slow or stalled client cannot park the slot. On
+// shed it responds 503 and does not invoke the wrapped handler. A disabled or nil limiter is a
+// pass-through with zero overhead. Handlers that have a cheap no-payload branch (the FDv2
+// polls) should instead call AcquireInitSlot directly, after that branch, so the cheap reply
+// is not charged.
+func LimitConcurrency(limiter *concurrency.Limiter, maxHold time.Duration) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		if !limiter.Enabled() {
+			return next
+		}
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			release, ok := AcquireInitSlot(limiter, w, r)
+			if !ok {
+				return
+			}
+			defer release()
+			defer clearWriteDeadline(w)
+			next.ServeHTTP(initwrite.Wrap(w, maxHold), r)
+		})
+	}
+}
+
+// clearWriteDeadline removes any write deadline the progress-aware writer armed on the
+// connection, so it cannot linger on a kept-alive connection and fire during a later request.
+// This server sets no http.Server.WriteTimeout, so net/http does not reset it for us.
+func clearWriteDeadline(w http.ResponseWriter) {
+	_ = http.NewResponseController(w).SetWriteDeadline(time.Time{})
+}
+
+// ProvideInitLimiter makes the shared initialization-delivery limiter available to a
+// downstream handler via the request context, without acquiring a slot, and wraps the response
+// in the progress-aware writer. It is used for the FDv2 poll endpoints, whose handlers acquire
+// lazily (via AcquireInitSlotFromContext) only on the full-basis branch, so a cheap up-to-date
+// reply is never charged. A disabled or nil limiter is a pass-through with zero overhead.
+func ProvideInitLimiter(limiter *concurrency.Limiter, maxHold time.Duration) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		if !limiter.Enabled() {
+			return next
+		}
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			defer clearWriteDeadline(w)
+			ctx := context.WithValue(r.Context(), initLimiterCtxKey{}, initLimiterHolder{limiter: limiter, maxHold: maxHold})
+			next.ServeHTTP(initwrite.Wrap(w, maxHold), r.WithContext(ctx))
+		})
+	}
+}
+
+// AcquireInitSlotFromContext acquires a slot from the limiter installed by ProvideInitLimiter,
+// with the same semantics as AcquireInitSlot. If no limiter was provided (or it is disabled)
+// it admits with a no-op release, so a handler can call it unconditionally on its full-basis
+// path.
+func AcquireInitSlotFromContext(w http.ResponseWriter, r *http.Request) (release func(), ok bool) {
+	holder, _ := r.Context().Value(initLimiterCtxKey{}).(initLimiterHolder)
+	return AcquireInitSlot(holder.limiter, w, r)
+}
+
+// retryAfterSeconds returns a small jittered Retry-After (in seconds) so a shed herd does
+// not retry in lockstep and re-synchronize the next burst.
+func retryAfterSeconds() int {
+	return 2 + rand.IntN(4) //nolint:gosec // Retry-After jitter is not security-sensitive; a fast PRNG is fine. 2..5 seconds
+}

internal/middleware/concurrency_test.go

@@ -0,0 +1,117 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/launchdarkly/ld-relay/v9/internal/concurrency"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLimitConcurrencyDisabledIsPassThrough(t *testing.T) {
+	limiter := concurrency.New("t", concurrency.Params{MaxConcurrent: 0}) // disabled
+	called := false
+	h := LimitConcurrency(limiter, time.Second)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		called = true
+		w.WriteHeader(http.StatusOK)
+	}))
+	rr := httptest.NewRecorder()
+	h.ServeHTTP(rr, httptest.NewRequest("GET", "/sdk/flags", nil))
+	assert.True(t, called)
+	assert.Equal(t, http.StatusOK, rr.Code)
+}
+
+func TestLimitConcurrencyShedsWhenBudgetFull(t *testing.T) {
+	limiter := concurrency.New("t", concurrency.Params{MaxConcurrent: 1, MaxQueued: 0})
+	// Occupy the only slot so the wrapped request must shed.
+	release, ok := limiter.Acquire(context.Background())
+	require.True(t, ok)
+	defer release()
+
+	called := false
+	h := LimitConcurrency(limiter, time.Second)(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
+		called = true
+	}))
+	rr := httptest.NewRecorder()
+	h.ServeHTTP(rr, httptest.NewRequest("GET", "/sdk/flags", nil))
+
+	assert.False(t, called, "handler must not run when shed")
+	assert.Equal(t, http.StatusServiceUnavailable, rr.Code)
+	assert.Equal(t, "application/json", rr.Header().Get("Content-Type"))
+	assert.Contains(t, rr.Body.String(), "concurrency limit reached")
+
+	// Retry-After is a small positive integer, jittered into a range so a shed herd does not
+	// retry in lockstep.
+	ra, err := strconv.Atoi(rr.Header().Get("Retry-After"))
+	require.NoError(t, err)
+	assert.GreaterOrEqual(t, ra, 2)
+	assert.LessOrEqual(t, ra, 5)
+}
+
+func TestLimitConcurrencyReleasesSlotWhenHandlerReturns(t *testing.T) {
+	limiter := concurrency.New("t", concurrency.Params{MaxConcurrent: 1, MaxQueued: 0})
+	h := LimitConcurrency(limiter, time.Second)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+	// Two sequential requests both succeed, which is only possible if the single slot is
+	// released when each handler returns.
+	for range 2 {
+		rr := httptest.NewRecorder()
+		h.ServeHTTP(rr, httptest.NewRequest("GET", "/sdk/flags", nil))
+		assert.Equal(t, http.StatusOK, rr.Code)
+	}
+	assert.Equal(t, int64(2), limiter.Stats().Admitted)
+}
+
+func TestAcquireInitSlotFromContextChargesOnlyWhenProvided(t *testing.T) {
+	limiter := concurrency.New("t", concurrency.Params{MaxConcurrent: 1, MaxQueued: 0})
+
+	// With no limiter installed in the context, callers are admitted with a no-op release and
+	// the budget is untouched. This is what lets an FDv2 up-to-date reply avoid a slot.
+	rr := httptest.NewRecorder()
+	release, ok := AcquireInitSlotFromContext(rr, httptest.NewRequest("GET", "/sdk/poll", nil))
+	require.True(t, ok)
+	release()
+	assert.Equal(t, int64(0), limiter.Stats().Admitted, "absent limiter must not charge a slot")
+
+	// When ProvideInitLimiter installed the limiter, a full-basis handler charges a slot.
+	provided := ProvideInitLimiter(limiter, time.Second)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		rel, ok := AcquireInitSlotFromContext(w, r)
+		require.True(t, ok)
+		defer rel()
+		w.WriteHeader(http.StatusOK)
+	}))
+	rr2 := httptest.NewRecorder()
+	provided.ServeHTTP(rr2, httptest.NewRequest("GET", "/sdk/poll", nil))
+	assert.Equal(t, http.StatusOK, rr2.Code)
+	assert.Equal(t, int64(1), limiter.Stats().Admitted)
+}
+
+func TestAcquireInitSlotFromContextShedsWhenBudgetFull(t *testing.T) {
+	limiter := concurrency.New("t", concurrency.Params{MaxConcurrent: 1, MaxQueued: 0})
+	release, ok := limiter.Acquire(context.Background())
+	require.True(t, ok)
+	defer release()
+
+	handlerRan := false
+	provided := ProvideInitLimiter(limiter, time.Second)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		rel, ok := AcquireInitSlotFromContext(w, r)
+		if !ok {
+			return
+		}
+		defer rel()
+		handlerRan = true
+		w.WriteHeader(http.StatusOK)
+	}))
+	rr := httptest.NewRecorder()
+	provided.ServeHTTP(rr, httptest.NewRequest("GET", "/sdk/poll", nil))
+
+	assert.False(t, handlerRan, "full-basis handler must not proceed once shed")
+	assert.Equal(t, http.StatusServiceUnavailable, rr.Code)
+}

internal/streams/sse_logger_test.go

@@ -0,0 +1,51 @@
+package streams
+
+import (
+	"context"
+	"errors"
+	"log/slog"
+	"os"
+	"testing"
+
+	"github.com/launchdarkly/ld-relay/v9/internal/basictypes"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type capturingHandler struct{ recs *[]slog.Record }
+
+func (h capturingHandler) Enabled(context.Context, slog.Level) bool { return true }
+func (h capturingHandler) Handle(_ context.Context, r slog.Record) error {
+	*h.recs = append(*h.recs, r)
+	return nil
+}
+func (h capturingHandler) WithAttrs([]slog.Attr) slog.Handler { return h }
+func (h capturingHandler) WithGroup(string) slog.Handler      { return h }
+
+func TestSSELoggerDistinguishesCutFromDisconnect(t *testing.T) {
+	var recs []slog.Record
+	l := sseLogger{log: slog.New(capturingHandler{&recs})}
+
+	// A write-deadline cut (the limiter reclaiming a slot) surfaces as a deadline error and is
+	// logged at warn; an ordinary client disconnect is logged at debug.
+	l.Println(os.ErrDeadlineExceeded)
+	l.Println(errors.New("write: broken pipe"))
+
+	require.Len(t, recs, 2)
+	assert.Equal(t, slog.LevelWarn, recs[0].Level)
+	assert.Contains(t, recs[0].Message, "write deadline exceeded")
+	assert.Equal(t, slog.LevelDebug, recs[1].Level)
+}
+
+func TestWithLoggerSetsServerSideSSELogger(t *testing.T) {
+	sp := NewStreamProvider(basictypes.ServerSideStream, 0, 0, WithLogger(slog.Default())).(*serverSideStreamProvider)
+	defer sp.Close()
+	assert.NotNil(t, sp.fdv1Server.Logger, "server-side FDv1 SSE server should have a logger set")
+	assert.NotNil(t, sp.fdv2Server.Logger, "server-side FDv2 SSE server should have a logger set")
+
+	// Without WithLogger, the logger stays unset (unchanged default behavior).
+	sp2 := NewStreamProvider(basictypes.ServerSideStream, 0, 0).(*serverSideStreamProvider)
+	defer sp2.Close()
+	assert.Nil(t, sp2.fdv1Server.Logger)
+}