contexts.py

# Copyright 2016 Quora, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asynq
import asyncio

from .asynq_to_async import is_asyncio_mode
from . import debug
from ._debug import options as _debug_options


ASYNCIO_CONTEXT_FIELD = "_asynq_contexts"
ASYNCIO_CONTEXT_ACTIVE_FIELD = "_asynq_contexts_active"


class NonAsyncContext(object):
    """Indicates that context can't contain yield statements.

    It means while a NonAsyncContext is active, async tasks cannot
    yield the control back to scheduler.

    Contexts should subclass this class if they want it throw an AssertionError
    if async tasks end up yielding within them.

    Note: Remember to call super inside your __enter__/__exit__.

    """

    def __enter__(self):
        if is_asyncio_mode():
            enter_context_asyncio(self)
        else:
            self._active_task = enter_context(self)

    def __exit__(self, typ, val, tb):
        if is_asyncio_mode():
            leave_context_asyncio(self)
        else:
            leave_context(self, self._active_task)

    def pause(self):
        assert False, "Task %s cannot yield while %s is active" % (
            self._active_task,
            self,
        )

    def resume(self):
        assert False, "Task %s cannot yield while %s is active" % (
            self._active_task,
            self,
        )


def enter_context(context):
    # perf optimization: inline get_active_task
    active_task = asynq.scheduler._state.current.active_task
    if active_task is not None:
        active_task._enter_context(context)
    return active_task


def leave_context(context, active_task):
    if active_task is not None:
        active_task._leave_context(context)


def enter_context_asyncio(context):
    if _debug_options.DUMP_CONTEXTS:
        debug.write("@async: +context: %s" % debug.str(context))

    # since we are in asyncio mode, there is an active task
    task = asyncio.current_task()

    if hasattr(task, ASYNCIO_CONTEXT_FIELD):
        getattr(task, ASYNCIO_CONTEXT_FIELD)[id(context)] = context
    else:
        setattr(task, ASYNCIO_CONTEXT_FIELD, {id(context): context})


def leave_context_asyncio(context):
    if _debug_options.DUMP_CONTEXTS:
        debug.write("@async: -context: %s" % debug.str(context))

    task = asyncio.current_task()
    getattr(task, ASYNCIO_CONTEXT_FIELD, {}).pop(id(context), None)  # type: ignore


def pause_contexts_asyncio(task):
    if getattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False):
        setattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False)
        for ctx in reversed(list(getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values())):
            ctx.pause()


def resume_contexts_asyncio(task):
    if not getattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, True):
        setattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, True)
        for ctx in getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values():
            ctx.resume()


class AsyncContext(object):
    """Base class for contexts that should pause and resume during an async's function execution.

    Your context should subclass this class and implement pause and resume (at least).

    That would make the context pause and resume each time the execution of the async function
    within this context is paused and resumed.

    Additionally, you can also subclass __enter__ and __exit__ if you want to customize its
    behaviour. Remember to call super in that case.

    NOTE: __enter__/__exit__ methods automatically call resume/pause so the overridden
    __enter__/__exit__ methods shouldn't do that explicitly.

    """

    def __enter__(self):
        if is_asyncio_mode():
            enter_context_asyncio(self)
        else:
            self._active_task = enter_context(self)

        self.resume()
        return self

    def __exit__(self, ty, value, tb):
        if is_asyncio_mode():
            leave_context_asyncio(self)
            self.pause()
        else:
            leave_context(self, self._active_task)
            self.pause()
            del self._active_task

    def resume(self):
        raise NotImplementedError()

    def pause(self):
        raise NotImplementedError()