#!/usr/bin/env python3

"""
This application is a stress tester for libjio. It's not a traditional stress
test like fsx (which can be used to test libjio using the preloading library),
but uses fault injection to check how the library behaves under random
failures.
"""

import sys
import os
import random
import traceback
import libjio

try:
	import fiu
except ImportError:
	print()
	print("Error: unable to load fiu module. This test needs libfiu")
	print("support. Please install libfiu and recompile libjio with FI=1.")
	print()
	raise

#
# Auxiliary stuff
#

gbcount = 0
def getbytes(n):
	global gbcount
	gbcount = (gbcount + 1) % 10
	return bytes(str(gbcount) * n, 'ascii')

def randfrange(maxend, maxsize):
	start = random.randint(0, maxend - 1)
	size = random.randint(0, (maxend - 1) - start) % maxsize
	return start, start + size

def randfloat(min, max):
	return min + random.random() % (max - min)

class ConsistencyError (Exception):
	pass

def jfsck(fname):
	try:
		r = libjio.jfsck(fname)
		return r
	except IOError as e:
		if e.args[0] == libjio.J_ENOJOURNAL:
			return { 'total': 0 }
		else:
			raise

def comp_cont(bytes):
	"'aaaabbcc' -> [ ('a', 4), ('b', 2), ('c', 2) ]"
	l = []
	prev = bytes[0]
	c = 1
	for b in bytes[1:]:
		if (b == prev):
			c += 1
			continue

		l.append((prev, c))
		prev = b
		c = 1
	return l


#
# The test itself
#

class Stresser:
	def __init__(self, fname, fsize, nops, use_fi, use_as):
		self.fname = fname
		self.fsize = fsize
		self.nops = nops
		self.use_fi = use_fi
		self.use_as = use_as

		self.maxoplen = min(int(self.fsize / 256),
					64 * 1024)

		jflags = 0
		if use_as:
			jflags = libjio.J_LINGER

		self.jf = libjio.open(fname, libjio.O_RDWR | libjio.O_CREAT,
				0o600, jflags)
		self.f = open(fname, mode = 'rb')

		self.jf.truncate(fsize)

		if use_as:
			self.jf.autosync_start(5, 10 * 1024 * 1024)

		# data used for consistency checks
		self.current_range = (0, 0)
		self.prev_data = b""
		self.new_data = b""

	def pread(self, start, end):
		ppos = self.f.tell()
		self.f.seek(start, 0)
		r = bytes()
		c = 0
		total = end - start
		while c < total:
			n = self.f.read(total - c)
			if (n == ''):
				break
			c += len(n)
			r += n
		self.f.seek(ppos, 0)
		assert c == end - start
		return r

	def prep_randwrite(self):
		start, end = randfrange(self.fsize, self.maxoplen)

		# read an extended range so we can check we
		# only wrote what we were supposed to
		estart = max(0, start - 32)
		eend = min(self.fsize, end + 32)
		self.current_range = (estart, eend)
		self.prev_data = self.pread(estart, eend)

		nd = getbytes(end - start)
		self.new_data = self.prev_data[:start - estart] \
			+ nd + self.prev_data[- (eend - end):]
		return nd, start

	def randwrite(self, nd, start):
		self.jf.pwrite(nd, start)
		return True

	def randwrite_fork(self):
		# do the prep before the fork so we can verify() afterwards
		nd, start = self.prep_randwrite()
		sys.stdout.flush()
		pid = os.fork()
		if pid == 0:
			# child
			try:
				self.fiu_enable()
				self.randwrite(nd, start)
				self.fiu_disable()
			except (IOError, MemoryError):
				try:
					self.reopen()
				except (IOError, MemoryError):
					pass
				except:
					self.fiu_disable()
					traceback.print_exc()
				self.fiu_disable()
				sys.exit(1)
			except MemoryError:
				self.fiu_disable()
				sys.exit(1)
			except:
				self.fiu_disable()
				traceback.print_exc()
				sys.exit(1)
			sys.exit(0)
		else:
			# parent
			id, status = os.waitpid(pid, 0)
			if not os.WIFEXITED(status):
				raise RuntimeError(status)

			if os.WEXITSTATUS(status) != 0:
				return False
			return True

	def verify(self):
		# NOTE: must not use self.jf
		real_data = self.pread(self.current_range[0],
				self.current_range[1])
		if real_data not in (self.prev_data, self.new_data):
			print('Corruption detected')
			print('Range:', self.current_range)
			print('Real:', comp_cont(real_data))
			print('Prev:', comp_cont(self.prev_data))
			print('New: ', comp_cont(self.new_data))
			print()
			raise ConsistencyError

	def reopen(self):
		self.jf = None
		r = jfsck(self.fname)

		self.verify()

		self.jf = libjio.open(self.fname,
			libjio.O_RDWR | libjio.O_CREAT, 0o600)
		return r

	def fiu_enable(self):
		if not self.use_fi:
			return

		# To improve code coverage, we randomize the probability each
		# time we enable failure points
		fiu.enable_random('jio/*',
				probability = randfloat(0.0005, 0.005))
		fiu.enable_random('linux/*',
				probability = randfloat(0.005, 0.03))
		fiu.enable_random('posix/*',
			probability = randfloat(0.005, 0.03))
		fiu.enable_random('libc/mm/*',
			probability = randfloat(0.003, 0.07))
		fiu.enable_random('libc/str/*',
			probability = randfloat(0.005, 0.07))

	def fiu_disable(self):
		if self.use_fi:
			fiu.disable('libc/mm/*')
			fiu.disable('posix/*')
			fiu.disable('jio/*')
			fiu.disable('linux/*')

	def run(self):
		nfailures = 0
		sys.stdout.write("  ")
		for i in range(1, self.nops + 1):
			sys.stdout.write(".")
			if i % 10 == 0:
				sys.stdout.write(" ")
			if i % 50 == 0:
				sys.stdout.write(" %d\n" % i)
				sys.stdout.write("  ")
			sys.stdout.flush()

			if self.use_fi:
				r = self.randwrite_fork()
			else:
				nd, start = self.prep_randwrite()
				r = self.randwrite(nd, start)
			if not r:
				nfailures += 1
				r = self.reopen()
			self.verify()
		sys.stdout.write("\n")
		sys.stdout.flush()
		return nfailures


#
# Main
#

def usage():
	print("""
Use: jiostress <file name> <file size in Mb> [<number of operations>]
	[--fi] [--as]

If the number of operations is not provided, the default (1000) will be
used.

If the "--fi" option is passed, the test will perform fault injection. This
option conflicts with "--as".

If the "--as" option is passed, lingering transactions will be used, along
with the automatic syncing thread. This option conflicts with "--fi".
""")


def main():
	try:
		fname = sys.argv[1]
		fsize = int(sys.argv[2]) * 1024 * 1024
		nops = 1000
		if len(sys.argv) >= 4 and sys.argv[3].isnumeric():
			nops = int(sys.argv[3])

		use_fi = False
		if '--fi' in sys.argv:
			use_fi = True

		use_as = False
		if '--as' in sys.argv:
			use_as = True
	except:
		usage()
		sys.exit(1)

	if use_fi and use_as:
		print("Error: --fi and --as cannot be used together")
		sys.exit(1)

	s = Stresser(fname, fsize, nops, use_fi, use_as)
	print("Running stress test")
	nfailures = s.run()
	del s
	print("Stress test completed")
	print("  %d operations" % nops)
	print("  %d simulated failures" % nfailures)

	r = jfsck(fname)
	print("Final check completed")
	#os.unlink(fname)


if __name__ == '__main__':
	main()

