feat: Add environment ID support for hooks

lib/sdk/server/contract-tests/service/src/main/java/sdktest/Representations.java

@@ -255,6 +255,7 @@ public static class EvaluationSeriesContextParam {
     LDContext context;
     LDValue defaultValue;
     String method;
+    String environmentId;
   }
 
   public static class IdentifyEventParams {

lib/sdk/server/contract-tests/service/src/main/java/sdktest/TestHook.java

@@ -44,6 +44,7 @@ public Map<String, Object> beforeEvaluation(EvaluationSeriesContext seriesContex
     seriesContextParam.context = seriesContext.context;
     seriesContextParam.defaultValue = seriesContext.defaultValue;
     seriesContextParam.method = seriesContext.method;
+    seriesContextParam.environmentId = seriesContext.environmentId;
     params.evaluationSeriesContext = seriesContextParam;
 
     params.evaluationSeriesData = data;
@@ -72,6 +73,7 @@ public Map<String, Object> afterEvaluation(EvaluationSeriesContext seriesContext
     seriesContextParam.context = seriesContext.context;
     seriesContextParam.defaultValue = seriesContext.defaultValue;
     seriesContextParam.method = seriesContext.method;
+    seriesContextParam.environmentId = seriesContext.environmentId;
     params.evaluationSeriesContext = seriesContextParam;
 
     params.evaluationSeriesData = data;

lib/sdk/server/contract-tests/service/src/main/java/sdktest/TestService.java

@@ -34,6 +34,7 @@ public class TestService {
     "event-gzip",
     "event-sampling",
     "filtering",
+    "hook-environment-id",
     "inline-context-all",
     "migrations",
     "optional-event-gzip",

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/DataModelDependencies.java

@@ -131,7 +131,7 @@ public static FullDataSet<ItemDescriptor> sortAllCollections(FullDataSet<ItemDes
       DataKind kind = entry.getKey();
       builder.put(kind, sortCollection(kind, entry.getValue()));
     }
-    return new FullDataSet<>(builder.build().entrySet(), allData.shouldPersist());
+    return new FullDataSet<>(builder.build().entrySet(), allData.shouldPersist(), allData.getEnvironmentId());
   }
 
   /**

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/DataSourceSynchronizerAdapter.java

@@ -131,20 +131,25 @@ public void close()  {
      */
     private static class ConvertingUpdateSink implements DataSourceUpdateSink {
         private final IterableAsyncQueue<FDv2SourceResult> resultQueue;
+        // Only a full data set carries an environment ID; it is retained for subsequent partial updates.
+        private volatile String environmentId = null;
 
         public ConvertingUpdateSink(IterableAsyncQueue<FDv2SourceResult> resultQueue) {
             this.resultQueue = resultQueue;
         }
 
         @Override
         public boolean init(DataStoreTypes.FullDataSet<ItemDescriptor> allData) {
+            if (allData.getEnvironmentId() != null && !allData.getEnvironmentId().isEmpty()) {
+                environmentId = allData.getEnvironmentId();
+            }
             // Convert the full data set into a ChangeSet and emit it
             ChangeSet<Iterable<Map.Entry<DataKind, KeyedItems<ItemDescriptor>>>> changeSet =
                     new ChangeSet<>(
                             ChangeSetType.Full,
                             Selector.EMPTY,
                             allData.getData(),
-                            null,
+                            environmentId,
                             allData.shouldPersist()
                         );
             resultQueue.put(FDv2SourceResult.changeSet(changeSet, false));
@@ -166,7 +171,7 @@ public boolean upsert(DataKind kind, String key, ItemDescriptor item) {
                             ChangeSetType.Partial,
                             Selector.EMPTY,
                             data,
-                            null,
+                            environmentId,
                             true // default to true as this adapter is used for adapting FDv1 data sources which are always persistent
                         );
             resultQueue.put(FDv2SourceResult.changeSet(changeSet, false));

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/DataSourceUpdatesImpl.java

@@ -429,7 +429,8 @@ private boolean applyToLegacyStore(ChangeSet<Iterable<Map.Entry<DataKind, KeyedI
 
   private boolean applyFullChangeSetToLegacyStore(ChangeSet<Iterable<Map.Entry<DataKind, KeyedItems<ItemDescriptor>>>> unsortedChangeset) {
     // Convert ChangeSet to FullDataSet for legacy init path, preserving shouldPersist flag
-    return init(new FullDataSet<>(unsortedChangeset.getData(), unsortedChangeset.shouldPersist()));
+    return init(new FullDataSet<>(unsortedChangeset.getData(), unsortedChangeset.shouldPersist(),
+        unsortedChangeset.getEnvironmentId()));
   }
 
   private boolean applyPartialChangeSetToLegacyStore(ChangeSet<Iterable<Map.Entry<DataKind, KeyedItems<ItemDescriptor>>>> changeSet) {

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/DataSystem.java

@@ -55,6 +55,13 @@ interface DataSystem {
    * @return the data store status provider
    */
   DataStoreStatusProvider getDataStoreStatusProvider();
+
+  /**
+   * Returns the ID of the LaunchDarkly environment the data came from, or null if it is not known.
+   *
+   * @return the environment ID, or null
+   */
+  String getEnvironmentId();
 }
 
 /**

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/DefaultFeatureRequestor.java

@@ -128,7 +128,8 @@ public FullDataSet<ItemDescriptor> getAllData(boolean returnDataEvenIfCached)
 
       JsonReader jr = new JsonReader(response.body().charStream());
       // Polling data from LaunchDarkly should be persisted
-      return new FullDataSet<>(parseFullDataSet(jr), true);
+      return new FullDataSet<>(parseFullDataSet(jr), true,
+          response.header(HeaderConstants.ENVIRONMENT_ID.getHeaderName()));
     }
   }
 }

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/EvaluatorWithHooks.java

@@ -11,6 +11,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * An {@link EvaluatorInterface} that will invoke the evaluation series methods of the provided {@link Hook} when
@@ -21,13 +22,17 @@ class EvaluatorWithHooks implements EvaluatorInterface {
   private final EvaluatorInterface underlyingEvaluator;
   private final List<Hook> hooks;
   private final LDLogger logger;
+  private final Supplier<String> environmentIdSupplier;
 
   /**
-   * @param underlyingEvaluator that will do the actual flag evaluation
-   * @param hooks               that will be invoked at various stages of the evaluation series
-   * @param hooksLogger         that will be used to log
+   * @param underlyingEvaluator   that will do the actual flag evaluation
+   * @param hooks                 that will be invoked at various stages of the evaluation series
+   * @param hooksLogger           that will be used to log
+   * @param environmentIdSupplier provides the environment ID reported by LaunchDarkly, if known
    */
-  EvaluatorWithHooks(EvaluatorInterface underlyingEvaluator, List<Hook> hooks, LDLogger hooksLogger) {
+  EvaluatorWithHooks(EvaluatorInterface underlyingEvaluator, List<Hook> hooks, LDLogger hooksLogger,
+      Supplier<String> environmentIdSupplier) {
+    this.environmentIdSupplier = environmentIdSupplier;
     this.underlyingEvaluator = underlyingEvaluator;
     this.hooks = hooks;
     this.logger = hooksLogger;
@@ -40,7 +45,8 @@ public EvalResultAndFlag evalAndFlag(String method, String featureKey, LDContext
     int size = hooks.size();
     List<Map> seriesDataList = new ArrayList<>(size);
 
-    EvaluationSeriesContext seriesContext = new EvaluationSeriesContext(method, featureKey, context, defaultValue);
+    EvaluationSeriesContext seriesContext = new EvaluationSeriesContext(method, featureKey, context, defaultValue,
+        environmentIdSupplier.get());
     Map<String, Object> emptyMap = Collections.emptyMap();
     for (int i = 0; i < size; i++) {
       Hook currentHook = hooks.get(i);

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/FDv1DataSystem.java

@@ -155,6 +155,11 @@ public DataStoreStatusProvider getDataStoreStatusProvider() {
     return dataStoreStatusProvider;
   }
 
+  @Override
+  public String getEnvironmentId() {
+    return dataStore.getEnvironmentId();
+  }
+
   @Override
   public void close() throws IOException {
     if (disposed) {

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/FDv2DataSystem.java

@@ -226,6 +226,11 @@ public DataStoreStatusProvider getDataStoreStatusProvider() {
     return dataStoreStatusProvider;
   }
 
+  @Override
+  public String getEnvironmentId() {
+    return store.getEnvironmentId();
+  }
+
   @Override
   public void close() throws IOException {
     if (disposed) {

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/InMemoryDataStore.java

@@ -33,10 +33,11 @@ class InMemoryDataStore implements DataStore, TransactionalDataStore, CacheExpor
   private final Object selectorLock = new Object();
   private volatile Selector selector = Selector.EMPTY;
   private volatile boolean shouldPersist = false;
+  private volatile String environmentId = null;
 
   @Override
   public void init(FullDataSet<ItemDescriptor> allData) {
-    applyFullPayload(allData.getData(), null, Selector.EMPTY, allData.shouldPersist());
+    applyFullPayload(allData.getData(), allData.getEnvironmentId(), Selector.EMPTY, allData.shouldPersist());
   }
 
   @Override
@@ -128,7 +129,8 @@ public void apply(ChangeSet<Iterable<Map.Entry<DataKind, KeyedItems<ItemDescript
         applyFullPayload(changeSet.getData(), changeSet.getEnvironmentId(), changeSet.getSelector(), changeSet.shouldPersist());
         break;
       case Partial:
-        applyPartialData(changeSet.getData(), changeSet.getSelector(), changeSet.shouldPersist());
+        applyPartialData(changeSet.getData(), changeSet.getEnvironmentId(), changeSet.getSelector(),
+            changeSet.shouldPersist());
         break;
       case None:
         break;
@@ -152,8 +154,19 @@ private void setSelector(Selector newSelector) {
     }
   }
 
+  @Override
+  public String getEnvironmentId() {
+    return this.environmentId;
+  }
+
+  private void setEnvironmentId(String newEnvironmentId) {
+    if (newEnvironmentId != null && !newEnvironmentId.isEmpty()) {
+      this.environmentId = newEnvironmentId;
+    }
+  }
+
   private void applyPartialData(Iterable<Map.Entry<DataKind, KeyedItems<ItemDescriptor>>> data,
-      Selector selector, boolean shouldPersist) {
+      String environmentId, Selector selector, boolean shouldPersist) {
     synchronized (this.writeLock) {
       // Build the complete updated dictionary before assigning to Items for transactional update
       ImmutableMap.Builder<DataKind, Map<String, ItemDescriptor>> itemsBuilder = ImmutableMap.builder();
@@ -192,6 +205,7 @@ private void applyPartialData(Iterable<Map.Entry<DataKind, KeyedItems<ItemDescri
 
       this.allData = itemsBuilder.build();
       this.shouldPersist = shouldPersist;
+      setEnvironmentId(environmentId);
       setSelector(selector);
     }
   }
@@ -214,6 +228,7 @@ private void applyFullPayload(Iterable<Map.Entry<DataKind, KeyedItems<ItemDescri
       this.allData = newItems;
       this.initialized = true;
       this.shouldPersist = shouldPersist;
+      setEnvironmentId(environmentId);
       setSelector(selector);
     }
   }
@@ -231,7 +246,7 @@ public FullDataSet<ItemDescriptor> exportAll() {
       }
 
       // Preserve the shouldPersist value that was set when data was provided to this store
-      return new FullDataSet<>(builder.build(), this.shouldPersist);
+      return new FullDataSet<>(builder.build(), this.shouldPersist, this.environmentId);
     }
   }
 }

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/LDClient.java

@@ -233,8 +233,10 @@ public LDClient(String sdkKey, LDConfig config) {
       this.evaluator = evaluator;
       this.migrationEvaluator = new MigrationStageEnforcingEvaluator(evaluator, evaluationLogger);
     } else {
-      this.evaluator = new EvaluatorWithHooks(evaluator, allHooks, this.baseLogger.subLogger(Loggers.HOOKS_LOGGER_NAME));
-      this.migrationEvaluator = new EvaluatorWithHooks(new MigrationStageEnforcingEvaluator(evaluator, evaluationLogger), allHooks, this.baseLogger.subLogger(Loggers.HOOKS_LOGGER_NAME));
+      this.evaluator = new EvaluatorWithHooks(evaluator, allHooks, this.baseLogger.subLogger(Loggers.HOOKS_LOGGER_NAME),
+          this.dataSystem::getEnvironmentId);
+      this.migrationEvaluator = new EvaluatorWithHooks(new MigrationStageEnforcingEvaluator(evaluator, evaluationLogger), allHooks,
+          this.baseLogger.subLogger(Loggers.HOOKS_LOGGER_NAME), this.dataSystem::getEnvironmentId);
     }
 
     // Create FlagTracker using the dataSystem's flag change notifier

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/PersistentDataStoreConverter.java

@@ -40,7 +40,7 @@ static FullDataSet<SerializedItemDescriptor> toSerializedFormat(
     }
 
     // Preserve shouldPersist flag when converting formats
-    return new FullDataSet<>(builder.build(), inMemoryData.shouldPersist());
+    return new FullDataSet<>(builder.build(), inMemoryData.shouldPersist(), inMemoryData.getEnvironmentId());
   }
 
   /**

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/PersistentDataStoreWrapper.java

@@ -63,6 +63,7 @@ final class PersistentDataStoreWrapper implements DataStore, SettableCache, Disa
   // themselves remain alive until GC reclaims them; the LoadingCache loaders
   // are short-circuited because every touch site checks this flag first.
   private volatile boolean cacheDisabled;
+  private volatile String environmentId = null;
 
   PersistentDataStoreWrapper(
       final PersistentDataStore core,
@@ -204,7 +205,8 @@ public void init(FullDataSet<ItemDescriptor> allData) {
       KeyedItems<SerializedItemDescriptor> items = PersistentDataStoreConverter.serializeAll(kind, e0.getValue());
       allBuilder.add(new AbstractMap.SimpleEntry<>(kind, items));
     }
-    RuntimeException failure = initCore(new FullDataSet<>(allBuilder.build(), allData.shouldPersist()));
+    RuntimeException failure = initCore(new FullDataSet<>(allBuilder.build(), allData.shouldPersist(),
+        allData.getEnvironmentId()));
     if (itemCache != null && allCache != null && !cacheDisabled) {
       itemCache.invalidateAll();
       allCache.invalidateAll();
@@ -226,6 +228,10 @@ public void init(FullDataSet<ItemDescriptor> allData) {
     }
     if (failure == null || cacheIndefinitely) {
       inited.set(true);
+      String newEnvironmentId = allData.getEnvironmentId();
+      if (newEnvironmentId != null && !newEnvironmentId.isEmpty()) {
+        environmentId = newEnvironmentId;
+      }
     }
     if (failure != null) {
       throw failure;
@@ -372,6 +378,11 @@ public CacheStats getCacheStats() {
         itemStats.evictionCount() + allStats.evictionCount());
   }
 
+  @Override
+  public String getEnvironmentId() {
+    return environmentId;
+  }
+
   private ItemDescriptor getAndDeserializeItem(DataKind kind, String key) {
     SerializedItemDescriptor maybeSerializedItem = core.get(kind, key);
     return maybeSerializedItem == null ? null : PersistentDataStoreConverter.deserialize(kind, maybeSerializedItem);

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/StreamProcessor.java

@@ -31,6 +31,7 @@
 import com.launchdarkly.sdk.server.interfaces.DataStoreStatusProvider;
 import com.launchdarkly.sdk.server.subsystems.DataSource;
 import com.launchdarkly.sdk.server.subsystems.DataSourceUpdateSink;
+import com.launchdarkly.sdk.server.subsystems.DataStoreTypes.FullDataSet;
 import com.launchdarkly.sdk.server.subsystems.DataStoreTypes.ItemDescriptor;
 import com.launchdarkly.sdk.server.subsystems.SerializationException;
 
@@ -269,7 +270,7 @@ private void handleMessage(MessageEvent event, CompletableFuture<Void> initFutur
     try {
       switch (event.getEventName()) {
         case PUT:
-          handlePut(event.getDataReader(), initFuture);
+          handlePut(event.getDataReader(), environmentIdOf(event), initFuture);
           break;
 
         case PATCH:
@@ -331,12 +332,19 @@ private static boolean exceptionHasCause(Throwable e, Class<?> c) {
     return e.getCause() != null && exceptionHasCause(e.getCause(), c);
   }
 
-  private void handlePut(Reader eventData, CompletableFuture<Void> initFuture)
+  private static String environmentIdOf(MessageEvent event) {
+    return event.getHeaders() == null ? null
+        : event.getHeaders().value(HeaderConstants.ENVIRONMENT_ID.getHeaderName());
+  }
+
+  private void handlePut(Reader eventData, String environmentId, CompletableFuture<Void> initFuture)
       throws StreamInputException, StreamStoreException {
     recordStreamInit(false);
     esStarted = 0;
     PutData putData = parseStreamJson(StreamProcessorEvents::parsePutData, eventData);
-    if (!dataSourceUpdates.init(putData.data)) {
+    FullDataSet<ItemDescriptor> data = new FullDataSet<>(putData.data.getData(), putData.data.shouldPersist(),
+        environmentId);
+    if (!dataSourceUpdates.init(data)) {
       throw new StreamStoreException();
     }
     if (!initialized.getAndSet(true)) {

lib/sdk/server/src/main/java/com/launchdarkly/sdk/server/WriteThroughStore.java

@@ -131,6 +131,11 @@ public Selector getSelector() {
     return txMemoryStore.getSelector();
   }
 
+  @Override
+  public String getEnvironmentId() {
+    return memoryStore.getEnvironmentId();
+  }
+
   @Override
   public void close() throws IOException {
     memoryStore.close();
@@ -182,7 +187,8 @@ private boolean applyToLegacyPersistence(ChangeSet<Iterable<Map.Entry<DataKind,
    */
   private void applyFullChangeSetToLegacyStore(ChangeSet<Iterable<Map.Entry<DataKind, KeyedItems<ItemDescriptor>>>> sortedChangeSet) {
     // Preserve shouldPersist flag when converting ChangeSet to FullDataSet
-    persistentStore.init(new FullDataSet<>(sortedChangeSet.getData(), sortedChangeSet.shouldPersist()));
+    persistentStore.init(new FullDataSet<>(sortedChangeSet.getData(), sortedChangeSet.shouldPersist(),
+        sortedChangeSet.getEnvironmentId()));
   }
 
   /**