#
# File:: memory_profiler.rb
# Description:: 
#
# Author:: David Muir <david.muir@rockstarnorth.com>
# Date:: 9 April 2009
#

#----------------------------------------------------------------------------
# Uses
#----------------------------------------------------------------------------
# None

#----------------------------------------------------------------------------
# Implementation
#----------------------------------------------------------------------------
module Pipeline
module Util

	#
	# == Description
	# 
	# === Example Usage
	#
	# === References
	# The original which this profiler code is based can be found at:
	# http://scottstuff.net/blog/articles/2006/08/17/memory-leak-profiling-with-rails
	#
	class MemoryProfiler
		DEFAULTS = { 
			:delay => 10,
			:string_debug => true,
			:log_file => 'x:\memory_profiler.log',
			:string_log_file => 'x:\memory_profiler_strings.log'
		}

		UNIT_BYTES = 1
		UNIT_KILOBYTES = 1024
		UNIT_MEGABYTES = (1024*1024)

		def MemoryProfiler::stop( )
			@@running = false
			@@profiler_thread.join( )
		end

		def MemoryProfiler::current( )
			@@current
		end

		def MemoryProfiler::previous( )
			@@previous
		end

		def MemoryProfiler::display( sort_key = :size, max_num = 20, io = STDOUT, units = UNIT_KILOBYTES )
			begin
				io.puts "#{Time.now.strftime( '%Y/%m/%d %H:%M:%S' )}"
				io.puts "Current Top #{max_num}"
				io.puts 
				io.puts "sortkey #{sort_key.class} #{sort_key.to_s}"
				current().sort_by { |k,v| -v[sort_key].abs }[0..max_num-1].sort_by { |k,v| -v[sort_key]}.each do |k,v|
					io.printf "%+8d%s: %s (%d)\n", v[:size]/units, UNIT_STRINGS[units], k.name, current()[k][:count] unless v == 0
				end
					
				80.times do io.print( '-' ); end
				io.puts
				io.flush
			rescue Exception => ex
				STDERR.puts "Unhandled exception: #{ex.message}"
				STDERR.puts ex.backtrace.join( "\n" )
			end
		end
				
		def MemoryProfiler::capture()
			current().clear( )
			@@curr_strings = [] if @@opt[:string_debug]

			ObjectSpace.each_object do |o|
				current()[o.class] = { } unless ( current().has_key?( o.class ) )
				current()[o.class][:count] = 0 unless ( current()[o.class].has_key?( :count ) )
				current()[o.class][:size] = 0 unless ( current()[o.class].has_key?( :size ) )
				current()[o.class][:count] += 1
				current()[o.class][:size] += Marshal.dump(o).size rescue 1
				if ( @@opt[:string_debug] and ( String == o.class ) )
					@@curr_strings.push o
				end
			end
		end

		def MemoryProfiler::write_string_debug( )
			if ( @@opt[:string_debug] )
				File.open( "#{@@opt[:string_log_file]}.#{Time.now.to_i}", 'w' ) do |f|
					@@curr_strings.sort.each do |s|
						f.puts s
					end
				end
				@@curr_strings.clear
			end
		end

		def MemoryProfiler::calc_deltas( delta )
			delta.clear( )
			current().keys.each do |k,v|
				delta[k] = { } unless ( delta.has_key?( k ) )
				delta[k][:count] = 0 unless ( delta[k].has_key?( :count ) )
				delta[k][:size] = 0 unless ( delta[k].has_key?( :size ) )
				delta[k][:size] = ( current()[k][:size] - previous()[k][:size] ) if ( previous().has_key?( k ) )
				delta[k][:size] = current()[k][:size] unless ( previous().has_key?( k ) )
				delta[k][:count] = ( current()[k][:count] - previous()[k][:count] ) if ( previous().has_key?( k ) )
				delta[k][:count] = current()[k][:count] unless ( previous().has_key?( k ) )
			end
		end
		
		def MemoryProfiler::snapshot( sort_key = :size )
			MemoryProfiler::capture( ) 
			MemoryProfiler::display( sort_key )
		end
		
		#
		#
		#
		def MemoryProfiler::start( opt = {} )
			opt = DEFAULTS.dup.merge(opt)
		  	
		  	@@opt = opt
			@@running = true
			@@profiler_thread = Thread.new do
			Thread.current[:name] = PROFILER_THREAD_NAME
			@@previous = { }
			@@current = { }
			@@curr_strings = []
			delta = { }
			begin
				file = File::open( opt[:log_file], 'w' )
			rescue Exception => err
				STDERR.puts "** memory_profiler file open error: #{err}"
				STDERR.puts err.backtrace.join( "\n" )			
			end
			
			while ( @@running ) do
				puts "Log memory to #{file.path}" 				
				begin
					GC.start( )
					
					MemoryProfiler::capture( ) 
					 
					MemoryProfiler::write_string_debug( )
	
					MemoryProfiler::calc_deltas( delta )					 
					 
					MemoryProfiler::display( :size, 20 , file )
					  
					delta.clear
					previous().clear
					previous().update current()
					GC.start( )
				rescue Exception => err
					STDERR.puts "** memory_profiler error: #{err}"
					STDERR.puts err.backtrace.join( "\n" )
				end
				sleep opt[:delay]
				end
			end
		end

		#--------------------------------------------------------------------
		# Private
		#--------------------------------------------------------------------
	private
		PROFILER_THREAD_NAME = '__MemoryProfiler_Thread'
		UNIT_STRINGS = {
			UNIT_BYTES => 'B',
			UNIT_KILOBYTES => 'KB',
			UNIT_MEGABYTES => 'MB'
		}
	end

end # Util module
end # Pipeline module


if ( __FILE__ == $0 ) then


	Pipeline::Util::MemoryProfiler::start( )
	s = 'STRING STRING 1'
	s = 'STRIN GSTRING 2'
	s = 'STRING STRING 1'
	while true do
		s = 'STRIN GSTRING 2'
	end
	Pipeline::Util::MemoryProfiler::stop( )
end

# memory_profiler.rb