# Copyright (c) 2021 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/service_heartbeat' require 'contrast/api/communication/messaging_queue' module Contrast module Agent # This class used to ensure that our worker threads are running in multi-process environments # # @attr_reader heapdump_util [Contrast::Utils::HeapDumpUtil] # @attr_reader heartbeat [Contrast::Agent::ServiceHeartbeat] # @attr_reader messaging_queue [Contrast::Api::Communication::MessagingQueue] class ThreadWatcher include Contrast::Components::Logger::InstanceMethods attr_reader :heapdump_util, :heartbeat, :messaging_queue def initialize @pids = {} @heapdump_util = Contrast::Utils::HeapDumpUtil.new @heartbeat = Contrast::Agent::ServiceHeartbeat.new @messaging_queue = Contrast::Api::Communication::MessagingQueue.new end def startup! return unless ::Contrast::AGENT.enabled? unless heartbeat.running? logger.debug('Attempting to start heartbeat thread') heartbeat.start_thread! end heartbeat_result = heartbeat.running? logger.debug('Heartbeat thread status', alive: heartbeat_result) unless messaging_queue.running? logger.debug('Attempting to start messaging queue thread') messaging_queue.start_thread! end messaging_result = messaging_queue.running? logger.debug('Messaging thread status', alive: messaging_result) logger.debug('ThreadWatcher started threads') @pids[Process.pid] = messaging_result && heartbeat_result end def ensure_running? return if @pids[Process.pid] == true logger.debug('ThreadWatcher - threads not running') startup! end def shutdown! heartbeat.stop! messaging_queue.stop! heapdump_util.stop! end end end end