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

verbose = False
def debug( *args, **kargs ):
   if verbose:
      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

parser = optparse.OptionParser()

blockSize = 4096

syncTimeout = 30
flushTimeout = 30
readTimeout = 120

default_patterns = [0xffffffff,0x5a5a5a5a,0xa5a5a5a5,0x00000000]
default_pattern_str = ", ".join( ["%08x" % i for i in default_patterns])
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")
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:
   file( '/proc/sys/vm/drop_caches', 'w' )
except IOError, e:
   print >>sys.stderr, "Error: unable to open /proc/sys/vm/drop_caches for writing."
   print >>sys.stderr, \
      "I need to open this file in order to flush filesystem caches."
   if not os.getuid() == 0:
      print >>sys.stderr, "I noticed you are not running as user root."\
            "  You probably need to be root."
   sys.exit(2)


# --------------------------------
# Process options.size
# --------------------------------
m = re.match( "(\d+)([kKMm])?", options.size )
if not m:
   parser.error( "invalid size (-s/--size) specification: " + options.size )
bytes = int( m.group(1) )
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 ):
   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):
   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
   file( '/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 ):
   global i
   startPhase( "Doing pattern 0x%08x" % struct.unpack("L",pattern))

   tf = file( path, "w")

   # Write the file
   startPhase( "writing" )
   for i in xrange( 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 = file( path, 'r' )
   for i in xrange( numBlocks ):
      data = rf.read( len( block ) )
      if data != block:
         from Hexdump import hexdump
         print >>sys.stderr, "Error reading block", i, \
               "expected:", pattern, \
               "saw:\n", hexdump( data )
         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, e:
            if options.keep:
               print "Test failed: file is saved in ", path
            else:
               if e.errno != 2: raise
try:
   if options.iterations:
      for i in xrange(int(options.iterations)):
         doOneIteration()
   else:
      while True:
         doOneIteration()
except KeyboardInterrupt:
   print "Interrupted after",ii,"successful iterations"
            


