Mail Stress Testing

mailanalysis.py

import sys
import glob
import re
import pandas as pd
import numpy as np
import datetime

STARTTIMEREGEX = re.compile(r'''^\#([0-9]{2}):([0-9]{2}):([0-9]{2})''')
COMMENTREGEX = re.compile(r'''^//''')
LINEREGEX = re.compile(r'''^\{([0-9]{2})\:([0-9]{2})\:([0-9]{2})\} (.+?@.+?):(.+?) \[([0-9\.]{6,15})\]$''')
BYTEREGEX = re.compile(r'''(append|retr)\((\d+)\)''')

def _convtodataframe(protodata):
  protoseriesdata = {}
  for cmd, l in protodata.items():
    protoseriesdata[cmd] = pd.Series(l)
  
  frame = pd.DataFrame(protoseriesdata)
  return frame

def _addtoprotodata(cmd, duration, protodata):
  match2 = BYTEREGEX.search(cmd)
  if match2:
    cmd = match2.group(1).lower()+"perbyte"
    bytes = int(match2.group(2))
    if cmd not in protodata:
      protodata[cmd] = []
    protodata[cmd].append(duration/bytes)
  else:
    if cmd not in protodata:
      protodata[cmd] = []
    protodata[cmd].append(duration)

def process(proto, rampup):
  allprotodata = {}
  
  if rampup:
    rampupdonetime = None
    postprotodata = {}
    preprotodata = {}
  
  protofiles = glob.glob("./" + proto + "*.stresslog")
  for protofile in protofiles:
    #print "Opening " + protofile
    with open(protofile) as f:
      startlogtime = None
      for line in f:
        #check for start time
        starttime = STARTTIMEREGEX.search(line)
        if starttime:
          if rampup and not startlogtime:
            hour = int(starttime.group(1))
            minute = int(starttime.group(2))
            second = int(starttime.group(3))
            startlogtime = datetime.time(hour, minute, second)
            if not rampupdonetime:
              logdatetime = datetime.datetime.combine(datetime.date.today(), startlogtime)
              rampupdonetime = (logdatetime + datetime.timedelta(seconds=rampup)).time()
          continue
        
        #check for comment
        comment = COMMENTREGEX.search(line)
        if comment:
          continue
        
        #move on to real processing
        match = LINEREGEX.search(line)
        if match:
          #extract data from regex
          hour = int(match.group(1))
          minute = int(match.group(2))
          second = int(match.group(3))
          logtime = datetime.time(hour, minute, second)
          user = match.group(4)
          cmd = match.group(5)
          duration = float(match.group(6))
          
          #add to allprotodata
          _addtoprotodata(cmd, duration, allprotodata)
          
          if rampup:
            if logtime <= rampupdonetime:
              _addtoprotodata(cmd, duration, preprotodata)
            else:
              _addtoprotodata(cmd, duration, postprotodata)
        else:
          sys.stderr.write("Line didn't match: %s" % line)
  
  alldataframe = _convtodataframe(allprotodata)
  
  if rampup:
    predataframe = _convtodataframe(preprotodata)
    postdataframe = _convtodataframe(postprotodata)
  
  with open(proto+"results.txt", "w") as f:
    f.write("ALL DATA\n")
    f.write(alldataframe.describe().to_string(justify="left"))
    f.write("\n-----\n")
    if rampup:
      f.write("PRE RAMP UP DATA\n")
      f.write(predataframe.describe().to_string(justify="left"))
      f.write("\n-----\n")
      f.write("POST RAMP UP DATA\n")
      f.write(postdataframe.describe().to_string(justify="left"))
      f.write("\n-----\n")

if __name__ == "__main__":
  try:
    rampup = int(sys.argv[-1])
  except ValueError:
    rampup = None
  process("imap", rampup)
  process("pop", rampup)
  

mailstress.py

#!/usr/bin/env python

import sys
import json
import subprocess
import shlex
import socket
import time
import json
import netifaces
import glob
import os
import os.path
import errno
import signal
from optparse import OptionParser

import stress

ISPARENT = True

