#!/usr/bin/ruby
#
# 1) generate I/O patterns
# 2) play I/O traces
#
# 2006, 2007	Fengguang Wu

PLAYIO_VERSION='0.1'

require 'optparse'
require 'ostruct'
require 'mmap'

CMDLINE = ARGV.join ' '

PAGE_SHIFT = 12
PAGE_SIZE = 1 << PAGE_SHIFT
UNITS = {?g => 30, ?m => 20, ?p => PAGE_SHIFT, ?k => 10, ?b => 0 }
$base_shift = PAGE_SHIFT

$options = OpenStruct.new
$options.unit = ?p
$options.align_mask = 0xffffffffff
$options.pos_bounds = [0, 10 << 20]		# min, max
$options.len_bounds = [PAGE_SIZE, PAGE_SIZE]
$options.skip_bounds = [0, 0]
$options.max_bytes = 1 << 48
$options.max_time = 1000
$options.max_requests = 1 << 48
$options.pattern = 'sequential'
$options.think_time = 0
$options.pdf = 'uniform'

$pages = Hash.new

def s2i(s)
	i = s.to_i
	shift = UNITS[s[-1]] || UNITS[$options.unit]
	$base_shift = shift if $base_shift > shift
	i <<= shift
	i
end

def i2s(i)
	UNITS.each do |suffix, shift|
		if i & ((1 << shift) - 1) == 0 and $base_shift >= shift
			return (i >> shift).to_s + suffix.chr
		end
	end
	i.to_s
end

class IOTrace

	def initialize
		@io = Array.new
	end

	def io
		@io
	end

	def create
		srand Time::now.usec
		pat_first = self.method("#{$options.pattern}_first")
		pat_next  = self.method("#{$options.pattern}_next")
		bytes = 0
		request = pat_first.call
		loop do
			bytes += request[1]
			@io << request
			break if bytes >= $options.max_bytes
			break if @io.size >= $options.max_requests
			request = pat_next.call
			break if request == nil
		end
	end

	def load(file)
		File.open(file).each_line do |line|
			pos, len = line.split
			next if pos[0] == ?#
			len ||= '1p'
			@io << [s2i(pos), s2i(len)]
		end
	end

	def save(file)
		File.open(file, 'w') do |f|
			f.puts "# iotrace.rb #{CMDLINE}"
			@io.each do |r|
				pos, len = r
				f.print i2s(pos)
				f.print "\t", i2s(len) if len != (1 << $base_shift)
				f.print "\n"
			end
		end
	end

	def play(bdev)
		system("fadvise #{bdev} 0 0 dontneed")
		ppos = -1
		File.open(bdev) do |f|
			seeks = 0
			bytes = 0
			start = Time.now
			@io.each do |r|
				pos, len = r
				if pos != ppos
					f.sysseek(pos)
					seeks += 1
				end
				f.sysread(len)
				ppos = pos + len
				bytes += len
				break if bytes >= $options.max_bytes
				break if Time.now - start >= $options.max_time
				sleep 0.001 * $options.think_time
			end
			stop = Time.now
			throughput = bytes / (stop - start)
			runlen = bytes / seeks
			printf "%fMB/s %fs %dx %db %s\n", throughput/(1<<20), stop - start, seeks, runlen, CMDLINE
		end
	end

	def mplay(bdev)
		system("fadvise #{bdev} 0 0 dontneed")
		mmap_file = Mmap.new(bdev,
				     mode = "r",
				     protection = Mmap::MAP_SHARED,
				     options = {})
		# 'length' => 1*1024*1024*1024
		# 'advice' => Mmap::MADV_NORMAL, Mmap::MADV_RANDOM, Mmap::MADV_SEQUENTIAL, Mmap::MADV_WILLNEED, Mmap::MADV_DONTNEED
		data = 0
		bytes = 0
		@io.each do |r|
			pos, len = r
			mmap_file[pos, len]
			bytes += len
			break if bytes >= $options.max_bytes
			# pgoff = -1
			# pos.upto(pos+len) { |offset|
				# next if pgoff == offset / 4096
				# pgoff = offset / 4096
				# data = data | mmap_file[offset]
			# }
			sleep 0.001 * $options.think_time
		end
	end

	def break_up(file)
		bio = Array.new
		time = 1
		File.open(file, 'w') do |f|
			f.puts "# playio.rb #{CMDLINE}"
			@io.each do |r|
				pos, len = r
				loop do
					l = (len < alen ? len : alen)
					bio << [pos, l]
					f.print time
					f.print "\t", i2s(pos)
					f.print "\t", i2s(l) if l != (1 << $base_shift)
					f.print "\n"
					len -= l
					break if len <= 0
					pos += l
				end
				time = time + 1
			end
		end
		@io = bio
	end

	def round_up
		@io.each_index do |i|
			pos, len = @io[i]
			epos = (pos + len + ~$options.align_mask) & $options.align_mask
			pos &= $options.align_mask
			@io[i] = [pos, epos - pos ]
		end
	end

	def shuffle(count, scope)
		i = rand(scope)
		count.times do
			j = (i + rand(scope)) % (@io.size)
			tmp = @io[i]
			@io[i] = @io[j]
			@io[j] = tmp
			i = j
		end
	end

	def interleave(b)
		aio = Array.new
		i = 0
		j = 0
		n = @io.size + b.io.size
		loop do
			r = rand(n)
			if r < @io.size and i < @io.size
				aio << @io[i]
				i += 1
			elsif j < b.io.size
				aio << b.io[j]
				j += 1
			elsif i < @io.size
				aio << @io[i]
				i += 1
			else
				break
			end
		end
		@io = aio
	end

	def arand(bounds)
		n = bounds[0]
		len = bounds[1] - bounds[0] + 1
		if $options.pdf == 'uniform'
			n += rand(len)
		elsif $options.pdf == 'clustered'
			n = len * (0.5 + (0.5 - rand()) * rand()).to_i
		end
		n &= $options.align_mask
	end
	def apos
		arand $options.pos_bounds
	end
	def alen
		arand $options.len_bounds
	end
	def askip
		arand $options.skip_bounds
	end

	def random_first
		[apos, alen]
	end
	alias random_next random_first

	# out-of-order: (non-overlapping) non-repeating random
	def ooo_first
		pos = 0
		loop do
			pos = apos
			if not $pages[pos]
				$pages[pos] = 1
				break
			end
		end
		[pos, alen]
	end
	alias ooo_next ooo_first

	def sequential_first
		[$options.pos_bounds[0], alen]
	end

	def sequential_next
		pos = @io.last[0] + @io.last[1] + askip
		[pos, alen] if pos < $options.pos_bounds[1]
	end

	def backward_first
		[$options.pos_bounds[1], alen]
	end

	def backward_next
		len = alen
		pos = @io.last[0] - len - askip
		[pos, len] if pos >= $options.pos_bounds[0]
	end

