#!/usr/bin/python3
#+
# inotipy example: try creating and deleting lots of files in a temporary
# directory, and see if all the events get correctly reported. Invoke
# this script as follows:
#
#    stress_test --lifetime=«lifetime» [--maxrate=«maxrate»] \
#        [--processes=«nr_processes»] --totalfiles=«totalfiles» [--verbose]
#
# where «lifetime» is how many seconds each file should exist before
# being deleted (fractions allowed), «totalfiles» is how many files in
# total to create, «maxrate» is the maximum number of files to create
# per second (if omitted, they are created as fast as possible), and
# «nr_processes» is how many concurrent child processes to run to perform
# the file creations and deletions (if omitted, the default is 1).
#
# At the end of the run, the number of creation, deletion and overflow
# events are reported: ideally the first two should equal «totalfiles»
# and the last one should be zero. But if things happen too fast,
# creation/deletion events will be missed and overflows will be reported
# instead. The best way to create such an overload condition is to
# run multiple child processes.
#
# For example, an invocation like this
#
#     stress_test --lifetime=0.1 --totalfiles=100000 --processes=10
#
# runs for less than 15 seconds, but puts sufficient load on my system
# to lead to overflows during the run.
#
# Copyright 2018 by Lawrence D'Oliveiro <ldo@geek-central.gen.nz>. This
# script is licensed CC0
# <https://creativecommons.org/publicdomain/zero/1.0/>; do with it
# what you will.
#-

import sys
import os
import math
import time
import ctypes as ct
import multiprocessing
import threading
import tempfile
#import shutil
import asyncio
import getopt
import inotify
from inotify import \
    IN, \
    EVENT_BIT

#+
# Useful stuff
#-

def call_blocking_async(func, funcargs = (), timeout = None, abort = None, loop = None) :
    "invokes func on a separate temporary thread and returns a Future that" \
    " waits either for it to complete (and returns its result) or the timeout" \
    " (in seconds) to elapse, in which case it raises a TimeoutError exception" \
    " (if not specified, the timeout is infinite). This allows easy invocation" \
    " of blocking I/O functions in an asyncio-compatible fashion. But note that" \
    " the operation cannot be cancelled if the timeout elapses; instead, you can" \
    " specify an abort callback which will be invoked with whatever result is" \
    " eventually returned from func."

    if loop == None :
        loop = inotify.get_event_loop()
    #end if

    awaiting = loop.create_future()
    timeout_task = None

    async def func_done(result) :
        if not awaiting.done() :
            awaiting.set_result(result)
            if timeout_task != None :
                timeout_task.cancel()
            #end if
        else :
            if abort != None :
                abort(result)
            #end if
        #end if
    #end func_done

    def do_func_timedout() :
        if not awaiting.done() :
            awaiting.set_exception(TimeoutError())
            # Python doesn’t give me any (easy) way to cancel the thread running the
            # do_func() call, so just let it run to completion, whereupon func_done()
            # will get rid of the result. Even if I could delete the thread, can I be sure
            # that would clean up memory and OS/library resources properly?
        #end if
    #end do_func_timedout

    def do_func() :
        # makes the blocking call on a separate thread.
        result = func(*funcargs)
        # A Future is not itself threadsafe, but I can thread-safely
        # create a coroutine on the main thread to set it.
        asyncio.run_coroutine_threadsafe(func_done(result), loop)
    #end do_func

#begin call_blocking_async
    subthread = threading.Thread(target = do_func)
    subthread.start()
    if timeout != None :
        timeout_task = loop.call_later(timeout, do_func_timedout)
    #end if
    return \
        awaiting
#end call_blocking_async

#+
# Mainline
#-

class Counter(ct.Structure) :
    "shared-memory structure for communicating with child processes."
    _fields_ = \
        [
          # from parent to child:
            ("step", ct.c_uint), # step between filename indexes
            ("total", ct.c_uint), # how many files to create & delete
            ("create_index", ct.c_uint), # starting filename index
          # from child to parent:
            ("created", ct.c_uint), # how many created so far
            ("deleted", ct.c_uint), # how many deleted so far
        ]
#end Counter

loop = inotify.get_event_loop()

queue_factor = 2
progress_interval = 2
verbose = False
nr_processes = 1
totalfiles = None
min_queue = None
max_queue = None
max_rate = None
lifetime = None
opts, args = getopt.getopt \
  (
    sys.argv[1:],
    "",
    ["lifetime=", "maxqueue=", "maxrate=", "minqueue=", "processes=", "totalfiles=", "verbose"]
  )
for keyword, value in opts :
    if keyword == "--lifetime" :
        lifetime = float(value)
        assert lifetime > 0
    elif keyword == "--maxqueue" :
        max_queue = int(value)
        assert max_queue > 0
    elif keyword == "--maxrate" :
        max_rate = float(value)
        assert max_rate > 0
    elif keyword == "--minqueue" :
        min_queue = int(value)
        assert min_queue > 0
    elif keyword == "--processes" :
        nr_processes = int(value)
        assert nr_processes > 0
    elif keyword == "--totalfiles" :
        totalfiles = int(value)
        assert totalfiles > 0
    elif keyword == "--verbose" :
        verbose = True
    #end if
