launchdarkly · kinyoklion · Mar 20, 2025 · Mar 4, 2025 · Mar 4, 2025 · Mar 4, 2025
@@ -25,18 +25,26 @@ public sealed class EvaluationSeriesContext {
         /// </summary>
         public string Method { get;  }
 
+        /// <summary>
+        /// The environment ID for the evaluation, or null if not available.
+        /// </summary>
+        public string EnvironmentId { get; }
+
         /// <summary>
         /// Constructs a new EvaluationSeriesContext.
         /// </summary>
         /// <param name="flagKey">the flag key</param>
         /// <param name="context">the context</param>
         /// <param name="defaultValue">the default value</param>
         /// <param name="method">the variation method</param>
-        public EvaluationSeriesContext(string flagKey, Context context, LdValue defaultValue, string method) {
+        /// <param name="environmentId">the environment ID</param>
+        public EvaluationSeriesContext(string flagKey, Context context, LdValue defaultValue, string method,
+            string environmentId = null) {
             FlagKey = flagKey;
             Context = context;
             DefaultValue = defaultValue;
             Method = method;
+            EnvironmentId = environmentId;
         }
     }
 }
@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Threading.Tasks;
@@ -22,7 +23,7 @@ namespace LaunchDarkly.Sdk.Server.Internal.DataSources
     /// This component is also responsible for receiving updates to the data source status, broadcasting
     /// them to any status listeners, and tracking the length of any period of sustained failure.
     /// </remarks>
-    internal sealed class DataSourceUpdatesImpl : IDataSourceUpdates
+    internal sealed class DataSourceUpdatesImpl : IDataSourceUpdates, IDataSourceUpdatesHeaders
     {
         #region Private fields
 
@@ -84,7 +85,7 @@ internal DataSourceUpdatesImpl(
                 StateSince = DateTime.Now,
                 LastError = null
             };
-            _status = new StateMonitor<DataSourceStatus, StateAndError>(initialStatus, MaybeUpdateStatus, _log);            
+            _status = new StateMonitor<DataSourceStatus, StateAndError>(initialStatus, MaybeUpdateStatus, _log);
         }
 
         #endregion
@@ -93,47 +94,7 @@ internal DataSourceUpdatesImpl(
 
         public bool Init(FullDataSet<ItemDescriptor> allData)
         {
-            ImmutableDictionary<DataKind, ImmutableDictionary<string, ItemDescriptor>> oldData = null;
-
-            try
-            {
-                if (HasFlagChangeListeners())
-                {
-                    // Query the existing data if any, so that after the update we can send events for
-                    // whatever was changed
-                    var oldDataBuilder = ImmutableDictionary.CreateBuilder<DataKind,
-                        ImmutableDictionary<string, ItemDescriptor>>();
-                    foreach (var kind in DataModel.AllDataKinds)
-                    {
-                        var items = _store.GetAll(kind);
-                        oldDataBuilder.Add(kind, items.Items.ToImmutableDictionary());
-                    }
-                    oldData = oldDataBuilder.ToImmutable();
-                }
-                _store.Init(DataStoreSorter.SortAllCollections(allData));
-                _lastStoreUpdateFailed = false;
-            }
-            catch (Exception e)
            {
-                ReportStoreFailure(e);
-                return false;
-            }
-
-            // Calling Init implies that the data source is now in a valid state.
-            UpdateStatus(DataSourceState.Valid, null);
-
-            // We must always update the dependency graph even if we don't currently have any event listeners, because if
-            // listeners are added later, we don't want to have to reread the whole data store to compute the graph
-            UpdateDependencyTrackerFromFullDataSet(allData);
-
-            // Now, if we previously queried the old data because someone is listening for flag change events, compare
-            // the versions of all items and generate events for those (and any other items that depend on them)
-            if (oldData != null)
-            {
-                SendChangeEvents(ComputeChangedItemsForFullDataSet(oldData, FullDataSetToMap(allData)));
-            }
-
-            return true;
+            return InitWithHeaders(allData, null);
         }
 
         public bool Upsert(DataKind kind, string key, ItemDescriptor item)
@@ -326,5 +287,67 @@ private void ReportStoreFailure(Exception e)
         }
 
         #endregion
