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

cs__scoped_require 'objspace'
cs__scoped_require 'contrast/components/interface'

module Contrast
  module Utils
    # Implementation of a heap dump util to automate generation
    class HeapDumpUtil
      include Contrast::Components::Interface
      access_component :heap_dump, :logging

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

      class << self
        def run
          return unless heap_dump_enabled?

          log_enabled_warning
          dir = heap_dump_control[:path]
          Dir.mkdir(dir) unless Dir.exist?(dir)
          return unless File.writable?(dir)

          delay = heap_dump_control[:delay]
          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
          dir    = heap_dump_control[:path]
          window = heap_dump_control[:window]
          count  = heap_dump_control[:count]
          delay  = heap_dump_control[:delay]
          clean  = heap_dump_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
          dir    = heap_dump_control[:path]
          window = heap_dump_control[:window]
          count  = heap_dump_control[:count]
          clean  = heap_dump_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 + 1, max: count)
            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)
              logger.info('FINISHING HEAP DUMP PASS', current_pass: i + 1, max: count)
            ensure
              file.close
            end
            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 # We weren't kidding!
        end
      end
    end
  end
end