end

iotrace = IOTrace.new


opts = OptionParser.new do |opts|
	opts.banner = "Usage: playio.rb [$options]"

	opts.separator ""
	opts.separator "Parameters:"

	opts.on("--unit p|b|k|m|g", "set default size unit: Page/Byte/Kilo/Mega/Giga") do |u|
		$options.unit = u[0]
	end

	opts.on("--max-requests REQUESTS", "define total I/O requests limit") do |n|
		$options.max_requests = n
	end

	opts.on("--max-bytes BYTES", "define total I/O bytes limit") do |n|
		$options.max_bytes = s2i(n)
	end

	opts.on("--max-seconds SECONDS", "define total I/O time limit") do |n|
		$options.max_time = n.to_f
	end

	opts.on("--offset [BEGIN=0,]END", "define I/O offset range") do |list|
		a, b = list.split ','
		if b == nil
			b = a
			a = 0
		end
		a = s2i(a)
		b = s2i(b)
		$options.pos_bounds = [a, b]
	end

	opts.on("--size MIN[,MAX=MIN]", "define I/O size range") do |list|
		a, b = list.split ','
		b ||= a
		a = s2i(a)
		b = s2i(b)
		$options.len_bounds = [a, b]
	end

	opts.on("--skip MIN[,MAX=MIN]", "define I/O skip range") do |list|
		a, b = list.split ','
		b ||= a
		a = s2i(a)
		b = s2i(b)
		$options.skip_bounds = [a, b]
	end

	opts.on("--align BOUNDARY", "align I/O offset/size to BOUNDARY") do |n|
		$options.align_mask = ~(s2i(n) - 1)
	end

	opts.on("--pattern PATTERN", "define I/O pattern") do |pat|
		$options.pattern = pat
	end

	opts.on("--pdf DISTRIBUTION", "define random pattern") do |pdf|
		$options.pdf = pdf
	end

	opts.on("--think-time MSEC", "define think time between I/O requests") do |t|
		$options.think_time = t.to_f
	end

	opts.separator ""
	opts.separator "Actions:"

	opts.on("-c", "--create", "create I/O trace") do |file|
		iotrace.create
	end

	opts.on("-b", "--break-up [FILE]", "break up to new size and save time-offset-size") do |file|
		iotrace.break_up file || '/dev/null'
	end

	opts.on("-r", "--round-up", "round up to new align") do
		iotrace.round_up
	end

	opts.on("-x", "--shuffle [N,M]", "shuffle N times with locality M") do |list|
		n, m = list.split ','
		iotrace.shuffle((n.to_i || iotrace.io.size), (m.to_i || iotrace.io.size))
	end

	opts.on("-l", "--load FILE", "load I/O trace from FILE") do |file|
		iotrace.load file
	end

	opts.on("-i", "--interleave FILE", "interleave I/O with FILE") do |file|
		b = IOTrace.new
		b.load file
		iotrace.interleave b
	end

	opts.on("-s", "--save FILE", "save I/O trace to FILE") do |file|
		iotrace.save file
	end

	opts.on("-p", "--play [FILE]", "play I/O trace on FILE") do |file|
		iotrace.play file || '/tmp/sparse/100G'
	end

	opts.on("--mplay [FILE]", "play I/O trace on FILE") do |file|
		iotrace.mplay file || '/tmp/sparse/100G'
	end

	opts.on("--sleep SEC", "sleep SEC seconds") do |sec|
		sleep sec.to_i
	end

	opts.separator ""
	opts.separator "Others:"
      
	opts.on_tail("--pid FILE", "Save pid to FILE") do |file|
		File.open(file, 'w') do |f|
			f.puts Process.pid
		end
	end

	opts.on_tail("-h", "--help", "Show this message") do
		puts opts
		exit
	end

	opts.on_tail("-v", "--version", "Show version") do
		puts PLAYIO_VERSION
		exit
	end

end

opts.parse!(ARGV)


# random
# ruby -e 'File.open("iotrace", "w") { |f| 100000.downto(1) { f.puts rand(1000000) } }'
#
# backward
# seq 10000 -1 0 > iotrace
#
# clear
# fadvise sparse 0 10000000000 dontneed