fix: paginate environment-document for local evaluation

Author: higordasneves

client.go

@@ -2,9 +2,11 @@ package flagsmith
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log/slog"
 	"net/http"
+	"net/url"
 	"reflect"
 	"runtime"
 	"strings"
@@ -14,6 +16,7 @@ import (
 	"github.com/Flagsmith/flagsmith-go-client/v5/flagengine"
 	"github.com/Flagsmith/flagsmith-go-client/v5/flagengine/engine_eval"
 	"github.com/Flagsmith/flagsmith-go-client/v5/flagengine/environments"
+	"github.com/Flagsmith/flagsmith-go-client/v5/flagengine/identities"
 	"github.com/Flagsmith/flagsmith-go-client/v5/flagengine/segments"
 	"github.com/go-resty/resty/v2"
 )
@@ -433,28 +436,61 @@ func (c *Client) pollThenStartRealtime(ctx context.Context) {
 
 func (c *Client) UpdateEnvironment(ctx context.Context) error {
 	var env environments.EnvironmentModel
-	resp, err := c.client.NewRequest().
-		SetContext(ctx).
-		SetResult(&env).
-		ForceContentType("application/json").
-		Get(c.config.baseURL + "environment-document/")
+	nextPage := ""
+	pageCount := 0
+	var totalOverridesBytes int
 
-	if err != nil {
-		msg := fmt.Sprintf("flagsmith: error performing request to Flagsmith API: %s", err)
-		f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
-		if c.errorHandler != nil {
-			c.errorHandler(f)
+	for {
+		var page environments.EnvironmentModel
+		req := c.client.NewRequest().
+			SetContext(ctx).
+			SetResult(&page).
+			ForceContentType("application/json")
+		if nextPage != "" {
+			req = req.SetQueryParam("page_id", nextPage)
 		}
-		return f
-	}
-	if resp.StatusCode() != 200 {
-		msg := fmt.Sprintf("flagsmith: unexpected response from Flagsmith API: %s", resp.Status())
-		f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
-		if c.errorHandler != nil {
-			c.errorHandler(f)
+
+		resp, err := req.Get(c.config.baseURL + "environment-document/")
+		if err != nil {
+			msg := fmt.Sprintf("flagsmith: error performing request to Flagsmith API: %s", err)
+			f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
+			if c.errorHandler != nil {
+				c.errorHandler(f)
+			}
+			return f
+		}
+		if resp.StatusCode() != 200 {
+			msg := fmt.Sprintf("flagsmith: unexpected response from Flagsmith API: %s", resp.Status())
+			f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
+			if c.errorHandler != nil {
+				c.errorHandler(f)
+			}
+			return f
+		}
+
+		pageCount++
+		nextPage = c.ExtractNextPage(resp.Header().Get("link"))
+
+		if pageCount == 1 {
+			env = page
+		} else {
+			ok, newTotal := c.checkEnvironmentMemoryAlloc(page.IdentityOverrides, pageCount, totalOverridesBytes)
+			if !ok {
+				break
+			}
+			totalOverridesBytes = newTotal
+			env.IdentityOverrides = append(env.IdentityOverrides, page.IdentityOverrides...)
+		}
+
+		if nextPage == "" {
+			break
+		}
+		if c.config.localEvaluationPageLimit > 0 && pageCount >= c.config.localEvaluationPageLimit {
+			c.log.Debug("page limit reached, stopping pagination", "limit", c.config.localEvaluationPageLimit)
+			break
 		}
-		return f
 	}
+
 	isNew := false
 	previousEnv := c.environment.Load()
 	if previousEnv == nil || env.UpdatedAt.After(previousEnv.(*environments.EnvironmentModel).UpdatedAt) {
@@ -469,5 +505,52 @@ func (c *Client) UpdateEnvironment(ctx context.Context) error {
 		c.log.Info("environment updated", "environment", env.APIKey, "updated_at", env.UpdatedAt)
 	}
 
+	c.log.Debug("IdentityOverrides", "len", len(env.IdentityOverrides))
+
 	return nil
 }
+
+// ExtractNextPage parses the Link header from the environment-document API and
+// returns the decoded page_id value when a next page exists, or empty string otherwise.
+// Expected format: </api/v1/environment-document/?page_id=xxx>; rel="next"
+func (c *Client) ExtractNextPage(linkHeader string) string {
+	parts := strings.SplitN(linkHeader, ">", 2)
+	if len(parts) == 0 {
+		return ""
+	}
+
+	u, err := url.Parse(strings.TrimPrefix(parts[0], "<"))
+	if err != nil {
+		return ""
+	}
+
+	pageID := u.Query().Get("page_id")
+	c.log.Debug("environment-document next page", "link", linkHeader, "page_id", pageID)
+
+	return pageID
+}
+
+// checkEnvironmentMemoryAlloc checks whether appending a new page's identity overrides
+// would exceed the configured byte limit. Returns ok=false and the unchanged
+// accumulated total when the limit would be breached; otherwise returns ok=true
+// and the updated total. When no limit is configured (0) it always returns ok=true.
+func (c *Client) checkEnvironmentMemoryAlloc(overrides []*identities.IdentityModel, pageCount, accumulated int) (ok bool, newTotal int) {
+	if c.config.localEvaluationMemoryAllocBytes == 0 {
+		return true, accumulated
+	}
+
+	pageBytes, err := json.Marshal(overrides)
+	if err != nil {
+		return true, accumulated
+	}
+	newTotal = accumulated + len(pageBytes)
+	if newTotal > c.config.localEvaluationMemoryAllocBytes {
+		c.log.Warn("memory limit reached, skipping page",
+			"page", pageCount,
+			"limit_bytes", c.config.localEvaluationMemoryAllocBytes,
+			"accumulated_bytes", accumulated,
+		)
+		return false, accumulated
+	}
+	return true, newTotal
+}

client_test.go

@@ -1236,3 +1236,148 @@ func TestCustomClientOptionsShoudPanic(t *testing.T) {
 		})
 	}
 }
+
+func TestExtractNextPage(t *testing.T) {
+	client := flagsmith.NewClient("test-key")
+
+	testCases := []struct {
+		name     string
+		header   string
+		expected string
+	}{
+		{
+			name:     "valid link header with encoded page_id",
+			header:   "</api/v1/environment-document/?page_id=" + fixtures.PageIDEncoded + ">; rel=\"next\"",
+			expected: fixtures.PageID,
+		},
+		{
+			name:     "empty header returns empty string",
+			header:   "",
+			expected: "",
+		},
+		{
+			name:     "header without page_id returns empty string",
+			header:   "</api/v1/environment-document/>; rel=\"next\"",
+			expected: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := client.ExtractNextPage(tc.header)
+			assert.Equal(t, tc.expected, result)
+		})
+	}
+}
+
+func TestUpdateEnvironmentPaginatesIdentityOverrides(t *testing.T) {
+	// Given
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.PaginatedEnvironmentDocumentHandler))
+	defer server.Close()
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey, flagsmith.WithLocalEvaluation(ctx),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"),
+		flagsmith.WithLocalEvaluationPageLimit(0))
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+
+	// Then
+	assert.NoError(t, err)
+
+	// Identity from page 1 should be found
+	flags, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifier, nil)
+	assert.NoError(t, err)
+	enabled, err := flags.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled, "identity from page 1 should have overridden feature disabled")
+
+	// Identity from page 2 should also be found
+	flags2, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifierPage2, nil)
+	assert.NoError(t, err)
+	enabled2, err := flags2.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled2, "identity from page 2 should have overridden feature disabled")
+}
+
+func TestUpdateEnvironmentSinglePageNoLinkHeader(t *testing.T) {
+	// Given
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.EnvironmentDocumentHandler))
+	defer server.Close()
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey, flagsmith.WithLocalEvaluation(ctx),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+
+	// Then — no pagination, identity from page 1 should still work
+	assert.NoError(t, err)
+
+	flags, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifier, nil)
+	assert.NoError(t, err)
+	enabled, err := flags.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled, "identity override should have feature disabled")
+}
+
+func TestUpdateEnvironmentRespectsPageLimit(t *testing.T) {
+	// Given: server has 2 pages but client is limited to 1 (the default)
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.PaginatedEnvironmentDocumentHandler))
+	defer server.Close()
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey, flagsmith.WithLocalEvaluation(ctx),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"),
+		flagsmith.WithLocalEvaluationPageLimit(1))
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+	assert.NoError(t, err)
+
+	// Then: page 1 identity is present
+	flags, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifier, nil)
+	assert.NoError(t, err)
+	enabled, err := flags.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled)
+
+	// Page 2 identity must NOT be present — falls back to env default (enabled=true)
+	flags2, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifierPage2, nil)
+	assert.NoError(t, err)
+	enabled2, err := flags2.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.True(t, enabled2, "page 2 identity should not have been loaded")
+}
+
+func TestUpdateEnvironmentRespectsMemoryAllocLimit(t *testing.T) {
+	// Given: memory limit of 1 byte — page 2 will always exceed it
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.PaginatedEnvironmentDocumentHandler))
+	defer server.Close()
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey, flagsmith.WithLocalEvaluation(ctx),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"),
+		flagsmith.WithLocalEvaluationPageLimit(0),
+		flagsmith.WithLocalEvaluationMemoryAllocLimit(1)) // 1 byte — guarantees page 2 is blocked
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+	assert.NoError(t, err)
+
+	// Then: page 1 identity is present (memory limit only blocks appending subsequent pages)
+	flags, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifier, nil)
+	assert.NoError(t, err)
+	enabled, err := flags.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled)
+
+	// Page 2 identity must NOT be present — memory limit blocked the append
+	flags2, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifierPage2, nil)
+	assert.NoError(t, err)
+	enabled2, err := flags2.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.True(t, enabled2, "page 2 should have been blocked by memory limit")
+}

