ktexter

#!/usr/bin/python2
"""

You can use ktexter to test /dev/ktext under different
readers/writers scenarios to simulate the starvation.
This script expects that /dev/ktext is available, so
please load the module first.
Moreover, make sure the user has read/write permissions
on the device (at this time, only root is allowed to
execute this script, see os.getuid()).
Writers are printed in purple.
Readers are printed in teal.

Example scenario:
       One reader argv[1]
       One writer argv[2]
       Reader frequency: 1Hz argv[3]
       Writer frequency: 1Hz argv[4]

       ktexter 1 1 1 1

   Optional arguments:
       --die=<seconds>: let the game end after <seconds>
       --rsleep=<seconds float>: make readers sleep for
           at most <seconds float>
       --wsleep=<seconds float>: make writers sleep for
           at most <seconds float>
       --sleep-randomize: randomize the sleep time a bit

XXX: ktexter 20 5 2 5 --wsleep=0.2

"""

import sys
import signal
import os
import threading
import time
import random
import errno
random.seed()

KILL_ALL = False

class GenericTask(threading.Thread):

    def __init__(self, callback, sleep_time):
        threading.Thread.__init__(self)
        self._sleep_time = sleep_time
        self._callback = callback

    def run(self):
        while True:
            sts = self._callback()
            if not sts:
                break
            time.sleep(self._sleep_time)

class Dispatcher(object):

    def __init__(self, n_readers, n_writers, readers_freq,
                 writers_freq, die_seconds, r_sleep_seconds,
                 w_sleep_seconds, sleep_randomize):
        self._n_readers = n_readers
        self._n_writers = n_writers
        self._readers_count = 0
        self._writers_count = 0
        self._readers_freq = readers_freq
        self._writers_freq = writers_freq
        self._readers_t = float(1)/self._readers_freq
        self._writers_t = float(1)/self._writers_freq
        self._print_generic("Dispatcher init: readers: %d, writers: %d, "
                         "readers_freq: %d, writers_freq: %d, die_seconds: %s" % (
                self._n_readers, self._n_writers, self._readers_freq,
                self._writers_freq, die_seconds))
        self._readers_pool = []
        self._writers_pool = []
        self._phrases = ["one", "two", "three", "four", "five", "six",
                         "seven", "eight", "nine", "ten"]
        self._device_file = "/dev/ktext"
        self._die_time = None
        # Amount of seconds to sleep between open() and close() -- readers
        self._readers_sleep_seconds = r_sleep_seconds
        # Amount of seconds to sleep between open() and close() -- writers
        self._writers_sleep_seconds = w_sleep_seconds
        # Add some entropy
        self._sleep_randomize = sleep_randomize

    def _get_random_phrase(self):
        return self._phrases[random.randint(0, len(self._phrases)-1)]

    def _print_teal(self, text):
        sys.stdout.write("\x1b[36m" + text + "\x1b[39;49;00m\n")

    def _print_purple(self, text):
        sys.stdout.write("\x1b[35m" + text + "\x1b[39;49;00m\n")

    def _print_generic(self, text):
        sys.stdout.write(text + "\n")

    def _reader_callback(self):
        th_id = threading.current_thread().ident
        if KILL_ALL:
            self._print_teal("[%s] reader dying" % (th_id,))
            return False

        t_start = time.time()
        secs = 0.0
        if self._readers_sleep_seconds:
            secs = self._readers_sleep_seconds
            if self._sleep_randomize:
                secs += random.random() / 10

        self._print_teal("[%s] reader called, will sleep: %.4f" % (th_id, secs,))
        text = None
        try:
            with open(self._device_file, "r") as dev_f:
                if secs:
                    time.sleep(secs)
                text = dev_f.read()
        except IOError as err:
            if err.errno == errno.EWOULDBLOCK:
                # ktext might raise EWOULDBLOCK
                # of blocking until reader can
                # read.
                self._print_teal("-- [%s] !!! reader would "
                                   "block (starvation?)" % (th_id,))
                return True
            raise
        t_total = time.time() - t_start
        self._print_teal("[%s] read: %s, taken: %.4f" % (
                th_id, text, t_total,))
        if self._die_time is None:
            return True
        if time.time() > self._die_time:
            self._print_teal("[%s] reader dying" % (th_id,))
            # goodbye
            return False
        if KILL_ALL:
            self._print_teal("[%s] reader dying" % (th_id,))
            return False
        return True

    def _writer_callback(self):
        th_id = threading.current_thread().ident
        if KILL_ALL:
            self._print_purple("[%s] writer dying" % (th_id,))
            return False

        t_start = time.time()
        secs = 0.0
        if self._writers_sleep_seconds:
            secs = self._writers_sleep_seconds
            if self._sleep_randomize:
                secs += random.random() / 10


        self._print_purple("[%s] writer called, will sleep: %.4f" % (th_id, secs,))
        text = self._get_random_phrase()
        try:
            with open(self._device_file, "w") as dev_f:
                if secs:
                    time.sleep(secs)
                dev_f.write(text)
        except IOError as err:
            if err.errno == errno.EWOULDBLOCK:
                # ktext might raise EWOULDBLOCK
                # of blocking until writer can
                # write.
                self._print_purple("-- [%s] !!! writer would "
                                   "block (starvation?)" % (th_id,))
                return True
            raise
        t_total = time.time() - t_start
        self._print_purple("[%s] write: %s, taken: %.4f" % (
                th_id, text, t_total,))
        if self._die_time is None:
            return True
        if time.time() > self._die_time:
            # goodbye
            self._print_purple("[%s] writer dying" % (th_id,))
            return False
        if KILL_ALL:
            self._print_purple("[%s] writer dying" % (th_id,))
            return False
        return True

    def run(self):
        for i in range(self._n_readers):
            th = GenericTask(self._reader_callback, self._readers_t)
            self._readers_pool.append(th)
        for i in range(self._n_writers):
            th = GenericTask(self._writer_callback, self._writers_t)
            self._writers_pool.append(th)

        # randomize with butterfly effect
        for th in self._readers_pool:
            self._print_generic("starting reader: %s" % (th,))
            th.start()
            rand_sleep = float(1) / random.randint(2, 10)
            time.sleep(rand_sleep)

        for th in self._writers_pool:
            self._print_generic("starting writer: %s" % (th,))
            th.start()
            rand_sleep = float(1) / random.randint(2, 10)
            time.sleep(rand_sleep)

        if die_seconds is not None:
            self._die_time = time.time() + die_seconds

        while True:
            alive = [x for x in self._readers_pool + \
                         self._writers_pool if x.is_alive()]
            if not alive:
                return
            # sleep 5 seconds and check again
            time.sleep(5)