+
+       #region IDataSourceUpdatesHeaders methods
+        public bool InitWithHeaders(FullDataSet<ItemDescriptor> allData, IEnumerable<KeyValuePair<string, IEnumerable<string>>> headers)
+        {
+            ImmutableDictionary<DataKind, ImmutableDictionary<string, ItemDescriptor>> oldData = null;
+
+            try
+            {
+                if (HasFlagChangeListeners())
+                {
+                    // Query the existing data if any, so that after the update we can send events for
+                    // whatever was changed
+                    var oldDataBuilder = ImmutableDictionary.CreateBuilder<DataKind,
+                        ImmutableDictionary<string, ItemDescriptor>>();
+                    foreach (var kind in DataModel.AllDataKinds)
+                    {
+                        var items = _store.GetAll(kind);
+                        oldDataBuilder.Add(kind, items.Items.ToImmutableDictionary());
+                    }
+                    oldData = oldDataBuilder.ToImmutable();
+                }
+
+                var sortedCollections = DataStoreSorter.SortAllCollections(allData);
+
+                if (_store is IDataStoreMetadata storeMetadata)
+                {
+                    var environmentId = headers?.FirstOrDefault((item) =>
+                            item.Key.ToLower() == HeaderConstants.EnvironmentId).Value
+                        ?.FirstOrDefault();
+                    storeMetadata.InitWithMetadata(sortedCollections, new InitMetadata(environmentId));
+                }
+                else
+                {
+                    _store.Init(sortedCollections);
+                }
+
+                _lastStoreUpdateFailed = false;
+            }
+            catch (Exception e)
+            {
+                ReportStoreFailure(e);
+                return false;
+            }
+
+            // Calling Init implies that the data source is now in a valid state.
+            UpdateStatus(DataSourceState.Valid, null);
+
+            // We must always update the dependency graph even if we don't currently have any event listeners, because if
+            // listeners are added later, we don't want to have to reread the whole data store to compute the graph
+            UpdateDependencyTrackerFromFullDataSet(allData);
+
+            // Now, if we previously queried the old data because someone is listening for flag change events, compare
+            // the versions of all items and generate events for those (and any other items that depend on them)
+            if (oldData != null)
+            {
+                SendChangeEvents(ComputeChangedItemsForFullDataSet(oldData, FullDataSetToMap(allData)));
+            }
+
+            return true;
+        }
+        #endregion
+
     }
 }
@@ -17,6 +17,8 @@
 
 namespace LaunchDarkly.Sdk.Server.Internal.DataSources
 {
+    using BytesWithHeaders = Tuple<byte[], IEnumerable<KeyValuePair<string, IEnumerable<string>>>>;
+
     internal class FeatureRequestor : IFeatureRequestor
     {
         private readonly Uri _allUri;
@@ -51,19 +53,19 @@ private void Dispose(bool disposing)
 
         // Returns a data set of the latest flags and segments, or null if they have not been modified. Throws an
         // exception if there was a problem getting data.
-        public async Task<FullDataSet<ItemDescriptor>?> GetAllDataAsync()
+        public async Task<DataSetWithHeaders> GetAllDataAsync()
         {
-            var json = await GetAsync(_allUri);
-            if (json is null)
+            var res = await GetAsync(_allUri);
+            if (res is null)
             {
                 return null;
             }
-            var data = ParseAllData(json);
+            var data = ParseAllData(res.Item1);
             Func<DataKind, int> countItems = kind =>
                 data.Data.FirstOrDefault(kv => kv.Key == kind).Value.Items?.Count() ?? 0;
             _log.Debug("Get all returned {0} feature flags and {1} segments",
                 countItems(DataModel.Features), countItems(DataModel.Segments));
-            return data;
+            return new DataSetWithHeaders(data, res.Item2);
         }
 
         private FullDataSet<ItemDescriptor> ParseAllData(byte[] json)
@@ -72,7 +74,7 @@ private FullDataSet<ItemDescriptor> ParseAllData(byte[] json)
             return StreamProcessorEvents.ParseFullDataset(ref r);
         }
 
-        private async Task<byte[]> GetAsync(Uri path)
+        private async Task<BytesWithHeaders> GetAsync(Uri path)
         {
             _log.Debug("Getting flags with uri: {0}", path.AbsoluteUri);
             var request = new HttpRequestMessage(HttpMethod.Get, path);
@@ -113,7 +115,7 @@ private async Task<byte[]> GetAsync(Uri path)
                             }
                         }
                         var content = await response.Content.ReadAsByteArrayAsync().ConfigureAwait(false);
-                        return content.Length == 0 ? null : content;
+                        return new BytesWithHeaders(content.Length == 0 ? null : content, response.Headers);
                     }
                 }
                 catch (TaskCanceledException tce)

