#!/usr/bin/env python3
# Copyright (c) 2007, 2011 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

# pylint: disable=wrong-import-position

verbose = False
def debug( *args, **kargs ): # pylint: disable=redefined-outer-name
   if verbose:
      # pylint: disable-next=consider-using-f-string
      sys.stdout.write( "%5.3f " % (time.time() - startTime) )
      sys.stdout.write( " ".join( [str(i) for i in args] ) + "\n" )

import time
startTime = time.time()
import sys, os, struct, signal, optparse, re # pylint: disable=deprecated-module

parser = optparse.OptionParser()

blockSize = 4096

syncTimeout = 30
flushTimeout = 30
readTimeout = 120

default_patterns = [0xffffffff,0x5a5a5a5a,0xa5a5a5a5,0x00000000]
# pylint: disable-next=consider-using-f-string
default_pattern_str = ", ".join( ["%08x" % i for i in default_patterns])
# pylint: disable-next=consider-using-f-string
parser.set_usage( usage="""
%%prog [options]

Example that writes and then reads back a 32M file 10 times to /mnt/flash:

  %%prog -d /mnt/flash -s 32M -i 10

This program writes SIZE bytes, %d bytes in each call to write().  It
waits until the file is written to disk, then flushes the filesystem's
buffer cache, then reads the file back and compares each block to the
expected value.  It repeats this for each of %d patterns:
   %s
and then unlinks the file and repeats the whole process ITERATIONS
times.  If ITERATIONS is 0, then it repeats continuously.

The program exits when all iterations are complete, or when control-c is hit.

""" % (blockSize,len(default_patterns), default_pattern_str ))
       

parser.add_option( "-d", "--dir", help="specify dir to write to (default %default)",
                   default="/tmp" )
parser.add_option( "--keep", action="store_true",
                   help="keep the file around in the event of a failure")
parser.add_option( "-c", "--continuous", action="store_true",
                   help="Run continuously")
parser.add_option( "-v", "--verbose", action="store_true",
                   help="Verbose output")
# pylint: disable-next=consider-using-f-string
parser.add_option( "-p", "--pattern", help="Pattern to use when writing (default %s"
                   % default_pattern_str )
parser.add_option( "-i", "--iterations",
                   help="How many times to run (default %default)",
                   default=1, type=int )
parser.add_option( "-s", "--size",
                   help="Size to write, in bytes (k and m suffixes allowed),"
                   " default=%default",
                   default=str(blockSize * 1024 ))

options,args = parser.parse_args()
if args:
   parser.usage()

if options.continuous and ( options.iterations is not None ):
   parser.error("Both -c and -i options cannot be specified simultaneously")

verbose = options.verbose

debug( "parsing is done" )

try:
   open( '/proc/sys/vm/drop_caches', 'w' ) # pylint: disable=consider-using-with
except OSError as e:
   print( "Error: unable to open /proc/sys/vm/drop_caches for writing.",
      file=sys.stderr )
   print( "I need to open this file in order to flush filesystem caches.",
      file=sys.stderr )
   if not os.getuid() == 0:
      print( "I noticed you are not running as user root."
             "  You probably need to be root.", file=sys.stderr )
   sys.exit(2)


# --------------------------------
# Process options.size
# --------------------------------
m = re.match( r"(\d+)([kKMm])?", options.size )
if not m:
   parser.error( "invalid size (-s/--size) specification: " + options.size )
bytes = int( m.group(1) ) # pylint: disable=redefined-builtin
suffix = m.group(2)
if suffix:
   if suffix in "kK":
      bytes *= 1024
   if suffix in "Mm":
      bytes *= 1024 * 1024

# --------------------------------
# Process options.dir
# --------------------------------
where = options.dir
if not os.path.exists( where ):
   parser.error( "Target directory (-d,--dir) %s does not exist " + where )
if not os.path.isdir( where ):
   parser.error(
      "Target directory (-d,--dir) %s exists, but is not a directory " + where )


