def test_json_marshalling_subscribers() -> None:
subscriptions = [
Subscription("a", True, timedelta(seconds=1)),
Subscription("b", True, timedelta(minutes=1)),
Subscription("c", True, timedelta(hours=1)),
Subscription("d", False, timedelta(days=1)),
Subscription("e", False, timedelta(weeks=1)),
]
roundtrip(Subscriber.from_list(SubscriberId("foo"), []))
roundtrip(Subscriber.from_list(SubscriberId("foo"), subscriptions))
async def test_subscribe(handler: SubscriptionHandler, in_mem_db: SubscriberDb) -> None:
# register first time
result = await handler.add_subscription(SubscriberId("foo"), "event_bla", True, timedelta(seconds=3))
assert len(result.subscriptions) == 1
assert result.subscriptions["event_bla"].message_type == "event_bla"
# should be persisted in database as well
assert len((await in_mem_db.get("foo")).subscriptions) == 1 # type: ignore
# register again is ignored
result = await handler.add_subscription(SubscriberId("foo"), "event_bla", True, timedelta(seconds=3))
assert len(result.subscriptions) == 1
assert result.subscriptions["event_bla"].message_type == "event_bla"
# should be persisted in database as well
assert len((await in_mem_db.get("foo")).subscriptions) == 1 # type: ignore
async def update_subscriber(self, request: Request) -> StreamResponse:
subscriber_id = SubscriberId(request.match_info["subscriber_id"])
body = await self.json_from_request(request)
subscriptions = from_js(body, List[Subscription])
sub = await self.subscription_handler.update_subscriptions(
subscriber_id, subscriptions)
return await single_result(request, to_json(sub))
async def test_wait_for_running_job(
task_handler: TaskHandlerService, test_workflow: Workflow, all_events: List[Message]
) -> None:
test_workflow.on_surpass = TaskSurpassBehaviour.Wait
task_handler.task_descriptions = [test_workflow]
# subscribe as collect handler - the workflow will need to wait for this handler
sub = await task_handler.subscription_handler.add_subscription(
SubscriberId("sub_1"), "collect", True, timedelta(seconds=30)
)
await task_handler.handle_event(Event("start me up"))
# check, that the workflow has started
running_before = await task_handler.running_tasks()
assert len(running_before) == 1
act: Action = await wait_for_message(all_events, "collect", Action)
# pull the same trigger: the workflow can not be started, since there is already one in progress -> wait
await task_handler.handle_event(Event("start me up"))
# report success of the only subscriber
await task_handler.handle_action_done(ActionDone("collect", act.task_id, act.step_name, sub.id, dict(act.data)))
# check overdue tasks: wipe finished tasks and eventually start waiting tasks
await task_handler.check_overdue_tasks()
# check, that the workflow has started
running_after = await task_handler.running_tasks()
assert len(running_after) == 1
t_before, t_after = running_before[0], running_after[0]
assert t_before.descriptor.id == t_after.descriptor.id and t_before.id != t_after.id
async def add_subscription(self, request: Request) -> StreamResponse:
subscriber_id = SubscriberId(request.match_info["subscriber_id"])
event_type = request.match_info["event_type"]
timeout = timedelta(seconds=int(request.query.get("timeout", "60")))
wait_for_completion = request.query.get("wait_for_completion",
"true").lower() != "false"
sub = await self.subscription_handler.add_subscription(
subscriber_id, event_type, wait_for_completion, timeout)
return await single_result(request, to_js(sub))
def workflow_instance(
test_workflow: Workflow,
) -> Tuple[RunningTask, Subscriber, Subscriber, Dict[str, List[Subscriber]]]:
td = timedelta(seconds=100)
sub1 = Subscription("start_collect", True, td)
sub2 = Subscription("collect", True, td)
sub3 = Subscription("collect_done", True, td)
s1 = Subscriber.from_list(SubscriberId("s1"), [sub1, sub2, sub3])
s2 = Subscriber.from_list(SubscriberId("s2"), [sub2, sub3])
subscriptions = {"start_collect": [s1], "collect": [s1, s2], "collect_done": [s1, s2]}
w, _ = RunningTask.empty(test_workflow, lambda: subscriptions)
w.received_messages = [
Action("start_collect", w.id, "start"),
ActionDone("start_collect", w.id, "start", s1.id),
ActionDone("start_collect", w.id, "start", s2.id),
Event("godot", {"a": 1, "b": 2}),
Action("collect", w.id, "collect"),
ActionDone("collect", w.id, "collect", s1.id),
]
w.move_to_next_state()
return w, s1, s2, subscriptions
def test_message_serialization() -> None:
task_id = TaskId("123")
subsctiber_id = SubscriberId("sub")
roundtrip(Event("test", {"a": "b", "c": 1, "d": "bla"}))
roundtrip(Action("test", task_id, "step_name"))
roundtrip(Action("test", task_id, "step_name", {"test": 1}))
roundtrip(ActionDone("test", task_id, "step_name", subsctiber_id))
roundtrip(
ActionDone("test", task_id, "step_name", subsctiber_id, {"test": 1}))
roundtrip(ActionError("test", task_id, "step_name", subsctiber_id, "oops"))
roundtrip(
ActionError("test", task_id, "step_name", subsctiber_id, "oops",
{"test": 23}))
async def __handle_events(self) -> None:
subscriber_id = SubscriberId("resotocore.config.update")
async with self.message_bus.subscribe(
subscriber_id, [CoreMessage.ConfigUpdated]) as events:
while True:
event = await events.get()
event_id = event.data.get("id")
if event_id in (ResotoCoreConfigId,
ResotoCoreCommandsConfigId):
log.info(
f"Core config was updated: {event_id} Restart to take effect."
)
# stop the process and rely on os to restart the service
self.exit_fn()
async def listen_to_message_bus() -> None:
async with self.message_bus.subscribe(SubscriberId("resotocore.task_handler")) as messages:
while True:
message = None
try:
message = await messages.get()
if isinstance(message, Event):
await self.handle_event(message)
elif isinstance(message, Action):
await self.handle_action(message)
elif isinstance(message, (ActionDone, ActionError)):
log.info(f"Ignore message via event bus: {message}")
except asyncio.CancelledError as ex:
# if we outer task is cancelled, give up
raise ex
except Exception as ex:
log.error(f"Could not handle event {message} - give up.", exc_info=ex)
async def test_unsubscribe(handler: SubscriptionHandler, in_mem_db: SubscriberDb) -> None:
# register first time
subscriber_id = SubscriberId("foo")
subs = [Subscription("event_bla"), Subscription("event_bar")]
result = await handler.update_subscriptions(subscriber_id, subs)
assert len(result.subscriptions) == 2
updated = await handler.remove_subscription(subscriber_id, "event_bla")
assert len(updated.subscriptions) == 1
# should be persisted in database as well
assert len((await in_mem_db.get(subscriber_id)).subscriptions) == 1 # type: ignore
# second time should be ignored
updated = await handler.remove_subscription(subscriber_id, "event_bla")
assert len(updated.subscriptions) == 1
# last subscription is removed
updated = await handler.remove_subscription(subscriber_id, "event_bar")
assert len(updated.subscriptions) == 0
# should be persisted in database as well
assert await in_mem_db.get(subscriber_id) is None
async def wait_for_subscriber() -> None:
subscriber_id = SubscriberId(f"resotocore.wait_for_actor_{uuid_str()}")
async with message_bus.subscribe(subscriber_id,
[CoreMessage.Connected]) as bus:
while True:
message = await bus.get()
maybe_workflow = workflow_if_actor(message)
if maybe_workflow:
log.info(
f"Subscribed actor for workflow {maybe_workflow.name} connected. "
f"Wait for {wait_after_connect} seconds.")
# wait before we start the workflow
await asyncio.sleep(wait_after_connect.total_seconds())
log.info(f"Start workflow {maybe_workflow.name}")
# start the workflow
await task_handler.start_task_by_descriptor_id(
maybe_workflow.id)
# exit the loop and destroy the listener
break
log.info("task done")
async def handle_subscribed(self, request: Request) -> StreamResponse:
subscriber_id = SubscriberId(request.match_info["subscriber_id"])
subscriber = await self.subscription_handler.get_subscriber(
subscriber_id)
if subscriber_id in self.message_bus.active_listener:
log.info(
f"There is already a listener for subscriber: {subscriber_id}. Reject."
)
return web.HTTPTooManyRequests(
text="Only one connection per subscriber is allowed!")
elif subscriber and subscriber.subscriptions:
pending = await self.workflow_handler.list_all_pending_actions_for(
subscriber)
return await self.listen_to_events(
request, subscriber_id, list(subscriber.subscriptions.keys()),
pending)
else:
return web.HTTPNotFound(
text=
f"No subscriber with this id: {subscriber_id} or no subscriptions"
)
async def get_subscriber(self, request: Request) -> StreamResponse:
subscriber_id = SubscriberId(request.match_info["subscriber_id"])
subscriber = await self.subscription_handler.get_subscriber(
subscriber_id)
return self.optional_json(subscriber,
f"No subscriber with id {subscriber_id}")
async def delete_subscriber(self, request: Request) -> StreamResponse:
subscriber_id = SubscriberId(request.match_info["subscriber_id"])
await self.subscription_handler.remove_subscriber(subscriber_id)
return web.HTTPNoContent()
async def wait_for(name: str, list: List[Message]) -> None:
async with message_bus.subscribe(SubscriberId("test"),
[name]) as events:
while True:
list.append(await events.get())
async def gather_events() -> None:
async with message_bus.subscribe(SubscriberId("test")) as event_queue:
while True:
events.append(await event_queue.get())
def subscribers() -> List[Subscriber]:
subs = [Subscription("foo", True) for _ in range(0, 10)]
return [
Subscriber.from_list(SubscriberId(str(a)), subs) for a in range(0, 10)
]
async def delete_subscription(self, request: Request) -> StreamResponse:
subscriber_id = SubscriberId(request.match_info["subscriber_id"])
event_type = request.match_info["event_type"]
sub = await self.subscription_handler.remove_subscription(
subscriber_id, event_type)
return await single_result(request, to_js(sub))
async def subscription_handler(message_bus: MessageBus) -> SubscriptionHandler:
in_mem = InMemoryDb[SubscriberId, Subscriber](Subscriber, lambda x: x.id)
result = SubscriptionHandler(in_mem, message_bus)
await result.add_subscription(SubscriberId("sub_1"), "test", True, timedelta(seconds=3))
return result
from resotocore.message_bus import MessageBus, Action
import logging
import asyncio
from asyncio import Task, Future
from typing import Optional
from contextlib import suppress
from datetime import timedelta
from resotocore.task.model import Subscriber
from resotocore.ids import SubscriberId
from resotocore.task.task_handler import TaskHandlerService
from resotocore.ids import TaskId
from resotocore.task.subscribers import SubscriptionHandler
log = logging.getLogger(__name__)
subscriber_id = SubscriberId("resotocore")
merge_outer_edges = "merge_outer_edges"
class MergeOuterEdgesHandler:
def __init__(
self,
message_bus: MessageBus,
subscription_handler: SubscriptionHandler,
task_handler_service: TaskHandlerService,
):
self.message_bus = message_bus
self.merge_outer_edges_listener: Optional[Task[None]] = None
self.subscription_handler = subscription_handler
self.subscriber: Optional[Subscriber] = None
self.task_handler_service = task_handler_service
async def handle_events(self, request: Request) -> StreamResponse:
show = request.query["show"].split(
",") if "show" in request.query else ["*"]
return await self.listen_to_events(request,
SubscriberId(str(uuid.uuid1())),
show)
async def test_recover_workflow(
running_task_db: RunningTaskDb,
job_db: JobDb,
message_bus: MessageBus,
event_sender: AnalyticsEventSender,
subscription_handler: SubscriptionHandler,
all_events: List[Message],
cli: CLI,
test_workflow: Workflow,
) -> None:
def handler() -> TaskHandlerService:
th = TaskHandlerService(
running_task_db,
job_db,
message_bus,
event_sender,
subscription_handler,
Scheduler(),
cli,
empty_config(),
)
th.task_descriptions = [test_workflow]
return th
await subscription_handler.add_subscription(SubscriberId("sub_1"), "start_collect", True, timedelta(seconds=30))
sub1 = await subscription_handler.add_subscription(SubscriberId("sub_1"), "collect", True, timedelta(seconds=30))
sub2 = await subscription_handler.add_subscription(SubscriberId("sub_2"), "collect", True, timedelta(seconds=30))
async with handler() as wf1:
# kick off a new workflow
await wf1.handle_event(Event("start me up"))
assert len(wf1.tasks) == 1
# expect a start_collect action message
a: Action = await wait_for_message(all_events, "start_collect", Action)
await wf1.handle_action_done(ActionDone(a.message_type, a.task_id, a.step_name, sub1.id, dict(a.data)))
# expect a collect action message
b: Action = await wait_for_message(all_events, "collect", Action)
await wf1.handle_action_done(ActionDone(b.message_type, b.task_id, b.step_name, sub1.id, dict(b.data)))
# subscriber 3 is also registering for collect
# since the collect phase is already started, it should not participate in this round
sub3 = await subscription_handler.add_subscription(SubscriberId("sub_3"), "collect", True, timedelta(seconds=30))
# simulate a restart, wf1 is stopped and wf2 needs to recover from database
async with handler() as wf2:
assert len(wf2.tasks) == 1
wfi = list(wf2.tasks.values())[0]
assert wfi.current_state.name == "act"
assert (await wf2.list_all_pending_actions_for(sub1)) == []
assert (await wf2.list_all_pending_actions_for(sub2)) == [Action("collect", wfi.id, "act", {})]
assert (await wf2.list_all_pending_actions_for(sub3)) == []
await wf2.handle_action_done(ActionDone("collect", wfi.id, "act", sub2.id, {}))
# expect an event workflow_end
await wait_for_message(all_events, "task_end", Event)
# all workflow instances are gone
assert len(wf2.tasks) == 0
# simulate a restart, wf3 should start from a clean slate, since all instances are done
async with handler() as wf3:
assert len(wf3.tasks) == 0
async def handler(in_mem_db: SubscriberDb) -> SubscriptionHandler:
result = SubscriptionHandler(in_mem_db, MessageBus())
await result.add_subscription(SubscriberId("sub_1"), "test", True, timedelta(seconds=3))
return result
async def test_get_subscriber(handler: SubscriptionHandler) -> None:
result = await handler.get_subscriber(SubscriberId("sub_1"))
assert result
assert result.id == "sub_1"