Use asyncio for broadcasting

This commit is contained in:
Mouse Reeve 2022-11-10 15:41:56 -08:00
parent cbd14a69ea
commit ddcaf8e3b8
3 changed files with 58 additions and 27 deletions

View file

@ -1,14 +1,15 @@
""" activitypub model functionality """
import asyncio
from base64 import b64encode
from collections import namedtuple
from functools import reduce
import json
import operator
import logging
from typing import List
from uuid import uuid4
import requests
from requests.exceptions import RequestException
import aiohttp
from Crypto.PublicKey import RSA
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256
@ -136,7 +137,7 @@ class ActivitypubMixin:
queue=queue,
)
def get_recipients(self, software=None):
def get_recipients(self, software=None) -> List[str]:
"""figure out which inbox urls to post to"""
# first we have to figure out who should receive this activity
privacy = self.privacy if hasattr(self, "privacy") else "public"
@ -506,19 +507,31 @@ def unfurl_related_field(related_field, sort_field=None):
@app.task(queue=MEDIUM)
def broadcast_task(sender_id, activity, recipients):
def broadcast_task(sender_id: int, activity: str, recipients: List[str]):
"""the celery task for broadcast"""
user_model = apps.get_model("bookwyrm.User", require_ready=True)
sender = user_model.objects.get(id=sender_id)
for recipient in recipients:
try:
sign_and_send(sender, activity, recipient)
except RequestException:
pass
sender = user_model.objects.select_related("key_pair").get(id=sender_id)
asyncio.run(async_broadcast(recipients, sender, activity))
def sign_and_send(sender, data, destination):
"""crpyto whatever and http junk"""
async def async_broadcast(recipients: List[str], sender, data: str):
"""Send all the broadcasts simultaneously"""
timeout = aiohttp.ClientTimeout(total=10)
async with aiohttp.ClientSession(timeout=timeout) as session:
tasks = []
for recipient in recipients:
tasks.append(
asyncio.ensure_future(sign_and_send(session, sender, data, recipient))
)
results = await asyncio.gather(*tasks)
return results
async def sign_and_send(
session: aiohttp.ClientSession, sender, data: str, destination: str
):
"""Sign the messages and send them in an asynchronous bundle"""
now = http_date()
if not sender.key_pair.private_key:
@ -527,20 +540,25 @@ def sign_and_send(sender, data, destination):
digest = make_digest(data)
response = requests.post(
destination,
data=data,
headers={
"Date": now,
"Digest": digest,
"Signature": make_signature(sender, destination, now, digest),
"Content-Type": "application/activity+json; charset=utf-8",
"User-Agent": USER_AGENT,
},
)
if not response.ok:
response.raise_for_status()
return response
headers = {
"Date": now,
"Digest": digest,
"Signature": make_signature(sender, destination, now, digest),
"Content-Type": "application/activity+json; charset=utf-8",
"User-Agent": USER_AGENT,
}
try:
async with session.post(destination, data=data, headers=headers) as response:
if not response.ok:
logger.exception(
"Failed to send broadcast to %s: %s", destination, response.reason
)
return response
except asyncio.TimeoutError:
logger.info("Connection timed out for url: %s", destination)
except aiohttp.ClientError as err:
logger.exception(err)
# pylint: disable=unused-argument

View file

@ -27,6 +27,7 @@ from bookwyrm import models
class BaseActivity(TestCase):
"""the super class for model-linked activitypub dataclasses"""
# pylint: disable=invalid-name
def setUp(self):
"""we're probably going to re-use this so why copy/paste"""
with patch("bookwyrm.suggested_users.rerank_suggestions_task.delay"), patch(

View file

@ -12,6 +12,7 @@ from bookwyrm.models import base_model
from bookwyrm.models.activitypub_mixin import (
ActivitypubMixin,
ActivityMixin,
broadcast_task,
ObjectMixin,
OrderedCollectionMixin,
to_ordered_collection_page,
@ -19,7 +20,7 @@ from bookwyrm.models.activitypub_mixin import (
from bookwyrm.settings import PAGE_LENGTH
# pylint: disable=invalid-name
# pylint: disable=invalid-name,too-many-public-methods
@patch("bookwyrm.activitystreams.add_status_task.delay")
@patch("bookwyrm.models.activitypub_mixin.broadcast_task.apply_async")
class ActivitypubMixins(TestCase):
@ -421,3 +422,14 @@ class ActivitypubMixins(TestCase):
self.assertEqual(page_2.id, "http://fish.com/?page=2")
self.assertEqual(page_2.orderedItems[0]["content"], "test status 14")
self.assertEqual(page_2.orderedItems[-1]["content"], "test status 0")
def test_broadcast_task(self, *_):
"""Should be calling asyncio"""
recipients = [
"https://instance.example/user/inbox",
"https://instance.example/okay/inbox",
]
with patch("bookwyrm.models.activitypub_mixin.asyncio.run") as mock:
broadcast_task(self.local_user.id, {}, recipients)
self.assertTrue(mock.called)
self.assertEqual(mock.call_count, 1)