# --------------------------------
# Figure out the set of patterns we're writing
# --------------------------------
patterns = [struct.pack("L",i)
            for i in default_patterns]
# Handle a user-specified pattern
p = options.pattern
if p:
   if p.startswith( "0x" ):
      patterns = [ struct.pack("L",int(p)) ]
   else:
      patterns = [p]

# Now start
numBlocks = int((bytes + (blockSize -1)) / blockSize)
debug( "writing", bytes, "bytes in", numBlocks, "blocks" )

blocks = [(blockSize // len(i)) * i for i in patterns]

# --------------------------------
# Set up an alarm signal handler so we don't hang in this test forever
# --------------------------------
phase = "initialization"
phaseTimeout = 0
def alarm( *args ): # pylint: disable=redefined-outer-name
   exitWithError( "Phase", phase, "of iteration", i,
                  "took too long and we gave up after", phaseTimeout, "seconds")
signal.signal( signal.SIGALRM, alarm )

# --------------------------------
# Helper functions and their global state
# --------------------------------
i = None
def startPhase( ph, timeout=0 ):
   global phase
   global phaseTimeout
   phaseTimeout = timeout
   phase = ph
   debug( "starting phase", phase )
   signal.alarm(timeout)

def exitWithError( *args): # pylint: disable=redefined-outer-name
   sys.stderr.write( " ".join( [ str(i) for i in args ] ) )
   sys.exit( 1 )

def flushFsCache():
   # 1 flushes pages
   # 2 flushes inodes and dentries
   # 3 flushes pages, inodes, and dentries
   # only clean objects are flushed, so sync must be called first
   # pylint: disable-next=consider-using-with
   open( '/proc/sys/vm/drop_caches', 'w' ).write( '1' )


#--------------------------------
# doPattern does most of the work.  It writes 'block' repeatedly,
# which is derived from a repeating sequence of "pattern" to file
# 'path'.
#--------------------------------
def doPattern( path, block, pattern ): # pylint: disable=redefined-outer-name
   global i
   # pylint: disable-next=consider-using-f-string
   startPhase( "Doing pattern 0x%08x" % struct.unpack("L",pattern))

   tf = open( path, "w" ) # pylint: disable=consider-using-with

   # Write the file
   startPhase( "writing" )
   for i in range( numBlocks ):
      tf.write( block )

   startPhase( "syncing", syncTimeout )
   # Sync the data to disk
   os.fdatasync( tf.fileno() )

   tf.close()

   startPhase( "flushing", flushTimeout )
   flushFsCache()

   startPhase( "reading", readTimeout )
   # Read the data back
   rf = open( path ) # pylint: disable=consider-using-with
   for i in range( numBlocks ):
      data = rf.read( len( block ) )
      if data != block:
         from Hexdump import hexdump # pylint: disable=import-outside-toplevel
         print( "Error reading block", i,
                "expected:", pattern,
                "saw:\n", hexdump( data ), file=sys.stderr )
         sys.exit( 1 )

#--------------------------------
# Run the test, making one call to doPattern per pattern
# We make sure to clean up the file unless --keep is specified.
#--------------------------------
path = os.path.join( where, "fstest" )

ii = 0
def doOneIteration():
   global ii
   ii += 1
   for (block,pattern) in zip(blocks,patterns):
      try:
         doPattern( path, block, pattern )
      finally:
         # Make this happen unless --save
         try:
            os.unlink( path )
         except OSError as e: # pylint: disable=redefined-outer-name
            if options.keep:
               print( "Test failed: file is saved in ", path )
            else:
               if e.errno != 2: raise # pylint: disable=multiple-statements
try:
   if options.iterations:
      for i in range( int( options.iterations ) ):
         doOneIteration()
   else:
      while True:
         doOneIteration()
except KeyboardInterrupt:
   print( "Interrupted after",ii,"successful iterations" )
            