if __name__ == "__main__":  
  parser = OptionParser("Usage: %prog [options]")
  parser.add_option("--host", dest="host", help="Target hostname to benchmark")
  parser.add_option("--rampup", dest="rampup", default=None, type="float", help="Time in seconds to ramp up (0 for full power instantly)")
  parser.add_option("--duration", dest="duration", type="float", help="Time in seconds to benchmark over (total)")
  parser.add_option("--imap", dest="imap", default=None, type="int", help="Number of IMAP connections")
  parser.add_option("--pop", dest="pop", default=None, type="int", help="Number of POP connections")
  parser.add_option("--iface-pattern", dest="ifacepattern", default="eth0", type="string", help="Pattern of interfaces to use")
  parser.add_option("--loginfile", dest="loginpath", default="logins.json", type="string", help="Filename of login information")
  parser.add_option("--sleep", dest="sleep", default=None, type="float", help="Amount to sleep after each operation on children connections")
  parser.add_option("--debug", dest="debug", action="count", default=0)
  #parser.add_option("--smtp", dest="smtp", default=0, type="int", help="Number of SMTP connections")
  
  (options, args) = parser.parse_args()
  
  if options.rampup and options.rampup < 0:
    parser.error("--rampup must be >=0, or not specified")
    sys.exit(2)
  
  if options.duration <= 0:
    parser.error("--duration must be >0")
    sys.exit(2)
  
  if options.sleep and options.sleep <= 0:
    parser.error("--sleep must be >0, or not specified")
    sys.exit(2)
  
  if options.pop and options.pop <= 0:
    parser.error("--pop must be >0, or not specified")
    sys.exit(2)
    
  if options.imap and options.imap <= 0:
    parser.error("--imap must be >0, or not specified")
    sys.exit(2)
  
  if not os.path.isfile(options.loginpath):
    parser.error("--loginfile must exist")
    sys.exit(2)
 
  try:
    with open(options.loginpath) as loginfile:
      logins = json.load(loginfile)["logins"]
  except IOError as e:
    if e.errno == errno.EACCES:
      parser.error("Insufficient permissions to read specified --loginfile %s" % options.loginpath)
    else:
      parser.error("Cannot open specified --loginfile %s" % options.loginpath)
    sys.exit(2)
  
  null = open("/dev/null", "w")
  imaptime = time.time()
  imapbabies = []
  imappids = []
  #imapfiles = []
  poptime = time.time()
  popbabies = []
  poppids = []
  #popfiles = []
  
  interfaces = [interface for interface in netifaces.interfaces() if options.ifacepattern in interface]
  
  def forkimap(n):
    interface = interfaces[n%len(interfaces)]
    ip = netifaces.ifaddresses(interface)[netifaces.AF_INET][0]["addr"]
    login = logins[n%len(logins)]
    
    pid = os.fork()
    if pid == 0:
      global ISPARENT
      ISPARENT = False
      stress.imap(n, ip, options.host, login[0], login[1], glob.glob("./*.msg"), options.sleep, options.debug)
      sys.exit(0)
    else:
      print "Forked IMAP process"
      imappids.append(pid)
  
  def forkpop(n):
    interface = interfaces[n%len(interfaces)]
    ip = netifaces.ifaddresses(interface)[netifaces.AF_INET][0]["addr"]
    login = logins[n%len(logins)]
    
    pid = os.fork()
    if pid == 0:
      global ISPARENT
      ISPARENT = False
      stress.pop(n, ip, options.host, login[0], login[1], options.sleep, options.debug)
      sys.exit(0)
    else:
      print "Forked POP process"
      poppids.append(pid)
  
  def killforks():
    for imappid in imappids:
      print "Terminating IMAP fork"
      os.kill(imappid, signal.SIGTERM)
    for poppid in poppids:
      print "Terminating POP fork"
      os.kill(poppid, signal.SIGTERM)
  
  try:
    starttime = time.time()
    #print starttime
    if options.rampup == None:
      if options.imap:
        for n in range(options.imap):
          forkimap(len(imappids))
      if options.pop:
        for n in range(options.pop):
          forkpop(len(poppids))
    
    while time.time() - starttime < options.duration:
      #print time.time() - starttime
      if options.rampup > 0:
        if len(imappids) < options.imap:
          imappersec = options.rampup/options.imap
          while time.time() - imaptime > imappersec:
            imaptime += imappersec
            forkimap(len(imappids))
        if len(poppids) < options.pop:
          poppersec = options.rampup/options.pop
          while time.time() - poptime > poppersec:
            poptime += poppersec
            forkpop(len(poppids))
      
      time.sleep(0.05)
  finally:
    if ISPARENT:
      killforks()
      print "Done."

messagefiller.py

#!/usr/bin/env python