config.go

@@ -28,15 +28,26 @@ type config struct {
 	useRealtime        bool
 	polling            bool
 	userProvidedClient bool
+
+	// localEvaluationPageLimit caps how many pages of the environment-document
+	// are fetched per update cycle. 0 means unlimited. Default is 1 (first page only)
+	// to preserve existing behaviour.
+	localEvaluationPageLimit int
+
+	// localEvaluationMemoryAllocBytes is the maximum total size (in bytes, measured
+	// as the marshaled JSON of identity_overrides) that may be accumulated across
+	// pages. 0 means no limit. Use multiples of 1024*1024 for MB-based limits.
+	localEvaluationMemoryAllocBytes int
 }
 
 // defaultConfig returns default configuration.
 func defaultConfig() config {
 	return config{
-		baseURL:            DefaultBaseURL,
-		timeout:            DefaultTimeout,
-		envRefreshInterval: time.Second * 60,
-		realtimeBaseUrl:    DefaultRealtimeBaseUrl,
-		userProvidedClient: false,
+		baseURL:                  DefaultBaseURL,
+		timeout:                  DefaultTimeout,
+		envRefreshInterval:       time.Second * 60,
+		realtimeBaseUrl:          DefaultRealtimeBaseUrl,
+		userProvidedClient:       false,
+		localEvaluationPageLimit: 1,
 	}
 }