from multiprocessing import Event, Process, cpu_count
from time import clock, sleep
from timeit import repeat


# Used by example() process target
def child(e):
    print('started')
    # This builds up sleeping_count & woken_count semaphore values in the
    # underlying Condition. Obviously in real-world usage wait might be
    # called many times with a small-ish timeout over the course of a
    # process' running time and the outer loop might be 'while not
    # e.is_set()'. This is just to demonstrate what affects both
    # both Condition.wait() and Event.wait().
    while not e.wait(0):
        pass
    print('finished')


# Illustrates Condition's semaphore counter build-up
def example():
    try:
        count = cpu_count()
    except NotImplementedError:
        count = 2

    event = Event()
    procs = [Process(target=child, args=(event,)) for _ in range(count)]
    for proc in procs:
        proc.start()

    # Allow for cummulative Event.wait() calls (with timeout) to build up via
    # child processes.
    sleep(10)

    print('stopping')
    start = clock()
    # This takes a while to decrement semaphores sleeping_count and woken_count
    # via loop in Condition.notify_all() which iterates for as many times as
    # wait() has been called:
    # https://github.com/python/cpython/blob/3.4/Lib/multiprocessing/synchronize.py#L294
    event.set()
    stop = clock()
    print('notified in %f (clock)' % (stop - start))

    for proc in procs:
        proc.join()


# For comparing performance before & after proposed change
def timing():
    et = repeat(setup='from multiprocessing import Event; e = Event()',
                stmt=('for _ in range(10000):\n'
                      '    e.wait(0)\n'
                      'e.set()'),
                number=1, repeat=3)

    print('Event timeit average: %.3f' % (sum(et) / len(et)))

    ct = repeat(setup='from multiprocessing import Condition; c = Condition()',
                stmt=('with c:\n'
                      '    for _ in range(10000):\n'
                      '        c.wait(0)\n'
                      '    c.notify_all()'),
                number=1, repeat=3)

    print('Condition timeit average: %.3f' % (sum(ct) / len(ct)))


if __name__ == "__main__":
    # from multiprocessing import set_start_method
    # set_start_method('forkserver')
    example()
    timing()