# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors
#
# This module is part of async and is released under
# the New BSD License: http://www.opensource.org/licenses/bsd-license.php
"""Pool testing"""
from .lib import (
        TestBase,
        py2
    )
from .task import (
        FixtureThreadTask,
        FixtureChannelThreadTask,
        FixtureFailureThreadTask,
        add_task_chain,        
        make_iterator_task
    )

from async.pool import (
        ThreadPool,
        PoolReader
    )
from async.thread import terminate_threads
from async.util import cpu_count

import time
import sys


class TestThreadPool(TestBase):

    max_threads = cpu_count()

    def _assert_single_task(self, p, async=False):
        """Performs testing in a synchronized environment"""
        sys.stderr.write("Threadpool: Starting single task (async = %i) with %i threads\n" % (async, p.size()))
        null_tasks = p.num_tasks()      # in case we had some before

        # add a simple task
        # it iterates n items
        ni = 1000
        assert ni % 2 == 0, "ni needs to be dividable by 2"
        assert ni % 4 == 0, "ni needs to be dividable by 4"

        make_task = lambda *args, **kwargs: make_iterator_task(ni, *args, **kwargs)

        task = make_task()

        assert p.num_tasks() == null_tasks
        rc = p.add_task(task)
        assert p.num_tasks() == 1 + null_tasks
        assert isinstance(rc, PoolReader)
        assert task._out_writer is not None

        # pull the result completely - we should get one task, which calls its
        # function once. In sync mode, the order matches
        items = rc.read()
        assert len(items) == ni
        task._assert(1, ni)
        if not async:
            assert items[0] == 0 and items[-1] == ni-1

        # as the task is done, it should have been removed - we have read everything
        assert task.is_done()
        del(rc)
        assert p.num_tasks() == null_tasks
        task = make_task()

        # pull individual items
        rc = p.add_task(task)
        assert p.num_tasks() == 1 + null_tasks
        st = time.time()
        for i in range(ni):
            items = rc.read(1)
            assert len(items) == 1

            # can't assert order in async mode
            if not async:
                assert i == items[0]
        # END for each item
        elapsed = time.time() - st
        sys.stderr.write("Threadpool: processed %i individual items, with %i threads, one at a time, in %f s ( %f items / s )\n" % (ni, p.size(), elapsed, ni / elapsed))

        # it couldn't yet notice that the input is depleted as we pulled exaclty
        # ni items - the next one would remove it. Instead, we delete our channel
        # which triggers orphan handling
        assert not task.is_done()
        assert p.num_tasks() == 1 + null_tasks
        del(rc)
        assert p.num_tasks() == null_tasks

        # test min count
        # if we query 1 item, it will prepare ni / 2
        task = make_task()
        task.min_count = ni / 2
        rc = p.add_task(task)
        items = rc.read(1)
        assert len(items) == 1 and items[0] == 0            # processes ni / 2
        items = rc.read(1)
        assert len(items) == 1 and items[0] == 1            # processes nothing
        # rest - it has ni/2 - 2 on the queue, and pulls ni-2
        # It wants too much, so the task realizes its done. The task
        # doesn't care about the items in its output channel
        nri = ni - 2
        items = rc.read(nri)
        assert len(items) == nri
        p.remove_task(task)
        assert p.num_tasks() == null_tasks
        task._assert(2, ni)                     # two chunks, ni calls

        # its already done, gives us no more, its still okay to use it though
        # as a task doesn't have to be in the graph to allow reading its produced
        # items
        # it can happen that a thread closes the channel just a tiny fraction of time
        # after we check this, so the test fails, although it is nearly closed.
        # When we start reading, we should wake up once it sends its signal
        # assert task.is_closed()
        assert len(rc.read()) == 0

        # test chunking
        # we always want 4 chunks, these could go to individual nodes
        task = make_task()
        task.min_count = ni // 2             # restore previous value
        task.max_chunksize = ni // 4         # 4 chunks
        rc = p.add_task(task)

        # must read a specific item count
        # count is still at ni / 2 - here we want more than that
        # 2 steps with n / 4 items, + 1 step with n/4 items to get + 2
        nri = ni // 2 + 2
        items = rc.read(nri)
        if py2:
            assert len(items) == nri
        # have n / 4 - 2 items on queue, want n / 4 in first chunk, cause 1 processing
        # ( 4 in total ). Still want n / 4 - 2 in second chunk, causing another processing
        nri = ni // 2 - 2
        items = rc.read(nri)
        assert len(items) == nri

        task._assert( 5, ni)

        # delete the handle first, causing the task to be removed and to be set
        # done. We check for the set-done state later. Depending on the timing,
        # The task is not yet set done when we are checking it because we were
        # scheduled in before the flag could be set.
        del(rc)
        assert task.is_done()
        assert p.num_tasks() == null_tasks  # depleted

        # but this only hits if we want too many items, if we want less, it could
        # still do too much - hence we set the min_count to the same number to enforce
        # at least ni / 4 items to be preocessed, no matter what we request
        task = make_task()
        task.min_count = None
        task.max_chunksize = ni / 4     # match previous setup
        rc = p.add_task(task)
        st = time.time()
        for i in range(ni):
            if async:
                assert len(rc.read(1)) == 1
            else:
                assert rc.read(1)[0] == i
            # END handle async mode
        # END pull individual items
        # too many processing counts ;)
        elapsed = time.time() - st
        sys.stderr.write("Threadpool: processed %i individual items in chunks of %i, with %i threads, one at a time, in %f s ( %f items / s)\n" % (ni, ni/4, p.size(), elapsed, ni / elapsed))

        task._assert(ni, ni)
        assert p.num_tasks() == 1 + null_tasks
        assert p.remove_task(task) is p     # del manually this time
        assert p.num_tasks() == null_tasks

        # now with we set the minimum count to reduce the number of processing counts
        task = make_task()
        task.min_count = ni / 4
        task.max_chunksize = ni / 4     # match previous setup
        rc = p.add_task(task)
        for i in range(ni):
            items = rc.read(1)
            assert len(items) == 1
            if not async:
                assert items[0] == i
        # END for each item
        task._assert(ni / task.min_count, ni)
        del(rc)
        assert p.num_tasks() == null_tasks

        # test failure
        # on failure, the processing stops and the task is finished, keeping
        # his error for later
        task = make_task()
        task.should_fail = True
        rc = p.add_task(task)
        assert len(rc.read()) == 0      # failure on first item

        assert isinstance(task.error(), AssertionError)
        assert task.is_done()           # on error, its marked done as well
        del(rc)
        if py2:
            assert p.num_tasks() == null_tasks

        # test failure after ni / 2 items
        # This makes sure it correctly closes the channel on failure to prevent blocking
        nri = ni/2
        task = make_task(FixtureFailureThreadTask, fail_after=ni/2)
        rc = p.add_task(task)
        assert len(rc.read()) == nri
        assert task.is_done()
        assert isinstance(task.error(), AssertionError)

        sys.stderr.write("done with everything\n")



    def _assert_async_dependent_tasks(self, pool):
        # includes failure in center task, 'recursive' orphan cleanup
        # This will also verify that the channel-close mechanism works
        # t1 -> t2 -> t3

        sys.stderr.write("Threadpool: starting async dependency test in %i threads\n" % pool.size())
        null_tasks = pool.num_tasks()
        ni = 1000
        count = 3
        aic = count + 2
        make_task = lambda *args, **kwargs: add_task_chain(pool, ni, count, *args, **kwargs)

        ts, rcs = make_task()
        assert len(ts) == aic
        assert len(rcs) == aic
        assert pool.num_tasks() == null_tasks + len(ts)

        # read(0)
        #########
        st = time.time()
        items = rcs[-1].read()
        elapsed = time.time() - st
        assert len(items) == ni
        del(rcs)
        if py2:
            assert pool.num_tasks() == 0        # tasks depleted, all done, no handles
        # wait a tiny moment - there could still be something unprocessed on the
        # queue, increasing the refcount
        time.sleep(0.15)
        assert sys.getrefcount(ts[-1]) == 2 # ts + call
        assert sys.getrefcount(ts[0]) == 2  # ts + call
        sys.stderr.write("Dependent Tasks: evaluated %i items of %i dependent in %f s ( %i items / s )\n" % (ni, aic, elapsed, ni / elapsed))


        # read(1)
        #########
        ts, rcs = make_task()
        st = time.time()
        for i in range(ni):
            items = rcs[-1].read(1)
            assert len(items) == 1
        # END for each item to pull
        elapsed_single = time.time() - st
        # another read yields nothing, its empty
        assert len(rcs[-1].read()) == 0
        sys.stderr.write("Dependent Tasks: evaluated %i items with read(1) of %i dependent in %f s ( %i items / s )\n" % (ni, aic, elapsed_single, ni / elapsed_single))


        # read with min-count size
        ###########################
        # must be faster, as it will read ni / 4 chunks
        # Its enough to set one task, as it will force all others in the chain
        # to min_size as well.
        ts, rcs = make_task()
        if py2:
            assert pool.num_tasks() == len(ts)
        nri = ni / 4
        ts[-1].min_count = nri
        st = time.time()
        for i in range(ni):
            items = rcs[-1].read(1)
            assert len(items) == 1
        # END for each item to read
        elapsed_minsize = time.time() - st
        # its empty
        assert len(rcs[-1].read()) == 0
        sys.stderr.write("Dependent Tasks: evaluated %i items with read(1), min_size=%i, of %i dependent in %f s ( %i items / s )\n" % (ni, nri, aic, elapsed_minsize, ni / elapsed_minsize))

        # it should have been a bit faster at least, and most of the time it is
        # Sometimes, its not, mainly because:
        # * The test tasks lock a lot, hence they slow down the system
        # * Each read will still trigger the pool to evaluate, causing some overhead
        #   even though there are enough items on the queue in that case. Keeping
        #   track of the scheduled items helped there, but it caused further inacceptable
        #   slowdown
        # assert elapsed_minsize < elapsed_single


        # read with failure
        ###################
        # it should recover and give at least fail_after items
        # t1 -> x -> t3
        fail_after = ni/2
        ts, rcs = make_task(fail_setup=[(0, fail_after)])
        items = rcs[-1].read()
        assert len(items) == fail_after


        # MULTI-POOL
        # If two pools are connected, this shold work as well.
        # The second one has just one more thread
        ts, rcs = make_task()

        # connect verifier channel as feeder of the second pool
        p2 = ThreadPool(0)      # don't spawn new threads, they have the tendency not to wake up on mutexes
        assert p2.size() == 0
        p2ts, p2rcs = add_task_chain(p2, ni, count, feeder_channel=rcs[-1], id_offset=count)
        assert p2ts[0] is None      # we have no feeder task
        assert rcs[-1].pool_ref()() is pool     # it didnt change the pool
        assert rcs[-1] is p2ts[1].reader()
        assert p2.num_tasks() == len(p2ts)-1    # first is None

        # reading from the last one will evaluate all pools correctly
        st = time.time()
        items = p2rcs[-1].read()
        elapsed = time.time() - st
        assert len(items) == ni

        sys.stderr.write("Dependent Tasks: evaluated 2 connected pools and %i items with read(0), of %i dependent tasks in %f s ( %i items / s )\n" % (ni, aic + aic-1, elapsed, ni / elapsed))


        # loose the handles of the second pool to allow others to go as well
        del(p2rcs); del(p2ts)
        assert p2.num_tasks() == 0

        # now we lost our old handles as well, and the tasks go away
        ts, rcs = make_task()
        if py2:
            assert pool.num_tasks() == len(ts)

        p2ts, p2rcs = add_task_chain(p2, ni, count, feeder_channel=rcs[-1], id_offset=count)
        assert p2.num_tasks() == len(p2ts) - 1

        # Test multi-read(1)
        reader = rcs[-1]
        st = time.time()
        for i in range(ni):
            items = reader.read(1)
            assert len(items) == 1
        # END for each item to get
        elapsed = time.time() - st
        del(reader)     # decrement refcount

        sys.stderr.write("Dependent Tasks: evaluated 2 connected pools and %i items with read(1), of %i dependent tasks in %f s ( %i items / s )\n" % (ni, aic + aic-1, elapsed, ni / elapsed))

        # another read is empty
        assert len(rcs[-1].read()) == 0

        # now that both are connected, I can drop my handle to the reader
        # without affecting the task-count, but whats more important:
        # They remove their tasks correctly once we drop our references in the
        # right order
        del(p2ts)
        assert p2rcs[0] is rcs[-1]
        del(p2rcs)
        assert p2.num_tasks() == 0
        del(p2)

        if py2:
            assert pool.num_tasks() == null_tasks + len(ts)


        del(ts)
        del(rcs)

        if py2:
            assert pool.num_tasks() == null_tasks


        # ASSERTION: We already tested that one pool behaves correctly when an error
        # occours - if two pools handle their ref-counts correctly, which they
        # do if we are here, then they should handle errors happening during
        # the task processing as expected as well. Hence we can safe this here



    @terminate_threads
    def test_base(self):
        p = ThreadPool()

        # default pools have no workers - and threading was removed entirely ... 
        assert p.size() == 0

        # SINGLE TASK SERIAL SYNC MODE
        ##############################
        # put a few unrelated tasks that we forget about - check ref counts and cleanup
        t1, t2 = FixtureThreadTask(iter(list()), "nothing1", None), FixtureThreadTask(iter(list()), "nothing2", None)
        urc1 = p.add_task(t1)
        urc2 = p.add_task(t2)
        assert p.num_tasks() == 2

        # test pool reader
        assert urc1.pool_ref()() is p
        assert urc1.task_ref()() is t1
        assert urc1.pool() == p
        assert urc1.task() == t1

        ## SINGLE TASK #################
        self._assert_single_task(p, False)
        if py2:
            assert p.num_tasks() == 2
        del(urc1)
        if py2:
            assert p.num_tasks() == 1

        p.remove_task(t2)
        if py2:
            assert p.num_tasks() == 0
        assert sys.getrefcount(t2) == 2

        t3 = FixtureChannelThreadTask(urc2, "channel", None)
        urc3 = p.add_task(t3)
        if py2:
            assert p.num_tasks() == 1
        del(urc3)
        if py2:
            assert p.num_tasks() == 0
        assert sys.getrefcount(t3) == 2


        # DEPENDENT TASKS SYNC MODE
        ###########################
        self._assert_async_dependent_tasks(p)
