# Copyright (c) 2021 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
# frozen_string_literal: true

require 'objspace'
require 'singleton'
require 'contrast/components/heap_dump'
require 'contrast/components/logger'

module Contrast
  module Utils
    # Implementation of a heap dump util to automate generation
    class HeapDumpUtil < Contrast::Agent::WorkerThread
      extend Contrast::Components::Logger::InstanceMethods
      include Contrast::Components::Logger::InstanceMethods
      extend Contrast::Components::HeapDump::InstanceMethods

      LOG_ERROR_DUMPS = 'Unable to generate heap dumps'
      FILE_WRITE_FLAGS = 'w'

      class << self
        def enabled?
          heap_dump_enabled?
        end

        def control
          heap_dump_control
        end
      end

      def start_thread!
        return unless Contrast::Utils::HeapDumpUtil.enabled?

        control = Contrast::Utils::HeapDumpUtil.control
        log_enabled_warning
        dir = control[:path]
        Dir.mkdir(dir) unless Dir.exist?(dir)
        return unless File.writable?(dir)

        delay = control[:delay]
        @_thread = Contrast::Agent::Thread.new do
          logger.info("HEAP DUMP THREAD INITIALIZED. WAITING #{ delay } SECONDS TO BEGIN.")
          sleep(delay)
          capture_heap_dump
        end
      rescue StandardError => e
        logger.info(LOG_ERROR_DUMPS, e)
        nil
      end

      def log_enabled_warning
        control = Contrast::Utils::HeapDumpUtil.control
        dir    = control[:path]
        window = control[:window]
        count  = control[:count]
        delay  = control[:delay]
        clean  = control[:clean]

        logger.info <<~WARNING
          *****************************************************
          ********      HEAP DUMP HAS BEEN ENABLED     ********
          *** APPLICATION PROCESS WILL EXIT UPON COMPLETION ***
          *****************************************************

          Heap dump is a debugging tool that snapshots the entire
          state of the Ruby VM. It is an exceptionally expensive
          process, and should only be used to debug especially
          pernicious errors.

          It will write multiple memory snaphots, which are liable
          to be multiple gigabytes in size.
          They will be named "[unix timestamp]-heap.dump",
          e.g.: 1020304050-heap.dump

          It will then call Ruby `exit()`.

          If this is not your specific intent, you can (and should)
          disable this option in your Contrast config file.

          HEAP DUMP PARAMETERS:
          \t[write files to this directory]             dir:     #{ dir    }
          \t[wait this many seconds in between dumps]   window:  #{ window }
          \t[heap dump this many times]                 count:   #{ count  }
          \t[wait this many seconds into app lifetime]  delay:   #{ delay  }
          \t[perform gc pass before dump]               clean:   #{ clean  }

          *****************************************************
          ********        YOU HAVE BEEN WARNED         ********
          *****************************************************
        WARNING
      end

      def capture_heap_dump
        control = Contrast::Utils::HeapDumpUtil.control
        dir    = control[:path]
        window = control[:window]
        count  = control[:count]
        clean  = control[:clean]
        logger.info('HEAP DUMP MAIN LOOP')
        ObjectSpace.trace_object_allocations_start
        count.times do |i|
          logger.info('STARTING HEAP DUMP PASS', current_pass: i, max: count)
          snapshot_heap(dir, clean)
          logger.info('FINISHING HEAP DUMP PASS', current_pass: i, max: count)
          sleep(window)
        end
      ensure
        ObjectSpace.trace_object_allocations_stop
        logger.info('*****************************************************')
        logger.info('********        HEAP DUMP HAS CONCLUDED      ********')
        logger.info('***     APPLICATION PROCESS WILL EXIT SHORTLY     ***')
        logger.info('*****************************************************')
        exit # rubocop:disable Rails/Exit We weren't kidding!
      end

      def snapshot_heap dir, clean
        output = "#{ Time.now.to_f }-heap.dump"
        output = File.join(dir, output)
        begin
          logger.info('OPENING HEADUMP FILE', dir: dir, file: output)
          file = File.new(output, FILE_WRITE_FLAGS)
          if clean
            logger.info('PERFORMING GARBAGE COLLECTION BEFORE HEAP DUMP')
            GC.start
          end
          ObjectSpace.dump_all(output: file)
        ensure
          file.close
        end
      end
    end
  end
end