fix(schema-polling): improve polling

packages/forest_admin_datasource_rpc/lib/forest_admin_datasource_rpc.rb

@@ -13,12 +13,12 @@ class Error < StandardError; end
   # @param options [Hash] Configuration options
   # @option options [String] :uri The URI of the RPC agent
   # @option options [String] :auth_secret The authentication secret (optional, will use cache if not provided)
-  # @option options [Integer] :schema_polling_interval Polling interval in seconds (optional)
+  # @option options [Integer] :schema_polling_interval_sec Polling interval in seconds (optional)
   #   - Default: 600 seconds (10 minutes)
-  #   - Can be overridden with ENV['SCHEMA_POLLING_INTERVAL']
+  #   - Can be overridden with ENV['SCHEMA_POLLING_INTERVAL_SEC']
   #   - Valid range: 1-3600 seconds
-  #   - Priority: options[:schema_polling_interval] > ENV['SCHEMA_POLLING_INTERVAL'] > default
-  #   - Example: SCHEMA_POLLING_INTERVAL=30 for development (30 seconds)
+  #   - Priority: options[:schema_polling_interval_sec] > ENV['SCHEMA_POLLING_INTERVAL_SEC'] > default
+  #   - Example: SCHEMA_POLLING_INTERVAL_SEC=30 for development (30 seconds)
   # @option options [Hash] :introspection Pre-defined schema introspection for resilient deployment
   #   - When provided, allows the datasource to start even if the RPC slave is unreachable
   #   - The introspection will be used as fallback when the slave connection fails
@@ -29,36 +29,35 @@ def self.build(options)
     uri = options[:uri]
     auth_secret = options[:auth_secret] || ForestAdminAgent::Facades::Container.cache(:auth_secret)
     provided_introspection = options[:introspection]
-    ForestAdminAgent::Facades::Container.logger.log('Info', "Getting schema from RPC agent on #{uri}.")
 
-    schema = nil
+    # Create schema polling client with configurable polling interval
+    # Priority: options[:schema_polling_interval_sec] > ENV['SCHEMA_POLLING_INTERVAL_SEC'] > default (600)
+    polling_interval = if options[:schema_polling_interval_sec]
+                         options[:schema_polling_interval_sec]
+                       elsif ENV['SCHEMA_POLLING_INTERVAL_SEC']
+                         ENV['SCHEMA_POLLING_INTERVAL_SEC'].to_i
+                       else
+                         600 # 10 minutes by default
+                       end
 
-    begin
-      rpc_client = Utils::RpcClient.new(uri, auth_secret)
-      response = rpc_client.fetch_schema('/forest/rpc-schema')
-      schema = response.body
-    rescue Faraday::ConnectionFailed => e
-      ForestAdminAgent::Facades::Container.logger.log(
-        'Error',
-        "Connection failed to RPC agent at #{uri}: #{e.message}\n#{e.backtrace.join("\n")}"
-      )
-    rescue Faraday::TimeoutError => e
-      ForestAdminAgent::Facades::Container.logger.log(
-        'Error',
-        "Request timeout to RPC agent at #{uri}: #{e.message}"
-      )
-    rescue ForestAdminAgent::Http::Exceptions::AuthenticationOpenIdClient => e
-      ForestAdminAgent::Facades::Container.logger.log(
-        'Error',
-        "Authentication failed with RPC agent at #{uri}: #{e.message}"
-      )
-    rescue StandardError => e
-      ForestAdminAgent::Facades::Container.logger.log(
-        'Error',
-        "Failed to get schema from RPC agent at #{uri}: #{e.class} - #{e.message}\n#{e.backtrace.join("\n")}"
-      )
+    polling_options = {
+      polling_interval: polling_interval
+    }
+
+    schema_polling = Utils::SchemaPollingClient.new(uri, auth_secret, polling_options, provided_introspection) do
+      # Callback when schema change is detected
+      logger = ForestAdminAgent::Facades::Container.logger
+      logger.log('Info', '[RPCDatasource] Schema change detected, reloading agent...')
+      ForestAdminAgent::Builder::AgentFactory.instance.reload!
     end
 