import getopt
import json
import smtplib
import sys
import time
import email.utils
import random
import string
from email.mime.text import MIMEText

ARGS = ""
LONG_ARGS = []

QUIET = False
VERBOSITY = 1

def debugwrite(s, l):
  if VERBOSITY >= l and not QUIET:
    sys.stderr.write(s)

def dumpRecipientsRefused(e):
  '''Takes in an SMTPRecipientsRefused exception and
pretty prints it if verbosity is set.'''
  if VERBOSITY >= 1 and not QUIET:
    sys.stderr.write("Recipient Address Refused.\n")
    for k, v in e.recipients.items():
      sys.stderr.write("%s: %s (%s)\n" % (k, v[0], v[1]))

def sendMessage(smtp, messagestr, env_from_address, sendingtoaddress):
  '''Sends the given message to the given address,
using the given SMTP instance.'''
  debugwrite("Sending a message.\n", 2)
  try:
    smtp.sendmail(env_from_address,
                  sendingtoaddress,
                  messagestr)
  except smtplib.SMTPRecipientsRefused, e:
    dumpRecipientsRefused(e)
  except smtplib.SMTPHeloError, e:
    debugwrite("SMTP server failed to respond correctly to HELO command.\n", 1)
  except smtplib.SMTPSenderRefused, e:
    debugwrite(
        "SMTP server rejected env-from address of %s.\n"
        % env_from_address, 1)
  except smtplib.SMTPDataError, e:
    debugwrite("Unknown STMP error.\n", 1)
  time.sleep(.05)

      
if __name__=="__main__":
  try:
    opts, args = getopt.getopt(sys.argv[1:],
                               ARGS,
                               LONG_ARGS)
  except getopt.GetoptError, err:
    #print help information and exit:
    #will print something like "option -a not recognized"
    print str(err) 
    usage()
    sys.exit(2)
    
  smtp = smtplib.SMTP(args[0])
  #smtp.login(SMTP_USER, SMTP_PASS)
  
  count = int(args[1])
  for n in range(count):
    messagedatafilename = "messagefiller.json"
    
    size = 0
    for nn in range(20):
      size += random.randint(1,6) + 4
    messageraw = ""
    for nn in range(size):
      messageraw = messageraw + random.choice(string.printable)

    dataf = open(messagedatafilename)
    data = json.load(dataf)
    dataf.close()
    
    message = MIMEText(messageraw)
    
    for header, value in data.items():
      message[header] = value

    if "Date" in message:
      del message["Date"]
    message["Date"] = email.utils.formatdate()

    if "From" not in message:
      message["From"] = "bob@fpcomm.net"
    if "Reply-To" not in message:
      message["Reply-To"] = "do-not-reply@fpcomm.net"

    #message["Subject"] = "Reply from FRII regarding your recent support request"
    #message["From"] = "support@frii.com"
    #message["Reply-To"] = "do-not-reply@frii.com"
    
    with open("logins.json") as listingf:
      listing = json.load(listingf)
    
      for emailaddr in [login[0] for login in listing["logins"]]:
        del message["To"]
        if data["To"] == "__RCPTTO__":
          message["To"] = emailaddr
        else:
          message["To"] = data["To"]
        sendMessage(smtp, message.as_string(), "bob@fpcomm.net", emailaddr.lower())
  
  smtp.close()
  

stats.sh

#!/bin/bash

echo "Starting stats.sh" > misc.statlog
date >> misc.statlog

date > iostat.statlog
#iostat -xh 5 >> iostat.statlog &
iostat -h 5 >> iostat.statlog &

date > vmstat.statlog
vmstat 5 >> vmstat.statlog &

sar 5 0 > sar.statlog &

while [ true ]; do
  uptime >> misc.statlog
  free -m >> misc.statlog
  #ps aux | grep dovecot >> misc.statlog
  ps auxww|awk '$1 ~ /dovecot/ || $11 ~ /(pop|imap)/' >> misc.statlog
  echo "-----" >> misc.statlog
  sleep 5
done

stress.py

import traceback
import imaplib
import poplib
import random
import string
import socket
import signal
import sys
import datetime

from timer import Timer

DONE = False

def write_comment(s):
  sys.stderr.write("//" + str(s) + "\n")

def write_comment_exception():
  excstr = traceback.format_exc()
  excstr = "\n".join(["//" + line for line in excstr.split("\n")][:-1])
  sys.stderr.write(excstr+"\n")

