task_scheduler.py

from typing import TYPE_CHECKING, Awaitable, Callable, Dict, List, Optional, Set

import attr

from synapse.api.errors import SynapseError
from synapse.metrics.background_process_metrics import run_as_background_process
from synapse.types import JsonMapping, ScheduledTask, TaskState
from synapse.util.stringutils import random_string

if TYPE_CHECKING:
    from synapse.server import HomeServer


class TaskSchedulerHandler:
    SCHEDULING_INTERVAL_MS = 10 * 60 * 1000  # 10mn

    def __init__(self, hs: "HomeServer"):
        self.store = hs.get_datastores().main
        self.clock = hs.get_clock()
        self._is_master = hs.config.worker.worker_app is None
        self.running_tasks: Set[str] = set()
        self.actions: Dict[
            str, Callable[[ScheduledTask], Awaitable[Optional[ScheduledTask]]]
        ] = {}

        if self._is_master:
            self.clock.looping_call(
                run_as_background_process,
                TaskSchedulerHandler.SCHEDULING_INTERVAL_MS,
                "scheduled_tasks_loop",
                self._scheduled_tasks_loop,
            )

    def bind_action(
        self,
        fct: Callable[[ScheduledTask], Awaitable[Optional[ScheduledTask]]],
        action_name: str,
    ) -> None:
        self.actions[action_name] = fct

    async def schedule_task(
        self,
        action: str,
        *,
        resource_id: Optional[str] = None,
        timestamp: Optional[int] = None,
        params: Optional[JsonMapping] = None,
    ) -> str:
        if action not in self.actions:
            # TODO
            raise SynapseError(400, "Test")
        task_id = random_string(16)
        state = TaskState.SCHEDULED
        if timestamp is None or timestamp < self.clock.time_msec():
            state = TaskState.RUNNING
            timestamp = self.clock.time_msec()

        task = ScheduledTask(
            task_id,
            action,
            state,
            resource_id,
            timestamp,
            params,
            None,
        )
        await self.store.upsert_scheduled_task(task)
        return task_id

    async def update_task_state(
        self,
        task: ScheduledTask,
        # error: Optional[str],
    ) -> None:
        await self.store.upsert_scheduled_task(task)

    async def get_task(self, id: str) -> Optional[ScheduledTask]:
        return await self.store.get_scheduled_task(id)

    async def get_tasks(
        self, action: str, resource_id: Optional[str]
    ) -> List[ScheduledTask]:
        return await self.store.get_scheduled_tasks(action, resource_id)

    async def _scheduled_tasks_loop(self) -> None:
        for task in await self.store.get_scheduled_tasks():
            if task.id not in self.running_tasks:
                state = task.state
                if (
                    state == TaskState.SCHEDULED
                    and task.timestamp is not None
                    and task.timestamp < self.clock.time_msec()
                ):
                    state = TaskState.RUNNING

                if state == TaskState.RUNNING:
                    await self.store.upsert_scheduled_task(task)
                    self._run_task(task)

    def _run_task(self, task: ScheduledTask) -> None:
        if task.action in self.actions:
            fct = self.actions[task.action]

            async def wrapper() -> None:
                updated_task = await fct(task)
                if updated_task is None:
                    updated_task = attr.evolve(task, state=TaskState.COMPLETE)
                await self.update_task_state(updated_task)

            run_as_background_process(task.action, wrapper)
            self.running_tasks.add(task.id)