+    # Start polling (includes initial synchronous schema fetch)
+    # The initial fetch is blocking, then async polling starts
+    schema_polling.start
+
+    # Get the schema from the polling client
+    schema = schema_polling.current_schema
+
     # Use provided introspection as fallback when slave is unreachable
     if schema.nil? && provided_introspection
       ForestAdminAgent::Facades::Container.logger.log(
@@ -73,66 +72,7 @@ def self.build(options)
       # return empty datasource for not breaking stack
       ForestAdminDatasourceToolkit::Datasource.new
     else
-      # Create schema polling client with configurable polling interval
-      # Priority: options[:schema_polling_interval] > ENV['SCHEMA_POLLING_INTERVAL'] > default (600)
-      polling_interval = if options[:schema_polling_interval]
-                           options[:schema_polling_interval]
-                         elsif ENV['SCHEMA_POLLING_INTERVAL']
-                           ENV['SCHEMA_POLLING_INTERVAL'].to_i
-                         else
-                           600 # 10 minutes by default
-                         end
-
-      polling_options = {
-        polling_interval: polling_interval
-      }
-
-      schema_polling = Utils::SchemaPollingClient.new(uri, auth_secret, polling_options) do
-        # Callback when schema change is detected
-        logger = ForestAdminAgent::Facades::Container.logger
-        logger.log('Info', '[RPCDatasource] Schema change detected, reloading agent...')
-        ForestAdminAgent::Builder::AgentFactory.instance.reload!
-      end
-      schema_polling.start
-
-      datasource = ForestAdminDatasourceRpc::Datasource.new(options, schema, schema_polling)
-
-      # Setup cleanup hooks for proper schema polling client shutdown
-      setup_cleanup_hooks(datasource)
-
-      datasource
-    end
-  end
-
-  def self.setup_cleanup_hooks(datasource)
-    # Register cleanup handler for graceful shutdown
-    at_exit do
-      datasource.cleanup
-    rescue StandardError => e
-      # Silently ignore errors during exit cleanup to prevent test pollution
-      warn "[RPCDatasource] Error during at_exit cleanup: #{e.message}" if $VERBOSE
-    end
-
-    # Handle SIGINT (Ctrl+C)
-    Signal.trap('INT') do
-      begin
-        ForestAdminAgent::Facades::Container.logger&.log('Info', '[RPCDatasource] Received SIGINT, cleaning up...')
-      rescue StandardError
-        # Logger might not be available
-      end
-      datasource.cleanup
-      exit(0)
-    end
-
-    # Handle SIGTERM (default kill signal)
-    Signal.trap('TERM') do
-      begin
-        ForestAdminAgent::Facades::Container.logger&.log('Info', '[RPCDatasource] Received SIGTERM, cleaning up...')
-      rescue StandardError
-        # Logger might not be available
-      end
-      datasource.cleanup
-      exit(0)
+      ForestAdminDatasourceRpc::Datasource.new(options, schema, schema_polling)
     end
   end
 end

packages/forest_admin_datasource_rpc/lib/forest_admin_datasource_rpc/Utils/schema_polling_client.rb

@@ -1,23 +1,25 @@
 require 'openssl'
 require 'json'
 require 'time'
+require 'digest'
 
 module ForestAdminDatasourceRpc
   module Utils
     class SchemaPollingClient
-      attr_reader :closed
+      attr_reader :closed, :current_schema
 
       DEFAULT_POLLING_INTERVAL = 600 # seconds (10 minutes)
       MIN_POLLING_INTERVAL = 1 # seconds (minimum safe interval)
       MAX_POLLING_INTERVAL = 3600 # seconds (1 hour max)
 
-      def initialize(uri, auth_secret, options = {}, &on_schema_change)
+      def initialize(uri, auth_secret, options = {}, previous_schema = nil, &on_schema_change)
         @uri = uri
         @auth_secret = auth_secret
         @polling_interval = options[:polling_interval] || DEFAULT_POLLING_INTERVAL
         @on_schema_change = on_schema_change
         @closed = false
-        @cached_etag = nil
+        @current_schema = previous_schema
+        @cached_etag = compute_etag(previous_schema) if previous_schema
         @polling_thread = nil
         @mutex = Mutex.new
         @connection_attempts = 0
@@ -29,12 +31,22 @@ def initialize(uri, auth_secret, options = {}, &on_schema_change)
         @rpc_client = RpcClient.new(@uri, @auth_secret)
       end
 
-      def start
-        return if @closed
+      # Start schema polling with initial synchronous fetch
+      # The first schema fetch is done synchronously (blocking) to ensure
+      # the schema is available immediately. Then async polling starts.
+      #
+      # @return [Boolean] true if started successfully, false if already running or closed
+      def start # rubocop:disable Naming/PredicateMethod
+        return false if @closed
 
         @mutex.synchronize do
-          return if @polling_thread&.alive?
+          return false if @polling_thread&.alive?
+
+          # Fetch initial schema synchronously before starting the polling thread
+          ForestAdminAgent::Facades::Container.logger&.log('Info', "Getting schema from RPC agent on #{@uri}.")
+          fetch_initial_schema_sync
 
+          # Start async polling thread
           @polling_thread = Thread.new do
             polling_loop
           rescue StandardError => e
@@ -49,6 +61,7 @@ def start
           'Info',
           "[Schema Polling] Polling started (interval: #{@polling_interval}s)"
         )
+        true
       end
 
       def stop
@@ -69,31 +82,85 @@ def stop
 
       private
 