@@ -1,12 +1,26 @@
 using System;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 
 using static LaunchDarkly.Sdk.Server.Subsystems.DataStoreTypes;
 
 namespace LaunchDarkly.Sdk.Server.Internal.DataSources
 {
+    internal class DataSetWithHeaders
+    {
+        public readonly FullDataSet<ItemDescriptor>? DataSet;
+        public readonly IEnumerable<KeyValuePair<string, IEnumerable<string>>> Headers;
+
+        public DataSetWithHeaders(FullDataSet<ItemDescriptor>? dataSet,
+            IEnumerable<KeyValuePair<string, IEnumerable<string>>> headers)
+        {
+            DataSet = dataSet;
+            Headers = headers;
+        }
+    }
+
     internal interface IFeatureRequestor : IDisposable
     {
-        Task<FullDataSet<ItemDescriptor>?> GetAllDataAsync();
+        Task<DataSetWithHeaders> GetAllDataAsync();
     }
 }
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -60,15 +61,15 @@ private async Task UpdateTaskAsync()
             _log.Info("Polling LaunchDarkly for feature flag updates");
             try
             {
-                var allData = await _featureRequestor.GetAllDataAsync();
-                if (allData is null)
+                var dataAndHeaders = await _featureRequestor.GetAllDataAsync();
+                if (dataAndHeaders.DataSet is null)
                 {
                     // This means it was cached, and alreadyInited was true
                     _dataSourceUpdates.UpdateStatus(DataSourceState.Valid, null);
                 }
                 else
                 {
-                    if (_dataSourceUpdates.Init(allData.Value)) // this also automatically sets the state to Valid
+                    if (InitWithHeaders(dataAndHeaders.DataSet.Value, dataAndHeaders.Headers)) // this also automatically sets the state to Valid
                     {
                         if (!_initialized.GetAndSet(true))
                         {
@@ -137,5 +138,16 @@ private void Dispose(bool disposing)
                 _featureRequestor.Dispose();
             }
         }
+
+        private bool InitWithHeaders(DataStoreTypes.FullDataSet<DataStoreTypes.ItemDescriptor> allData,
+            IEnumerable<KeyValuePair<string, IEnumerable<string>>> headers)
+        {
+            if (_dataSourceUpdates is IDataSourceUpdatesHeaders dataSourceUpdatesHeaders)
+            {
+                return dataSourceUpdatesHeaders.InitWithHeaders(allData, headers);
+            }
+
+            return _dataSourceUpdates.Init(allData);
+        }
     }
-}
+}
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Net.Http;
 using System.Text.Json;
@@ -43,6 +44,8 @@ internal class StreamingDataSource : IDataSource
         private volatile bool _lastStoreUpdateFailed = false;
         internal DateTime _esStarted; // exposed for testing
 
+        private IEnumerable<KeyValuePair<string, IEnumerable<string>>> _headers;
+
         internal delegate IEventSource EventSourceCreator(Uri streamUri,
             HttpConfiguration httpConfig);
 
@@ -139,6 +142,7 @@ private void RecordStreamInit(bool failed)
 
         private void OnOpen(object sender, EventSource.StateChangedEventArgs e)
         {
+            _headers = e.Headers;
             _log.Debug("EventSource Opened");
             RecordStreamInit(false);
         }
@@ -238,13 +242,24 @@ private void OnError(object sender, EventSource.ExceptionEventArgs e)
             }
         }
 
+        private bool InitWithHeaders(FullDataSet<ItemDescriptor> allData,
+            IEnumerable<KeyValuePair<string, IEnumerable<string>>> headers)
+        {
+            if (_dataSourceUpdates is IDataSourceUpdatesHeaders dataSourceUpdatesHeaders)
+            {
+                return dataSourceUpdatesHeaders.InitWithHeaders(allData, headers);
+            }
+
+            return _dataSourceUpdates.Init(allData);
+        }
+
         private void HandleMessage(string messageType, byte[] messageData)
         {
             switch (messageType)
             {
                 case PUT:
                     var putData = ParsePutData(messageData);
-                    if (!_dataSourceUpdates.Init(putData.Data)) // this also automatically sets the state to Valid
+                    if (!InitWithHeaders(putData.Data, _headers)) // this also automatically sets the state to Valid
                     {
                         throw new StreamStoreException("failed to write full data set to data store");
                     }