module PlainApm
  module Extensions
    module ThreadAllocations
      module ActiveSupportEvent
        def start!
          super
          @thread_allocation_count_start = now_thread_allocations
        end

        def finish!
          super
          @thread_allocation_count_finish = now_thread_allocations
        end

        def thread_allocations
          @thread_allocation_count_finish - @thread_allocation_count_start
        end

        if defined?(PlainApm::ObjectTracing)
          def now_thread_allocations
            PlainApm::ObjectTracing.total_thread_allocated_objects
          end
        else
          def now_thread_allocations
            0
          end
        end
      end
    end
  end
end