fix: Local evaluation silently drops identity overrides on large environments

Flagsmith.Client.Test/FlagsmithTest.cs

@@ -246,6 +246,56 @@ public async Task TestGetIdentityFlagsUsesLocalIdentityOverridesWhenAvailable()
             Assert.Equal("some-overridden-value", flag.Value);
         }
         [Fact]
+        public async Task TestLocalEvaluationFollowsEnvironmentDocumentPages()
+        {
+            // Given
+            HttpResponseMessage Page(JObject document, string nextPageId)
+            {
+                var response = new HttpResponseMessage
+                {
+                    StatusCode = HttpStatusCode.OK,
+                    Content = new StringContent(document.ToString())
+                };
+                if (nextPageId != null)
+                {
+                    response.Headers.Add("Link", $"</api/v1/environment-document/?page_id={Uri.EscapeDataString(nextPageId)}>; rel=\"next\"");
+                }
+                return response;
+            }
+            JObject OverridePage(string identifier, string value)
+            {
+                var identityOverride = (JObject)Fixtures.JsonObject["identity_overrides"][0].DeepClone();
+                identityOverride["identifier"] = identifier;
+                identityOverride["identity_features"][0]["feature_state_value"] = value;
+                return new JObject { ["identity_overrides"] = new JArray(identityOverride) };
+            }
+            var responses = new Dictionary<string, HttpResponseMessage>
+            {
+                { "/api/v1/environment-document/", Page(Fixtures.JsonObject, "identity_override:1:page-2") },
+                { "/api/v1/environment-document/?page_id=identity_override%3A1%3Apage-2", Page(OverridePage("page-2-id", "page-2-value"), "identity_override:1:page-3") },
+                { "/api/v1/environment-document/?page_id=identity_override%3A1%3Apage-3", Page(OverridePage("page-3-id", "page-3-value"), null) }
+            };
+            var mockHttpClient = HttpMocker.MockHttpResponse(responses);
+
+            // When
+            var flagsmithClientTest = new FlagsmithClient(new FlagsmithConfiguration
+            {
+                EnvironmentKey = Fixtures.ApiKey,
+                HttpClient = mockHttpClient.Object,
+                EnableLocalEvaluation = true
+            });
+
+            // Then
+            mockHttpClient.Verify(x => x.SendAsync(It.IsAny<HttpRequestMessage>(), It.IsAny<CancellationToken>()), Times.Exactly(3));
+            foreach (var (identifier, value) in new[] { ("overridden-id", "some-overridden-value"), ("page-2-id", "page-2-value"), ("page-3-id", "page-3-value") })
+            {
+                var flag = await (await flagsmithClientTest.GetIdentityFlags(identifier)).GetFlag("some_feature");
+                Assert.Equal(value, flag.Value);
+            }
+            var environmentFlag = await (await flagsmithClientTest.GetEnvironmentFlags()).GetFlag("some_feature");
+            Assert.Equal("some-value", environmentFlag.Value);
+        }
+        [Fact]
         public async Task TestRequestConnectionErrorRaisesFlagsmithApiError()
         {
             var mockHttpClient = HttpMocker.MockHttpThrowConnectionError();

Flagsmith.FlagsmithClient/FlagsmithClient.cs

@@ -3,9 +3,12 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Net.Http;
+using System.Net.Http.Headers;
 using System.Text;
+using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
 using Flagsmith.Cache;
@@ -196,11 +199,17 @@ private IdentityFlagListCache GetFlagListCacheByIdentity(IdentityWrapper identit
         }
 
         private async Task<string> GetJson(HttpMethod method, string url, string? body = null)
+        {
+            var (json, _) = await GetJsonAndHeaders(method, url, body).ConfigureAwait(false);
+            return json;
+        }
+
+        private async Task<(string Json, HttpResponseHeaders Headers)> GetJsonAndHeaders(HttpMethod method, string url, string? body = null)
         {
             try
             {
                 var policy = HttpPolicies.GetRetryPolicyAwaitable(_config.Retries);
-                return await (await policy.ExecuteAsync(async () =>
+                var response = await policy.ExecuteAsync(async () =>
                 {
                     HttpRequestMessage request = new HttpRequestMessage(method, url)
                     {
@@ -219,7 +228,8 @@ private async Task<string> GetJson(HttpMethod method, string url, string? body =
                     var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(_config.RequestTimeout ?? 100));
                     HttpResponseMessage response = await _config.HttpClient.SendAsync(request, cancellationTokenSource.Token).ConfigureAwait(false);
                     return response.EnsureSuccessStatusCode();
-                }).ConfigureAwait(false)).Content.ReadAsStringAsync().ConfigureAwait(false);
+                }).ConfigureAwait(false);
+                return (await response.Content.ReadAsStringAsync().ConfigureAwait(false), response.Headers);
             }
             catch (HttpRequestException e)
             {
@@ -233,12 +243,41 @@ private async Task<string> GetJson(HttpMethod method, string url, string? body =
             }
         }
 
+        private static string? GetNextPageId(HttpResponseHeaders headers)
+        {
+            if (!headers.TryGetValues("Link", out var links))
+            {
+                return null;
+            }
+            var match = Regex.Match(string.Join(",", links), "<[^>]*[?&]page_id=([^&>]*)[^>]*>;\\s*rel=\"next\"");
+            return match.Success ? Uri.UnescapeDataString(match.Groups[1].Value) : null;
+        }
+
         private async Task GetAndUpdateEnvironmentFromApi()
         {
             try
             {
-                var json = await GetJson(HttpMethod.Get, new Uri(_config.ApiUri, "environment-document/").AbsoluteUri).ConfigureAwait(false);
-                Environment = JsonConvert.DeserializeObject<EnvironmentModel>(json);
+                var stopwatch = Stopwatch.StartNew();
+                var environmentUrl = new Uri(_config.ApiUri, "environment-document/").AbsoluteUri;
+                var (json, headers) = await GetJsonAndHeaders(HttpMethod.Get, environmentUrl).ConfigureAwait(false);
+                var environment = JsonConvert.DeserializeObject<EnvironmentModel>(json);
+                var pageId = GetNextPageId(headers);
+                while (pageId != null)
+                {
+                    var (pageJson, pageHeaders) = await GetJsonAndHeaders(HttpMethod.Get, $"{environmentUrl}?page_id={Uri.EscapeDataString(pageId)}").ConfigureAwait(false);
+                    var page = JsonConvert.DeserializeObject<EnvironmentModel>(pageJson);
+                    environment!.IdentityOverrides ??= new List<IdentityModel>();
+                    environment.IdentityOverrides.AddRange(page?.IdentityOverrides ?? new List<IdentityModel>());
+                    pageId = GetNextPageId(pageHeaders);
+                }
+                if (stopwatch.Elapsed > _config.EnvironmentRefreshInterval)
+                {
+                    _config.Logger?.LogWarning(
+                        "Fetching the environment document took {Elapsed}, longer than the environment refresh interval of {RefreshInterval}; raise the refresh interval or reduce the environment size.",
+                        stopwatch.Elapsed,
+                        _config.EnvironmentRefreshInterval);
+                }
+                Environment = environment;
                 _evaluationContext = Mappers.MapEnvironmentDocumentToContext(Environment);
                 _config.Logger?.LogInformation("Local Environment updated: " + json);
             }