#end for
if len(args) != 0 :
    raise getopt.GetoptError("not expecting any args")
#end if
if lifetime == None :
    raise getopt.GetoptError("forgot to specify --lifetime=«lifetime»")
#end if
if totalfiles == None :
    raise getopt.GetoptError("forgot to specify --totalfiles=«totalfiles»")
#end if
nr_digits = math.ceil(math.log10(totalfiles + 1))
if max_queue == None :
    if min_queue == None :
        min_queue = 10
    #end if
    max_queue = min_queue * queue_factor
elif min_queue == None :
    min_queue = max(round(max_queue / queue_factor), 1)
elif max_queue < min_queue :
    raise getopt.GetoptError("--max-queue value must not be less than --min-queue value")
#end if

tempdir = tempfile.mkdtemp(prefix = "inotipy-stress-test-")

def create_file(filename) :
    open(os.path.join(tempdir, filename), "wb").close() # just create empty file
#end create_file

def delete_file(filename) :
    os.unlink(os.path.join(tempdir, filename))
#end delete_file

counters = multiprocessing.Array(Counter, nr_processes, lock = False)
triggers = list(multiprocessing.Event() for i in range(nr_processes))

def subprocess_main(index) :
    "mainline for child process."

    if max_rate != None :
        min_interval = nr_processes / max_rate # minimum interval between file creations
        last_create = None
    else :
        min_interval = None
    #end if

    async def subprocess_task() :
        "main task that runs in child process event loop."
        counter = counters[index]
        in_progress = []
        to_delete = []
        while True :
            awaiting = None
            now = time.time()
            while True :
                # pop in_progress entries that have completed
                if len(in_progress) == 0 :
                    break
                awaiting, filename, creating = in_progress[0]
                if not awaiting.done() :
                    break
                in_progress.pop(0)
                if creating :
                    to_delete.append((filename, now + lifetime))
                #end if
            #end while
            if (
                    len(in_progress) < min_queue
                and
                    (
                        counter.created < counter.total
                    or
                        counter.deleted < counter.total
                    or
                        len(to_delete) != 0
                    )
            ) :
                # can start more operations
                now = time.time()
                while \
                  (
                        len(to_delete) != 0
                    and
                        to_delete[0][1] <= now
                    and
                        len(in_progress) < max_queue
                  ) :
                    # can delete another file
                    filename = to_delete.pop(0)[0]
                    in_progress.append \
                      (
                        (
                            call_blocking_async
                              (
                                func = delete_file,
                                funcargs = (filename,),
                                loop = loop
                              ),
                            filename,
                            False,
                        )
                      )
                    counter.deleted += 1
                #end while
                while True :
                    if counter.created == counter.total or len(in_progress) >= max_queue :
                        break
                    if (
                            min_interval != None
                        and
                            last_create != None
                        and
                            now - last_create < min_interval
                    ) :
                        break
                    # can create another file
                    counter.created += 1
                    filename = "%%0%dd.jnk" % nr_digits % counter.create_index
                    counter.create_index += counter.step
                    in_progress.append \
                      (
                        (
                            call_blocking_async
                              (
                                func = create_file,
                                funcargs = (filename,),
                                loop = loop
                              ),
                            filename,
                            True,
                        )
                      )
                    if min_interval != None :
                        last_create = now
                    #end if
                #end while
            #end if
            awaiting = []
            if len(in_progress) != 0 :
                awaiting.append(in_progress[0][0])
            #end if
            if len(to_delete) != 0 :
                awaiting.append(loop.create_task(asyncio.sleep(max(to_delete[0][1] - time.time(), 0))))
            #end if
            if counter.created < counter.total and min_interval != None and last_create != None :
                awaiting.append(loop.create_task(asyncio.sleep(max(last_create + min_interval - time.time(), 0))))
            #end if
            if len(awaiting) != 0 :
                await asyncio.wait(awaiting, return_when = asyncio.FIRST_COMPLETED)
            else :
                assert counter.created == counter.total and counter.deleted == counter.total
                break # all done
            #end if
        #end while
    #end subprocess_task

#begin subprocess_main
    sys.stdout.write("child %d/%d starting\n" % (index + 1, nr_processes)) # debug
    triggers[index].wait()
    sys.stdout.write("child %d/%d is go\n" % (index + 1, nr_processes)) # debug
    loop.run_until_complete(subprocess_task())
    sys.stdout.write("child %d/%d finishing\n" % (index + 1, nr_processes)) # debug
#end subprocess_main

# Note that I avoid creating coroutines on the various process event loops until
# children have been forked. Not sure how important this is, but just to be safe.