def setup(n, proto, user, sleeptime, debug):
  sys.stdout = open("%s%s.stresslog" % (proto, n), "w")
  if debug == 0:
    sys.stderr = open("/dev/null", "w")
  elif debug == 1:
    sys.stderr = sys.stdout
  
  print "#" + datetime.datetime.now().strftime("%H:%M:%S")
  
  signal.signal(signal.SIGTERM, _signalhandler)
  signal.signal(signal.SIGCHLD, _signalhandler)
  
  return maketimermaker(user, sleeptime)

def _signalhandler(signum, frame):
  sys.stdout.flush()
  global DONE
  DONE = True

class BindingPOP(poplib.POP3, object):
  def __init__(self, source, host, port=110):
    self.source = source
    self.host = host
    self.port = port
    #self.sock = socket.create_connection((host, port), timeout)
    self.sock = socket.socket()
    self.sock.bind((self.source, 0))
    self.sock.connect((host, port))
    self.file = self.sock.makefile('rb')
    self._debugging = 0
    self.welcome = self._getresp()

class BindingIMAP(imaplib.IMAP4, object):
  def __init__(self, source, host = '', port = 143):
    self.source = source
    super(BindingIMAP, self).__init__(host, port)
  
  def open(self, host = '', port = 143):
    self.host = host
    self.port = port
    #print "Opening socket"
    self.sock = socket.socket()
    #print "Binding socket to %s" % self.source
    self.sock.bind((self.source, 0))
    #print "Connecting"
    self.sock.connect((host, port))
    self.file = self.sock.makefile('rb')

def maketimermaker(user, sleep=None):
  def timermaker(cmd):
    return Timer("%s:%s" % (user, cmd), sleep)
  return timermaker

def pop(n, source, host, user, passwd, sleeptime, debug):
  maketimer = setup(n, "pop", user, sleeptime, debug)
  
  while not DONE:
    deleted = 0
    try:
      with maketimer("socket"):
        pop = BindingPOP(source, host)
    except:
      write_comment_exception()
      sys.stdout.flush()
      sys.exit(1)
    try:
      with maketimer("login"):
        pop.user(user)
        pop.pass_(passwd)
      with maketimer("stat"):
        count, size = pop.stat()
      if count > 0:
        roll = random.randint(1, 3)
        if roll == 1:
          msgnum = random.randint(1, count)
          with maketimer("list"):
            response = pop.list(msgnum)
          size = response.split(" ")[-1]
          with maketimer("retr(%s)" % size):
            pop.retr(msgnum)
        elif roll == 2:
          msgnum = random.randint(1, count)
          with maketimer("dele"):
            pop.dele(msgnum)
          if random.randint(1, 20) <= 5:
            with maketimer("rset"):
              pop.rset()
              deleted = 0
          else:
            deleted += 1
        elif roll == 3:
          msgnum = random.randint(1, count)
          with maketimer("top"):
            pop.top(msgnum, 0)
      else:
        with maketimer("noop"):
          pop.noop()
    except poplib.error_proto:
      write_comment_exception()
    except:
      write_comment_exception()
      sys.stdout.flush()
      sys.exit(1)
    finally:
      with maketimer("logout(%s)" % deleted):
        pop.quit()
  
  write_comment("POP fork for user %s exiting at %s due to lack of messages" % (user, datetime.datetime.now()))
  sys.stdout.flush()