+      # Compute ETag from schema using same algorithm as RPC slave
+      # @param schema [Hash] The schema to hash
+      # @return [String, nil] SHA1 hexdigest of schema JSON, or nil if schema is nil
+      def compute_etag(schema)
+        return nil if schema.nil?
+
+        Digest::SHA1.hexdigest(schema.to_json)
+      end
+
+      # Fetch initial schema synchronously (called from start() in main thread)
+      # This is a blocking call that sets @current_schema and @cached_etag
+      def fetch_initial_schema_sync
+        result = @rpc_client.fetch_schema('/forest/rpc-schema')
+        @current_schema = result.body
+        @cached_etag = result.etag || compute_etag(@current_schema)
+        ForestAdminAgent::Facades::Container.logger&.log(
+          'Debug',
+          "[Schema Polling] Initial schema fetched successfully (ETag: #{@cached_etag})"
+        )
+      rescue Faraday::ConnectionFailed => e
+        ForestAdminAgent::Facades::Container.logger&.log(
+          'Error',
+          "Connection failed to RPC agent at #{@uri}: #{e.message}\n#{e.backtrace.join("\n")}"
+        )
+      rescue Faraday::TimeoutError => e
+        ForestAdminAgent::Facades::Container.logger&.log(
+          'Error',
+          "Request timeout to RPC agent at #{@uri}: #{e.message}"
+        )
+      rescue ForestAdminAgent::Http::Exceptions::AuthenticationOpenIdClient => e
+        ForestAdminAgent::Facades::Container.logger&.log(
+          'Error',
+          "Authentication failed with RPC agent at #{@uri}: #{e.message}"
+        )
+      rescue StandardError => e
+        ForestAdminAgent::Facades::Container.logger&.log(
+          'Error',
+          "Failed to get schema from RPC agent at #{@uri}: #{e.class} - #{e.message}\n#{e.backtrace.join("\n")}"
+        )
+      end
+
       def polling_loop
+        etag_status = @cached_etag ? "with ETag: #{@cached_etag}" : 'without initial ETag'
         ForestAdminAgent::Facades::Container.logger&.log(
           'Debug',
-          "[Schema Polling] Starting polling loop (interval: #{@polling_interval}s)"
+          "[Schema Polling] Starting polling loop (interval: #{@polling_interval}s, #{etag_status})"
         )
 
+        first_check = true
+
         loop do
           break if @closed
 
+          # Wait before checking (skip wait on first iteration to start polling immediately)
+          unless first_check
+            etag_info = @cached_etag || 'none'
+            ForestAdminAgent::Facades::Container.logger&.log(
+              'Debug',
+              "[Schema Polling] Waiting #{@polling_interval}s before next check (current ETag: #{etag_info})"
+            )
+            sleep_with_interrupt(@polling_interval)
+            break if @closed
+          end
+          first_check = false
+
+          # Check for schema changes
           begin
             check_schema
           rescue StandardError => e
             handle_error(e)
           end
+        end
+      end
 
-          # Sleep with interrupt check (check every second for early termination)
-          ForestAdminAgent::Facades::Container.logger&.log(
-            'Debug',
-            "[Schema Polling] Waiting #{@polling_interval}s before next check (current ETag: #{@cached_etag || "none"})"
-          )
-          remaining = @polling_interval
-          while remaining.positive? && !@closed
-            sleep([remaining, 1].min)
-            remaining -= 1
-          end
+      def sleep_with_interrupt(duration)
+        remaining = duration
+        while remaining.positive? && !@closed
+          sleep([remaining, 1].min)
+          remaining -= 1
         end
       end
 
@@ -162,22 +229,28 @@ def handle_schema_unchanged
       end
 
       def handle_schema_changed(result)
-        new_etag = result.etag
-        @cached_etag.nil? ? handle_initial_schema(new_etag) : handle_schema_update(result.body, new_etag)
-        @connection_attempts = 0
-      end
+        new_schema = result.body
+        new_etag = result.etag || compute_etag(new_schema)
 
-      def handle_initial_schema(etag)
-        @cached_etag = etag
-        ForestAdminAgent::Facades::Container.logger&.log(
-          'Debug',
-          "[Schema Polling] Initial schema loaded successfully (ETag: #{etag})"
-        )
+        if @cached_etag.nil?
+          # Initial schema fetch
+          @cached_etag = new_etag
+          @current_schema = new_schema
+          ForestAdminAgent::Facades::Container.logger&.log(
+            'Debug',
+            "[Schema Polling] Initial schema loaded successfully (ETag: #{new_etag})"
+          )
+        else
+          # Schema update detected
+          handle_schema_update(new_schema, new_etag)
+        end
+        @connection_attempts = 0
       end
 
       def handle_schema_update(schema, etag)
         old_etag = @cached_etag
         @cached_etag = etag
+        @current_schema = schema
         msg = "[Schema Polling] Schema changed detected (old ETag: #{old_etag}, new ETag: #{etag}), " \
               'triggering reload callback'
         ForestAdminAgent::Facades::Container.logger&.log('Info', msg)