lib/contrast/agent/thread_watcher.rb in contrast-agent-6.11.0 vs lib/contrast/agent/thread_watcher.rb in contrast-agent-6.12.0

@@ -1,6 +1,6 @@
-# Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
+# Copyright (c) 2023 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
 # frozen_string_literal: true
 
 require 'contrast/components/logger'
 require 'contrast/agent/reporting/report'
 require 'contrast/agent/reporting/reporting_workers/reporting_workers'
@@ -26,11 +26,10 @@
       attr_reader :reporter_settings_worker
       # @return [Contrast::Agent::ReportingWorkers::ApplicationServerWorker]
       attr_reader :reporter_app_settings_worker
 
       def initialize
-        @pids = {}
         @heapdump_util = Contrast::Utils::HeapDumpUtil.new
         @reporter = Contrast::Agent::Reporter.new
         @reporter_heartbeat = Contrast::Agent::ReportingWorkers::ReporterHeartbeat.new
         @reporter_settings_worker = Contrast::Agent::ReportingWorkers::ServerSettingsWorker.new
         @reporter_app_settings_worker = Contrast::Agent::ReportingWorkers::ApplicationServerWorker.new
@@ -38,52 +37,28 @@
         @worth_watching_analyzer = Contrast::Agent::Protect::WorthWatchingInputAnalyzer.new unless protect_disabled?
       end
 
       # @return [Hash, nil] map of process to thread startup status
       def startup!
+        check_before_start
+
         return unless ::Contrast::AGENT.enabled?
 
         logger.debug('ThreadWatcher started threads')
 
-        @pids[Process.pid] = reporter_threads
-        if Contrast::Agent::Telemetry::Base.enabled?
-          telemetry_status = init_thread(telemetry_queue)
-          @pids[Process.pid] = @pids[Process.pid] && telemetry_status
+        desired_threads.each do |thread|
+          init_thread(thread)
         end
-        unless protect_disabled?
-          worth_watching_analyzer_status = init_thread(worth_watching_analyzer)
-          @pids[Process.pid] = @pids[Process.pid] && worth_watching_analyzer_status
-        end
-        @pids
       end
 
-      def reporter_threads
-        init_thread(reporter) && init_thread(reporter_heartbeat) &&
-            init_thread(reporter_settings_worker) && init_thread(reporter_app_settings_worker)
-      end
-
       def ensure_running?
-        return if @pids[Process.pid] == true
+        return if desired_threads.all?(&:running?)
 
         logger.debug('ThreadWatcher - threads not running')
         startup!
       end
 
-      # This method will check the config and if the config is invalid - it will kill the agent
-      def self.check_before_start
-        return if Contrast::CONFIG.send(:validate)
-
-        @_diagnostics = Contrast::Agent::DiagnosticsConfig::Diagnostics.new(Contrast::AGENT.logger.path ||
-                                                                            Contrast::Components::Config::CONTRAST_LOG)
-        @_diagnostics.config.populate_fail_connection
-        @_diagnostics.write_to_file_logic(false, reset: false)
-        # kill agent
-        Contrast::AGENT.disable_agent!
-        # TODO: RUBY-1822
-        # set the in_application_scope to 1 so we short circuit it
-      end
-
       def shutdown!
         heapdump_util&.stop!
         telemetry_queue&.stop!
         reporter&.stop!
         reporter_heartbeat&.stop!
@@ -97,10 +72,20 @@
         @telemetry
       end
 
       private
 
+      # Retrieve a list of the desired threads.
+      #
+      # @return [Array<Contrast::Agent::WorkerThread>]
+      def desired_threads
+        [reporter, reporter_heartbeat, reporter_settings_worker, reporter_app_settings_worker].tap do |entries|
+          entries.push(telemetry_queue) if Contrast::Agent::Telemetry::Base.enabled?
+          entries.push(worth_watching_analyzer) unless protect_disabled?
+        end
+      end
+
       # Start the thread governed by the given watcher
       #
       # @param watcher [Contrast::Agent::ThreadWatcher]
       # @return [Boolean] if the watched thread started successfully
       def init_thread watcher
@@ -115,9 +100,21 @@
         result
       end
 
       def protect_disabled?
         ::Contrast::PROTECT.forcibly_disabled?
+      end
+
+      # This method will check the config and if the config is invalid - it will kill the agent
+      def check_before_start
+        return if Contrast::CONFIG.send(:validate)
+
+        @_diagnostics = Contrast::Agent::DiagnosticsConfig::Diagnostics.new(Contrast::AGENT.logger.path ||
+                                                                            Contrast::Components::Config::CONTRAST_LOG)
+        @_diagnostics.config.populate_fail_connection
+        @_diagnostics.write_to_file_logic(false, reset: false)
+        # kill agent
+        Contrast::AGENT.disable_agent!
       end
     end
   end
 end