def imap(n, source, host, user, passwd, mailfiles, sleeptime, debug):
  maketimer = setup(n, "imap", user, sleeptime, debug)
  
  def parsemailboxes(mailboxesraw):
    return [entry.split(" ")[-1] for entry in mailboxesraw[1]]  

  while not DONE:
    selected = False
    try:
      with maketimer("socket"):
        imap = BindingIMAP(source, host)
    except:
      write_comment_exception()
      sys.stdout.flush()
      sys.exit(1)
    try:
      with maketimer("login"):
        imap.login(user, passwd)
      with maketimer("list"):
        mailboxesraw = imap.list()
      mailboxes = parsemailboxes(mailboxesraw)
      write_comment(str(mailboxes))
      madefolders = []
      pushedmsgs = []
      for _ in range(1,random.randint(15,25)):
        roll = random.randint(1,8)
        if roll == 1:
          #upload a message to a random folder
          mailbox = random.choice(mailboxes)
          with maketimer("select"):
            count = int(imap.select(mailbox)[1][0])
          selected = True
          if len(pushedmsgs) >= len(mailfiles):
            #just no-op instead
            with maketimer("noop"):
              imap.noop()
          else:
            #continue uploading file
            #make sure we haven't pushed that file yet
            mailfile = random.choice([mf for mf in mailfiles if mf not in [msg[0] for msg in pushedmsgs]])
            with open(mailfile, "r") as f:
              strng = f.read()
            with maketimer("append(%s)" % len(strng)):
              imap.append(mailbox, None, None, strng)
            response = imap.response("EXISTS")
            write_comment(response)
            newid = response[1][1]
            pushedmsgs.append((mailfile, mailbox, newid))
        elif roll == 2:
          #just no-op
          with maketimer("noop"):
            imap.noop()
        elif roll == 3:
          #read a random message from a random folder
          mailbox = random.choice(mailboxes)
          with maketimer("select"):
            count = int(imap.select(mailbox)[1][0])
          selected = True
          response = imap.response("EXISTS")
          write_comment(response)
          if count > 1:
            msgnum = random.choice(range(1,count))
            with maketimer("fetch"):
              imap.fetch("%s:%s" % (msgnum, msgnum), "(UID BODY[TEXT])")
          else:
            #just no-op
            with maketimer("noop"):
              imap.noop()
        elif roll == 4:
          #make a random folder
          foldername = "_" + "".join(random.sample(string.hexdigits, random.randint(5,8)))
          #check for the folder on the server
          with maketimer("list"):
            mailboxesraw = imap.list("", foldername)
          if mailboxesraw[0] == None:
            #just no-op instead
            with maketimer("noop"):
              imap.noop()
          else:
            with maketimer("create"):
              imap.create(foldername)
            madefolders.append(foldername)
        elif roll == 5:
          #delete a folder we made
          if len(madefolders) == 0:
            #just no-op instead
            with maketimer("noop"):
              imap.noop()
          else:
            #delete a folder
            foldername = random.choice(madefolders)
            madefolders.remove(foldername)
            with maketimer("delete"):
              imap.delete(foldername)
        elif roll == 6:
          #delete a message we pushed
          if len(pushedmsgs) == 0:
            #just no-op instead
            with maketimer("noop"):
              imap.noop()
          else:
            write_comment(str(pushedmsgs))
            msg = random.choice(pushedmsgs)
            mailbox = msg[1]
            id = msg[2]
            sys.stderr.write(str(msg))
            with maketimer("select"):
              imap.select(mailbox)
            selected = True
            with maketimer("store"):
              imap.store(id, "+FLAGS", r"\Deleted")
            with maketimer("expunge"):
              imap.expunge()
            pushedmsgs.remove(msg)
        elif roll == 7:
          #move a message to a random folder
          srcmailbox = random.choice(mailboxes)
          destmailbox = random.choice(mailboxes)
          with maketimer("select"):
            count = int(imap.select(srcmailbox)[1][0])
          selected = True
          if count > 1:
            msgnum = random.choice(range(1,count))
            with maketimer("copy"):
              imap.copy(msgnum, destmailbox)
            with maketimer("store"):
              imap.store(msgnum, "+FLAGS", r"\Deleted")
            with maketimer("expunge"):
              imap.expunge()
            pushedmsgs = [msg for msg in pushedmsgs if msg[2] != msgnum] #remove from pushed messages if we move it
          else:
            #just no-op instead
            with maketimer("noop"):
              imap.noop()
        elif roll == 8:
          with maketimer("recent"):
            imap.recent()
    except imaplib.IMAP4.error:
      write_comment_exception()
    except:
      write_comment_exception()
      sys.stdout.flush()
      sys.exit(1)
    finally:
      if selected:
        with maketimer("close"):
          imap.close()
      with maketimer("logout"):
        imap.logout()
  
  sys.stdout.flush()

timer.py

import time
import sys
import datetime

class Timer(object):
  def __init__(self, label, sleep=None):
    self.label = label
    self.sleep = sleep
  
  def __enter__(self):
    self.starttime = time.time()
  
  def __exit__(self, exc_type, exc_value, traceback):
    duration = time.time() - self.starttime
    timestamp = datetime.datetime.now().strftime("%H:%M:%S")
    sys.stdout.write("{%s} %s [%.8f]\n" % (timestamp, self.label, duration))
    if self.sleep:
      time.sleep(self.sleep)