feat: Add environment ID support for hooks

.github/actions/check/action.yml

@@ -51,4 +51,4 @@ runs:
         test_service_port: 9000
         enable_persistence_tests: true
         token: ${{ inputs.token }}
-        version: v3.0.0-alpha.6
+        version: v3

contract-tests/hook.rb

@@ -34,6 +34,7 @@ def before_evaluation(evaluation_series_context, data)
         context: @context_filter.filter(evaluation_series_context.context),
         defaultValue: evaluation_series_context.default_value,
         method: evaluation_series_context.method,
+        environmentId: evaluation_series_context.environment_id,
       },
       evaluationSeriesData: data,
       stage: 'beforeEvaluation',
@@ -58,6 +59,7 @@ def after_evaluation(evaluation_series_context, data, detail)
         context: @context_filter.filter(evaluation_series_context.context),
         defaultValue: evaluation_series_context.default_value,
         method: evaluation_series_context.method,
+        environmentId: evaluation_series_context.environment_id,
       },
       evaluationSeriesData: data,
       evaluationDetail: {

contract-tests/service.rb

@@ -46,6 +46,7 @@
       'instance-id',
       'anonymous-redaction',
       'evaluation-hooks',
+      'hook-environment-id',
       'omit-anonymous-contexts',
       'client-prereq-events',
       'persistent-data-store-consul',

lib/ldclient-rb/impl/data_source.rb

@@ -7,6 +7,27 @@
 module LaunchDarkly
   module Impl
     module DataSource
+      LD_ENVID_HEADER = "X-LD-EnvID"
+
+      #
+      # Records the environment ID reported by LaunchDarkly, when the response headers provide one and the sink is
+      # able to hold it. Externally implemented sinks are not required to support this.
+      #
+      # @param sink [LaunchDarkly::Interfaces::DataSource::UpdateSink, nil]
+      # @param headers [#[], nil]
+      # @return [void]
+      #
+      def self.record_environment_id(sink, headers)
+        return if headers.nil? || !sink.respond_to?(:set_environment_id)
+
+        environment_id = headers[LD_ENVID_HEADER]
+        # The http gem returns arrays for repeated headers; normalize to a string.
+        environment_id = environment_id.first if environment_id.is_a?(Array)
+        return unless environment_id.is_a?(String) && !environment_id.empty?
+
+        sink.set_environment_id(environment_id)
+      end
+
       class StatusProvider
         include LaunchDarkly::Interfaces::DataSource::StatusProvider
 
@@ -41,12 +62,28 @@ def initialize(data_store, status_broadcaster, flag_change_broadcaster)
           @dependency_tracker = LaunchDarkly::Impl::DependencyTracker.new
 
           @mutex = Mutex.new
+          @environment_id = nil
           @current_status = LaunchDarkly::Interfaces::DataSource::Status.new(
             LaunchDarkly::Interfaces::DataSource::Status::INITIALIZING,
             Time.now,
             nil)
         end
 
+        #
+        # @return [String, nil] The environment ID reported by LaunchDarkly, if known
+        #
+        def environment_id
+          @mutex.synchronize { @environment_id }
+        end
+
+        #
+        # @param environment_id [String]
+        # @return [void]
+        #
+        def set_environment_id(environment_id)
+          @mutex.synchronize { @environment_id = environment_id }
+        end
+
         def init(all_data)
           old_data = nil
           monitor_store_update do

lib/ldclient-rb/impl/data_source/polling.rb

@@ -1,3 +1,4 @@
+require "ldclient-rb/impl/data_source"
 require "ldclient-rb/impl/repeating_task"
 require "ldclient-rb/impl/util"
 
@@ -35,7 +36,8 @@ def stop
 
         def poll
           begin
-            all_data = @requestor.request_all_data
+            all_data, headers = request_all_data
+            DataSource.record_environment_id(@config.data_source_update_sink, headers)
             if all_data
               update_sink_or_data_store.init(all_data)
               if @initialized.make_true
@@ -92,6 +94,17 @@ def poll
           @config.data_source_update_sink || @config.feature_store
         end
 
+        #
+        # Requestors provided by application code may not be able to report the response headers.
+        #
+        # @return [Array(Hash, HTTP::Headers, nil)]
+        #
+        private def request_all_data
+          return @requestor.request_all_data_with_headers if @requestor.respond_to?(:request_all_data_with_headers)
+
+          [@requestor.request_all_data, nil]
+        end
+
         #
         # @param [LaunchDarkly::Interfaces::DataSource::ErrorInfo, nil] error_info
         #

lib/ldclient-rb/impl/data_source/requestor.rb

@@ -40,8 +40,18 @@ def initialize(sdk_key, config)
         end
 
         def request_all_data()
-          all_data = JSON.parse(make_request("/sdk/latest-all"), symbolize_names: true)
-          Impl::Model.make_all_store_data(all_data, @config.logger)
+          request_all_data_with_headers.first
+        end
+
+        #
+        # Requests the full data set, also returning the headers of the response it was retrieved from.
+        #
+        # @return [Array(Hash, HTTP::Headers)]
+        #
+        def request_all_data_with_headers
+          body, headers = make_request("/sdk/latest-all")
+          all_data = JSON.parse(body, symbolize_names: true)
+          [Impl::Model.make_all_store_data(all_data, @config.logger), headers]
         end
 
         def stop
@@ -80,7 +90,7 @@ def make_request(path)
             etag = response.headers["etag"]
             @cache.write(uri, CacheEntry.new(etag, body)) unless etag.nil?
           end
-          body
+          [body, response.headers]
         end
 
         def fix_encoding(body, content_type)

lib/ldclient-rb/impl/data_source/stream.rb

@@ -1,3 +1,4 @@
+require "ldclient-rb/impl/data_source"
 require "ldclient-rb/impl/model/serialization"
 require "ldclient-rb/impl/util"
 require "ldclient-rb/in_memory_store"
@@ -54,6 +55,7 @@ def start
 
           uri = Impl::Util.add_payload_filter_key(@config.stream_uri + "/all", @config)
           @es = SSE::Client.new(uri, **opts) do |conn|
+            conn.on_connect { |response_headers| DataSource.record_environment_id(@data_source_update_sink, response_headers) }
             conn.on_event { |event| process_message(event) }
             conn.on_error { |err|
               log_connection_result(false)

lib/ldclient-rb/impl/data_system.rb

@@ -119,6 +119,15 @@ def set_diagnostic_accumulator(diagnostic_accumulator)
         raise NotImplementedError, "#{self.class} must implement #set_diagnostic_accumulator"
       end
 
+      #
+      # Returns the ID of the environment the SDK is connected to, if LaunchDarkly has reported one.
+      #
+      # @return [String, nil]
+      #
+      def environment_id
+        raise NotImplementedError, "#{self.class} must implement #environment_id"
+      end
+
       #
       # Represents the availability of data in the SDK.
       #

lib/ldclient-rb/impl/data_system/fdv1.rb

@@ -121,6 +121,11 @@ def flag_change_broadcaster
           @flag_change_broadcaster
         end
 
+        # (see DataSystem#environment_id)
+        def environment_id
+          @data_source_update_sink.environment_id
+        end
+
         #
         # (see DataSystem#data_availability)
         #

lib/ldclient-rb/impl/data_system/fdv2.rb

@@ -100,6 +100,9 @@ def initialize(sdk_key, config, data_system_config)
             @store.with_persistence(wrapper, writable, @data_store_status_provider)
           end
 
+          @environment_id = nil
+          @environment_id_lock = Mutex.new
+
           # Threading
           @stop_event = Concurrent::Event.new
           @ready_event = Concurrent::Event.new
@@ -185,6 +188,11 @@ def flag_change_broadcaster
           @flag_change_broadcaster
         end
 
+        # (see DataSystem#environment_id)
+        def environment_id
+          @environment_id_lock.synchronize { @environment_id }
+        end
+
         # (see DataSystem#data_availability)
         def data_availability
           return DataAvailability::REFRESHED if @store.selector.defined?
@@ -270,6 +278,7 @@ def run_initializers
               if basis_result.success?
                 basis = basis_result.value
                 @logger.info { "[LDClient] Initialized via #{initializer.name}" }
+                record_environment_id(basis.environment_id)
 
                 # Apply the basis to the store regardless of whether fallback was signalled.
                 # If the server returned a valid payload alongside the directive we still want
@@ -460,7 +469,10 @@ def consume_synchronizer_results(synchronizer, check_recovery: false)
               @store.apply(update.change_set, true) if update.change_set
 
               # Set ready event on valid update
-              @ready_event.set if update.state == LaunchDarkly::Interfaces::DataSource::Status::VALID
+              if update.state == LaunchDarkly::Interfaces::DataSource::Status::VALID
+                @ready_event.set
+                record_environment_id(update.environment_id)
+              end
 
               # Update status
               @data_source_status_provider.update_status(update.state, update.error)
@@ -481,6 +493,18 @@ def consume_synchronizer_results(synchronizer, check_recovery: false)
           SyncResult::REMOVE
         end
 
+        #
+        # Retains the environment ID reported by a successful connection to LaunchDarkly.
+        #
+        # @param environment_id [String, nil]
+        # @return [void]
+        #
+        def record_environment_id(environment_id)
+          return unless environment_id.is_a?(String) && !environment_id.empty?
+
+          @environment_id_lock.synchronize { @environment_id = environment_id }
+        end
+
         #
         # Determine if we should fallback to the next synchronizer.
         #

lib/ldclient-rb/interfaces/hooks.rb

@@ -70,17 +70,27 @@ class EvaluationSeriesContext
         attr_reader :default_value
         attr_reader :method
 
+        #
+        # The ID of the LaunchDarkly environment the SDK is connected to, if known. This is only available once the
+        # SDK has received data from LaunchDarkly.
+        #
+        # @return [String, nil]
+        #
+        attr_reader :environment_id
+
         #
         # @param key [String]
         # @param context [LaunchDarkly::LDContext]
         # @param default_value [any]
         # @param method [Symbol]
+        # @param environment_id [String, nil]
         #
-        def initialize(key, context, default_value, method)
+        def initialize(key, context, default_value, method, environment_id = nil)
           @key = key
           @context = context
           @default_value = default_value
           @method = method
+          @environment_id = environment_id
         end
       end
     end

lib/ldclient-rb/ldclient.rb

@@ -421,7 +421,7 @@ def variation_detail(key, context, default)
       # Hooks can be added and we want to ensure all correct stages for a given hook execute. For example, we do not
       # want to trigger the after_evaluation method without also triggering the before_evaluation method.
       hooks = @hooks.dup
-      evaluation_series_context = Interfaces::Hooks::EvaluationSeriesContext.new(key, context, default, method)
+      evaluation_series_context = Interfaces::Hooks::EvaluationSeriesContext.new(key, context, default, method, @data_system.environment_id)
 
       [hooks, evaluation_series_context]
     end

spec/impl/data_source/polling_spec.rb

@@ -84,6 +84,43 @@ def with_processor(store, initialize_to_valid = false)
       end
     end
 
+    describe 'environment ID' do
+      flag = Impl::Model::FeatureFlag.new({ key: 'flagkey', version: 1 })
+      all_data = {
+        Impl::DataStore::FEATURES => { flagkey: flag },
+        Impl::DataStore::SEGMENTS => {},
+      }
+
+      it 'is recorded from the response headers' do
+        allow(requestor).to receive(:request_all_data_with_headers).and_return([all_data, { "X-LD-EnvID" => "env-abc" }])
+        store = InMemoryFeatureStore.new
+        with_processor(store) do |processor|
+          config = processor.instance_variable_get(:@config)
+          processor.start.wait
+          expect(config.data_source_update_sink.environment_id).to eq("env-abc")
+        end
+      end
+
+      it 'is not recorded when the header is absent' do
+        allow(requestor).to receive(:request_all_data_with_headers).and_return([all_data, {}])
+        store = InMemoryFeatureStore.new
+        with_processor(store) do |processor|
+          config = processor.instance_variable_get(:@config)
+          processor.start.wait
+          expect(config.data_source_update_sink.environment_id).to be_nil
+        end
+      end
+
+      it 'is not recorded for an error response' do
+        allow(requestor).to receive(:request_all_data_with_headers).and_raise(Impl::DataSource::UnexpectedResponseError.new(503))
+        with_processor(InMemoryFeatureStore.new, true) do |processor|
+          config = processor.instance_variable_get(:@config)
+          processor.start.wait(1)
+          expect(config.data_source_update_sink.environment_id).to be_nil
+        end
+      end
+    end
+
     describe 'connection error' do
       it 'does not cause immediate failure, does not set initialized' do
         allow(requestor).to receive(:request_all_data).and_raise(StandardError.new("test error"))

spec/impl/data_source/requestor_spec.rb

@@ -46,6 +46,18 @@ def with_requestor(base_uri, opts = {})
         end
       end
 
+      it "returns the response headers" do
+        expected_data = DataSetBuilder.new.flag(FlagBuilder.new("x").build)
+        with_server do |server|
+          with_requestor(server.base_uri.to_s) do |requestor|
+            server.setup_ok_response("/", expected_data.to_json, "application/json", { "X-LD-EnvID" => "env-abc" })
+            data, headers = requestor.request_all_data_with_headers
+            expect(data).to eq expected_data.to_store_data
+            expect(headers["X-LD-EnvID"]).to eq "env-abc"
+          end
+        end
+      end
+
       it "logs debug output" do
         logger = ::Logger.new($stdout)
         logger.level = ::Logger::DEBUG

spec/impl/data_source/stream_spec.rb

@@ -68,6 +68,39 @@ module LaunchDarkly
       end
     end
 
+    describe 'environment ID' do
+      def with_connect_handler(processor)
+        connection = double("connection", on_event: nil, on_error: nil)
+        handler = nil
+        allow(connection).to receive(:on_connect) { |&block| handler = block }
+        allow(SSE::Client).to receive(:new) do |_uri, **_opts, &block|
+          block.call(connection)
+          double("SSE::Client", close: nil)
+        end
+
+        processor.start
+        begin
+          yield handler
+        ensure
+          processor.stop
+        end
+      end
+
+      it 'is recorded from the connection response headers' do
+        with_connect_handler(processor) do |handler|
+          handler.call({ "X-LD-EnvID" => "env-abc" })
+          expect(config.data_source_update_sink.environment_id).to eq("env-abc")
+        end
+      end
+
+      it 'is not recorded when the header is absent' do
+        with_connect_handler(processor) do |handler|
+          handler.call({})
+          expect(config.data_source_update_sink.environment_id).to be_nil
+        end
+      end
+    end
+
     describe '#log_connection_result' do
       it "logs successful connection when diagnostic_accumulator is provided" do
         diagnostic_accumulator = double("DiagnosticAccumulator")