def sigterm_handler(signum, frame):
    global KILL_ALL
    KILL_ALL = True
    sys.stderr.write("sigterm handler, be patient!\n")

if __name__ == "__main__":
    if os.getuid() != 0:
        sys.stderr.write("run as root\n")
        raise SystemExit(1)

    # Can't use argparse because it's not available in <Py2.7
    # And optparse sucks!
    args = sys.argv[1:]

    sleep_randomize = False
    if "--sleep-randomize" in args:
        sleep_randomize = True
        args.remove("--sleep-randomize")

    die_opts = [x for x in args if x.startswith("--die=")]
    for die_opt in die_opts:
        args.remove(die_opt)
    die_seconds = None
    if die_opts:
        die_opt = die_opts[-1]
        try:
            die_seconds = int(die_opt.split("=")[-1])
        except (IndexError, ValueError):
            sys.stderr.write("invalid --die= option\n")
            raise SystemExit(1)

    sleep_opts = [x for x in args if x.startswith("--rsleep=")]
    for sleep_opt in sleep_opts:
        args.remove(sleep_opt)
    readers_sleep_seconds = None
    if sleep_opts:
        sleep_opt = sleep_opts[-1]
        try:
            readers_sleep_seconds = float(sleep_opt.split("=")[-1])
        except (IndexError, ValueError):
            sys.stderr.write("invalid --rsleep= option\n")

    sleep_opts = [x for x in args if x.startswith("--wsleep=")]
    for sleep_opt in sleep_opts:
        args.remove(sleep_opt)
    writers_sleep_seconds = None
    if sleep_opts:
        sleep_opt = sleep_opts[-1]
        try:
            writers_sleep_seconds = float(sleep_opt.split("=")[-1])
        except (IndexError, ValueError):
            sys.stderr.write("invalid --wsleep= option\n")

    if not args or "--help" in args or "-h" in args:
        sys.stdout.write("%s <number of readers> <number of writers> "
                         "<readers call frequency in Hz> "
                         "<writers call frequency in Hz> "
                         "[--die=<seconds>] [--rsleep=<seconds float>] "
                         "[--wsleep=<seconds float>]" % (
                sys.argv[0],))
        if not args:
            raise SystemExit(1)
        raise SystemExit(0)

    if len(args) != 4:
        sys.stderr.write("invalid arguments, see --help\n")
        raise SystemExit(1)

    n_readers, n_writers, readers_freq, writers_freq = args

    try:
        n_readers = int(n_readers)
    except ValueError:
        sys.stderr.write("invalid numbers of readers\n")
        raise SystemExit(1)

    try:
        n_writers = int(n_writers)
    except ValueError:
        sys.stderr.write("invalid numbers of writers\n")
        raise SystemExit(1)

    try:
        readers_freq = int(readers_freq)
    except ValueError:
        sys.stderr.write("invalid readers frequency\n")
        raise SystemExit(1)

    try:
        writers_freq = int(writers_freq)
    except ValueError:
        sys.stderr.write("invalid writers frequency\n")
        raise SystemExit(1)

    if (n_readers < 1) and (n_writers < 1):
        sys.stderr.write("so, no readers and no writers?\n")
        raise SystemExit(1)

    signal.signal(signal.SIGALRM, sigterm_handler)

    disp = Dispatcher(n_readers, n_writers, readers_freq,
                      writers_freq, die_seconds,
                      readers_sleep_seconds,
                      writers_sleep_seconds,
                      sleep_randomize)

    try:
        disp.run() # blocks
    except KeyboardInterrupt:
        KILL_ALL = True
        # now pray
        sys.stderr.write("\n !!! waiting for all the threads to quit !!! \n\n")
        raise SystemExit(1)
    raise SystemExit(0)