gistfile1.py

#!/home/kgorman/python/bin/python
#
# simple load test script
#

import datetime, os, time, sys, random
import pymongo
import settings
from pymongo import ASCENDING, DESCENDING
from pymongo import Connection
from optparse import OptionParser

import gridfs
from pymongo.objectid import ObjectId

def db_load(keys,ss,fs):
  # gridFS
  random.jumpahead(ss)
  k=random.randint(0,500)
  o = keys[k]["oid"]
  gfs = fs.get(o)
  out = gfs.read()

def getkeys(myfiles):
    c=myfiles.find()
    return c
   
def main():

  # command line options
  parser = OptionParser()
  parser.set_defaults(host="localhost",port="27017",readonly=False)
  parser.add_option("-p","--processes",dest="processcount",help="number of processes to fork")
  parser.add_option("-l","--loops",dest="loops",help="number of times to execute code, (sustain the load)")
  parser.add_option("--port",dest="port",type=int,help="host port to connect to")
  parser.add_option("--host",dest="host",help="host to connect to")
  (options, args) = parser.parse_args()

  # connections to Mongo
  connection = Connection( options.host , options.port )
  db = connection.test
  gfs_test=db.gfs_test
  fs = gridfs.GridFS(db)
  myfiles=db.myfiles

  # get structure of oid's to fetch
  keys=getkeys(myfiles)
  
  before = db.command( { "serverStatus" : 1 } )["opcounters"]
  startTime = time.time()

  # main fork/loop
  pids = []
  didIFork = False
  for i in xrange(int(options.processcount)):
    pid = os.fork()
    if pid:
      pids.append( pid )
    else:
      didIFork = True

      for ii in xrange(int(options.loops)):
        db_load(keys,i,fs)
      
      break

  if not didIFork:
    for x in pids:
      os.waitpid( x , 0 )
    elapsed = time.time() - startTime
    print( "main done" )
    
    after = db.command( { "serverStatus" : 1 } )["opcounters"]
    for x in after:
      diff = after[x] - before[x]
      if diff == 0 :
        continue
      print( x + "/sec \t" + str( int( diff / elapsed ) ) );

  
if __name__ == "__main__":
  main()