def init_subprocesses() :
    "creates all the subprocesses that do the actual file creations and deletions."
    global processes
    per_process = totalfiles // nr_processes
    per_process = (per_process + totalfiles % nr_processes, per_process)
      # give remainder to first process
    for i in range(nr_processes) :
        counter = counters[i]
        counter.step = nr_processes
        counter.total = per_process[i != 0]
        counter.create_index = i + 1
        counter.created = 0
        counter.deleted = 0
    #end for
    processes = list \
      (
        multiprocessing.Process(target = subprocess_main, args = (i,))
        for i in range(nr_processes)
      )
    for proc in processes :
        proc.start() # start your engines ...
    #end for
#end init_subprocesses

init_subprocesses()

async def creator() :
    "gives a running report of all the file creation and deletion counters" \
    " while waiting for the subprocesses to terminate."

    async def progress() :
        "gives a running report on how many file creations and deletions" \
        " have actually been done."
        prev_created = None
        prev_deleted = None
        prev_time = None
        while True :
            now = time.time()
            created = 0
            deleted = 0
            total = 0
            for elt in counters :
                created += elt.created
                deleted += elt.deleted
                total += elt.total
            #end for
            sys.stdout.write("activity cre/del/total: %d/%d/%d" % (created, deleted, total))
            if prev_created != None :
                interval = now - prev_time
                sys.stdout.write \
                  (
                        " cre %d/s, del %d/s"
                    %
                        (
                            round((created - prev_created) / interval),
                            round((deleted - prev_deleted) / interval),
                        )
                  )
            #end if
            sys.stdout.write("\n")
            sys.stdout.flush()
            prev_created = created
            prev_deleted = deleted
            prev_time = now
            await asyncio.sleep(progress_interval)
        #end while
    #end progress

#begin creator
    processes_done = list \
      (
        call_blocking_async
          (
            func = lambda proc : proc.join(),
            funcargs = (proc,),
            loop = loop
          )
        for proc in processes
      )
    progress_task = loop.create_task(progress())
    await asyncio.wait(processes_done, return_when = asyncio.ALL_COMPLETED)
    progress_task.cancel()
    sys.stdout.write("creator done\n") # debug
#end creator

async def monitor() :
    "triggers the subprocesses to actually start file operations and" \
    " does the processing of resulting inotify events."

    created = 0
    deleted = 0
    overflowed = 0

    async def progress() :
        "gives a running report on how many file creation and deletion" \
        " notifications have been received."
        prev_created = None
        prev_deleted = None
        prev_overflowed = None
        prev_time = None
        while True :
            now = time.time()
            sys.stdout.write \
              (
                    "notified cre/del/ovf/total: %d/%d/%d/%d"
                %
                    (created, deleted, overflowed, totalfiles)
              )
            if prev_created != None :
                interval = now - prev_time
                sys.stdout.write \
                  (
                        " cre %d/s, del %d/s, ovf %d/s"
                    %
                        (
                            round((created - prev_created) / interval),
                            round((deleted - prev_deleted) / interval),
                            round((overflowed - prev_overflowed) / interval),
                        )
                  )
            #end if
            sys.stdout.write("\n")
            sys.stdout.flush()
            prev_created = created
            prev_deleted = deleted
            prev_overflowed = overflowed
            prev_time = now
            await asyncio.sleep(progress_interval)
        #end while
    #end progress

#begin monitor
    watcher = inotify.Watcher.create()
    watcher.watch(tempdir, IN.ALL_EVENTS)
    sys.stdout.write("monitor starting\n") # debug
    for trig in triggers :
        trig.set() # ... and they’re off
    #end for
    sys.stdout.write("parent told children to go\n") # debug
    progress_task = loop.create_task(progress())
    start_time = time.time()
    while True :
        event = await watcher.get(2)
        if event != None :
            if verbose or event.mask & EVENT_BIT.Q_OVERFLOW.mask != 0 :
                sys.stdout.write("Got event: %s\n" % repr(event))
            #end if
            if event.mask & EVENT_BIT.CREATE.mask != 0 :
                created += 1
            #end if
            if event.mask & EVENT_BIT.DELETE.mask != 0 :
                deleted += 1
            #end if
            if event.mask & EVENT_BIT.Q_OVERFLOW.mask != 0 :
                overflowed += 1
            #end if
        else :
            if creator_task.done() :
                break
            sys.stdout.write("No event.\n")
        #end if
    #end while
    progress_task.cancel()
    sys.stdout.write("time taken: %.3fs\n" % (time.time() - start_time))
    sys.stdout.write \
      (
            "events received: created: %d, deleted: %d, overflowed: %d\n"
        %
            (created, deleted, overflowed)
      )
#end monitor

creator_task = loop.create_task(creator())
loop.run_until_complete(monitor())

#shutil.rmtree(tempdir)
os.rmdir(tempdir) # should be empty
