split metadata and notifier into separate files

This commit is contained in:
holger krekel 2024-03-31 01:55:53 +01:00
parent a31d998e67
commit 3d054847a0
3 changed files with 214 additions and 199 deletions

View File

@ -1,7 +1,4 @@
import time
from pathlib import Path
from threading import Thread
from queue import PriorityQueue
from socketserver import (
UnixStreamServer,
StreamRequestHandler,
@ -10,9 +7,9 @@ from socketserver import (
import sys
import logging
import os
import requests
from .filedict import FileDict
from .notifier import Notifier
DICTPROXY_HELLO_CHAR = "H"
@ -23,146 +20,51 @@ DICTPROXY_SET_CHAR = "S"
DICTPROXY_COMMIT_TRANSACTION_CHAR = "C"
DICTPROXY_TRANSACTION_CHARS = "BSC"
class Metadata:
# each SETMETADATA on this key appends to a list of unique device tokens
# which only ever get removed if the upstream indicates the token is invalid
METADATA_TOKEN_KEY = "devicetoken"
class Notifier:
URL = "https://notifications.delta.chat/notify"
CONNECTION_TIMEOUT = 60.0 # seconds until http-request is given up
NOTIFICATION_RETRY_DELAY = 8.0 # seconds with exponential backoff
MAX_NUMBER_OF_TRIES = 6
# exponential backoff means we try for 8^5 seconds, approximately 10 hours
DEVICETOKEN_KEY = "devicetoken"
def __init__(self, vmail_dir):
self.vmail_dir = vmail_dir
self.notification_dir = vmail_dir / "pending_notifications"
if not self.notification_dir.exists():
self.notification_dir.mkdir()
self.retry_queues = [PriorityQueue() for _ in range(self.MAX_NUMBER_OF_TRIES)]
def get_metadata_dict(self, addr):
return FileDict(self.vmail_dir / addr / "metadata.json")
def add_token_to_addr(self, addr, token):
with self.get_metadata_dict(addr).modify() as data:
tokens = data.get(METADATA_TOKEN_KEY)
tokens = data.get(self.DEVICETOKEN_KEY)
if tokens is None:
data[METADATA_TOKEN_KEY] = [token]
data[self.DEVICETOKEN_KEY] = [token]
elif token not in tokens:
tokens.append(token)
def remove_token_from_addr(self, addr, token):
with self.get_metadata_dict(addr).modify() as data:
tokens = data.get(METADATA_TOKEN_KEY, [])
tokens = data.get(self.DEVICETOKEN_KEY, [])
if token in tokens:
tokens.remove(token)
def get_tokens_for_addr(self, addr):
return self.get_metadata_dict(addr).read().get(METADATA_TOKEN_KEY, [])
def new_message_for_addr(self, addr):
for token in self.get_tokens_for_addr(addr):
self.notification_dir.joinpath(token).write_text(addr)
self.add_token_for_retry(token)
def add_token_for_retry(self, token, retry_num=0):
if retry_num >= self.MAX_NUMBER_OF_TRIES:
return False
when = time.time()
if retry_num > 0:
# backup exponentially with number of retries
when += pow(self.NOTIFICATION_RETRY_DELAY, retry_num)
self.retry_queues[retry_num].put((when, token))
return True
def requeue_persistent_pending_tokens(self):
for token_path in self.notification_dir.iterdir():
self.add_token_for_retry(token_path.name)
def start_notification_threads(self):
self.requeue_persistent_pending_tokens()
threads = {}
for retry_num in range(len(self.retry_queues)):
num_threads = {0: 4}.get(retry_num, 2)
threads[retry_num] = []
for _ in range(num_threads):
threads[retry_num].append(NotifyThread(self, retry_num))
threads[retry_num][-1].start()
return threads
mdict = self.get_metadata_dict(addr).read()
return mdict.get(self.DEVICETOKEN_KEY, [])
class NotifyThread(Thread):
def __init__(self, notifier, retry_num):
super().__init__(daemon=True)
self.notifier = notifier
self.retry_num = retry_num
def stop(self):
self.notifier.retry_queues[self.retry_num].put((None, None))
def run(self):
requests_session = requests.Session()
while self.retry_one(requests_session):
pass
def retry_one(self, requests_session, sleep=time.sleep):
# takes the next token from the per-retry-number PriorityQueue
# which is ordered by "when" (as set by add_token_for_retry()).
# If the request to notification server fails the token is
# queued to the next retry-number's PriorityQueue
# until it finally is dropped if MAX_NUMBER_OF_TRIES is exceeded
when, token = self.notifier.retry_queues[self.retry_num].get()
if when is None:
return False
wait_time = when - time.time()
if wait_time > 0:
sleep(wait_time)
self.perform_request_to_notification_server(requests_session, token)
return True
def perform_request_to_notification_server(self, requests_session, token):
token_path = self.notifier.notification_dir.joinpath(token)
try:
timeout = self.notifier.CONNECTION_TIMEOUT
res = requests_session.post(self.notifier.URL, data=token, timeout=timeout)
except requests.exceptions.RequestException as e:
res = e
else:
if res.status_code in (200, 410):
if res.status_code == 410:
# 410 Gone: means the token is no longer valid.
try:
addr = token_path.read_text()
except FileNotFoundError:
logging.warning("no address for token %r:", token)
return
self.notifier.remove_token_from_addr(addr, token)
token_path.unlink(missing_ok=True)
return
logging.warning("Notification request failed: %r", res)
if not self.notifier.add_token_for_retry(token, retry_num=self.retry_num + 1):
token_path.unlink(missing_ok=True)
logging.warning("dropping token after %d tries: %r", self.retry_num, token)
def handle_dovecot_protocol(rfile, wfile, notifier):
def handle_dovecot_protocol(rfile, wfile, notifier, metadata):
transactions = {}
while True:
msg = rfile.readline().strip().decode()
if not msg:
break
res = handle_dovecot_request(msg, transactions, notifier)
res = handle_dovecot_request(msg, transactions, notifier, metadata)
if res:
wfile.write(res.encode("ascii"))
wfile.flush()
def handle_dovecot_request(msg, transactions, notifier):
def handle_dovecot_request(msg, transactions, notifier, metadata):
# see https://doc.dovecot.org/3.0/developer_manual/design/dict_protocol/
short_command = msg[0]
parts = msg[1:].split("\t")
@ -172,8 +74,8 @@ def handle_dovecot_request(msg, transactions, notifier):
if keyparts[0] == "priv":
keyname = keyparts[2]
addr = parts[1]
if keyname == METADATA_TOKEN_KEY:
res = " ".join(notifier.get_tokens_for_addr(addr))
if keyname == metadata.DEVICETOKEN_KEY:
res = " ".join(metadata.get_tokens_for_addr(addr))
return f"O{res}\n"
logging.warning("lookup ignored: %r", msg)
return "N\n"
@ -206,10 +108,10 @@ def handle_dovecot_request(msg, transactions, notifier):
keyname = parts[1].split("/")
value = parts[2] if len(parts) > 2 else ""
addr = transactions[transaction_id]["addr"]
if keyname[0] == "priv" and keyname[2] == METADATA_TOKEN_KEY:
notifier.add_token_to_addr(addr, value)
if keyname[0] == "priv" and keyname[2] == metadata.DEVICETOKEN_KEY:
metadata.add_token_to_addr(addr, value)
elif keyname[0] == "priv" and keyname[2] == "messagenew":
notifier.new_message_for_addr(addr)
notifier.new_message_for_addr(addr, metadata)
else:
# Transaction failed.
transactions[transaction_id]["res"] = "F\n"
@ -228,13 +130,17 @@ def main():
logging.error("vmail dir does not exist: %r", vmail_dir)
return 1
notifier = Notifier(vmail_dir)
notifier.start_notification_threads()
notification_dir = vmail_dir / "pending_notifications"
if not notification_dir.exists():
notification_dir.mkdir()
metadata = Metadata(vmail_dir)
notifier = Notifier(metadata, notification_dir)
notifier.start_notification_threads(metadata.remove_token_from_addr)
class Handler(StreamRequestHandler):
def handle(self):
try:
handle_dovecot_protocol(self.rfile, self.wfile, notifier)
handle_dovecot_protocol(self.rfile, self.wfile, notifier, metadata)
except Exception:
logging.exception("Exception in the dovecot dictproxy handler")
raise

View File

@ -0,0 +1,105 @@
import time
import logging
from threading import Thread
from queue import PriorityQueue
import requests
class Notifier:
URL = "https://notifications.delta.chat/notify"
CONNECTION_TIMEOUT = 60.0 # seconds until http-request is given up
NOTIFICATION_RETRY_DELAY = 8.0 # seconds with exponential backoff
MAX_NUMBER_OF_TRIES = 6
# exponential backoff means we try for 8^5 seconds, approximately 10 hours
def __init__(self, notification_dir):
self.notification_dir = notification_dir
self.retry_queues = [PriorityQueue() for _ in range(self.MAX_NUMBER_OF_TRIES)]
def new_message_for_addr(self, addr, metadata):
for token in metadata.get_tokens_for_addr(addr):
self.notification_dir.joinpath(token).write_text(addr)
self.add_token_for_retry(token)
def requeue_persistent_pending_tokens(self):
for token_path in self.notification_dir.iterdir():
self.add_token_for_retry(token_path.name)
def add_token_for_retry(self, token, retry_num=0):
if retry_num >= self.MAX_NUMBER_OF_TRIES:
return False
when = time.time()
if retry_num > 0:
# backup exponentially with number of retries
when += pow(self.NOTIFICATION_RETRY_DELAY, retry_num)
self.retry_queues[retry_num].put((when, token))
return True
def start_notification_threads(self, remove_token_from_addr):
self.requeue_persistent_pending_tokens()
threads = {}
for retry_num in range(len(self.retry_queues)):
num_threads = {0: 4}.get(retry_num, 2)
threads[retry_num] = []
for _ in range(num_threads):
thread = NotifyThread(self, retry_num, remove_token_from_addr)
threads[retry_num].append(thread)
thread.start()
return threads
class NotifyThread(Thread):
def __init__(self, notifier, retry_num, remove_token_from_addr):
super().__init__(daemon=True)
self.notifier = notifier
self.retry_num = retry_num
self.remove_token_from_addr = remove_token_from_addr
def stop(self):
self.notifier.retry_queues[self.retry_num].put((None, None))
def run(self):
requests_session = requests.Session()
while self.retry_one(requests_session):
pass
def retry_one(self, requests_session, sleep=time.sleep):
# takes the next token from the per-retry-number PriorityQueue
# which is ordered by "when" (as set by add_token_for_retry()).
# If the request to notification server fails the token is
# queued to the next retry-number's PriorityQueue
# until it finally is dropped if MAX_NUMBER_OF_TRIES is exceeded
when, token = self.notifier.retry_queues[self.retry_num].get()
if when is None:
return False
wait_time = when - time.time()
if wait_time > 0:
sleep(wait_time)
self.perform_request_to_notification_server(requests_session, token)
return True
def perform_request_to_notification_server(self, requests_session, token):
token_path = self.notifier.notification_dir.joinpath(token)
try:
timeout = self.notifier.CONNECTION_TIMEOUT
res = requests_session.post(self.notifier.URL, data=token, timeout=timeout)
except requests.exceptions.RequestException as e:
res = e
else:
if res.status_code in (200, 410):
if res.status_code == 410:
# 410 Gone: means the token is no longer valid.
try:
addr = token_path.read_text()
except FileNotFoundError:
logging.warning("no address for token %r:", token)
return
self.remove_token_from_addr(addr, token)
token_path.unlink(missing_ok=True)
return
logging.warning("Notification request failed: %r", res)
if not self.notifier.add_token_for_retry(token, retry_num=self.retry_num + 1):
token_path.unlink(missing_ok=True)
logging.warning("dropping token after %d tries: %r", self.retry_num, token)

View File

@ -5,16 +5,26 @@ import requests
from chatmaild.metadata import (
handle_dovecot_request,
handle_dovecot_protocol,
Metadata,
)
from chatmaild.notifier import (
Notifier,
NotifyThread,
)
@pytest.fixture
def notifier(tmp_path):
def notifier(metadata):
notification_dir = metadata.vmail_dir.joinpath("pending_notifications")
notification_dir.mkdir()
return Notifier(notification_dir)
@pytest.fixture
def metadata(tmp_path):
vmail_dir = tmp_path.joinpath("vmaildir")
vmail_dir.mkdir()
return Notifier(vmail_dir)
return Metadata(vmail_dir)
@pytest.fixture
@ -50,72 +60,76 @@ def get_mocked_requests(statuslist):
return ReqMock()
def test_notifier_persistence(tmp_path, testaddr, testaddr2):
notifier1 = Notifier(tmp_path)
notifier2 = Notifier(tmp_path)
assert not notifier1.get_tokens_for_addr(testaddr)
assert not notifier2.get_tokens_for_addr(testaddr)
def test_metadata_persistence(tmp_path, testaddr, testaddr2):
metadata1 = Metadata(tmp_path)
metadata2 = Metadata(tmp_path)
assert not metadata1.get_tokens_for_addr(testaddr)
assert not metadata2.get_tokens_for_addr(testaddr)
notifier1.add_token_to_addr(testaddr, "01234")
notifier1.add_token_to_addr(testaddr2, "456")
assert notifier2.get_tokens_for_addr(testaddr) == ["01234"]
assert notifier2.get_tokens_for_addr(testaddr2) == ["456"]
notifier2.remove_token_from_addr(testaddr, "01234")
assert not notifier1.get_tokens_for_addr(testaddr)
assert notifier1.get_tokens_for_addr(testaddr2) == ["456"]
metadata1.add_token_to_addr(testaddr, "01234")
metadata1.add_token_to_addr(testaddr2, "456")
assert metadata2.get_tokens_for_addr(testaddr) == ["01234"]
assert metadata2.get_tokens_for_addr(testaddr2) == ["456"]
metadata2.remove_token_from_addr(testaddr, "01234")
assert not metadata1.get_tokens_for_addr(testaddr)
assert metadata1.get_tokens_for_addr(testaddr2) == ["456"]
def test_remove_nonexisting(tmp_path, testaddr):
notifier1 = Notifier(tmp_path)
notifier1.add_token_to_addr(testaddr, "123")
notifier1.remove_token_from_addr(testaddr, "1l23k1l2k3")
assert notifier1.get_tokens_for_addr(testaddr) == ["123"]
def test_remove_nonexisting(metadata, tmp_path, testaddr):
metadata.add_token_to_addr(testaddr, "123")
metadata.remove_token_from_addr(testaddr, "1l23k1l2k3")
assert metadata.get_tokens_for_addr(testaddr) == ["123"]
def test_notifier_remove_without_set(notifier, testaddr):
notifier.remove_token_from_addr(testaddr, "123")
assert not notifier.get_tokens_for_addr(testaddr)
def test_notifier_remove_without_set(metadata, testaddr):
metadata.remove_token_from_addr(testaddr, "123")
assert not metadata.get_tokens_for_addr(testaddr)
def test_handle_dovecot_request_lookup_fails(notifier, testaddr):
res = handle_dovecot_request(f"Lpriv/123/chatmail\t{testaddr}", {}, notifier)
def test_handle_dovecot_request_lookup_fails(notifier, metadata, testaddr):
res = handle_dovecot_request(
f"Lpriv/123/chatmail\t{testaddr}", {}, notifier, metadata
)
assert res == "N\n"
def test_handle_dovecot_request_happy_path(notifier, testaddr, token):
def test_handle_dovecot_request_happy_path(notifier, metadata, testaddr, token):
transactions = {}
# set device token in a transaction
tx = "1111"
msg = f"B{tx}\t{testaddr}"
res = handle_dovecot_request(msg, transactions, notifier)
assert not res and not notifier.get_tokens_for_addr(testaddr)
res = handle_dovecot_request(msg, transactions, notifier, metadata)
assert not res and not metadata.get_tokens_for_addr(testaddr)
assert transactions == {tx: dict(addr=testaddr, res="O\n")}
msg = f"S{tx}\tpriv/guid00/devicetoken\t{token}"
res = handle_dovecot_request(msg, transactions, notifier)
res = handle_dovecot_request(msg, transactions, notifier, metadata)
assert not res
assert len(transactions) == 1
assert notifier.get_tokens_for_addr(testaddr) == [token]
assert metadata.get_tokens_for_addr(testaddr) == [token]
msg = f"C{tx}"
res = handle_dovecot_request(msg, transactions, notifier)
res = handle_dovecot_request(msg, transactions, notifier, metadata)
assert res == "O\n"
assert len(transactions) == 0
assert notifier.get_tokens_for_addr(testaddr) == [token]
assert metadata.get_tokens_for_addr(testaddr) == [token]
# trigger notification for incoming message
tx2 = "2222"
assert handle_dovecot_request(f"B{tx2}\t{testaddr}", transactions, notifier) is None
assert (
handle_dovecot_request(f"B{tx2}\t{testaddr}", transactions, notifier, metadata)
is None
)
msg = f"S{tx2}\tpriv/guid00/messagenew"
assert handle_dovecot_request(msg, transactions, notifier) is None
assert handle_dovecot_request(msg, transactions, notifier, metadata) is None
assert notifier.retry_queues[0].get()[1] == token
assert handle_dovecot_request(f"C{tx2}", transactions, notifier) == "O\n"
assert handle_dovecot_request(f"C{tx2}", transactions, notifier, metadata) == "O\n"
assert not transactions
assert notifier.notification_dir.joinpath(token).exists()
def test_handle_dovecot_protocol_set_devicetoken(notifier):
def test_handle_dovecot_protocol_set_devicetoken(metadata, notifier):
rfile = io.BytesIO(
b"\n".join(
[
@ -127,12 +141,12 @@ def test_handle_dovecot_protocol_set_devicetoken(notifier):
)
)
wfile = io.BytesIO()
handle_dovecot_protocol(rfile, wfile, notifier)
handle_dovecot_protocol(rfile, wfile, notifier, metadata)
assert wfile.getvalue() == b"O\n"
assert notifier.get_tokens_for_addr("user@example.org") == ["01234"]
assert metadata.get_tokens_for_addr("user@example.org") == ["01234"]
def test_handle_dovecot_protocol_set_get_devicetoken(notifier):
def test_handle_dovecot_protocol_set_get_devicetoken(metadata, notifier):
rfile = io.BytesIO(
b"\n".join(
[
@ -144,19 +158,19 @@ def test_handle_dovecot_protocol_set_get_devicetoken(notifier):
)
)
wfile = io.BytesIO()
handle_dovecot_protocol(rfile, wfile, notifier)
assert notifier.get_tokens_for_addr("user@example.org") == ["01234"]
handle_dovecot_protocol(rfile, wfile, notifier, metadata)
assert metadata.get_tokens_for_addr("user@example.org") == ["01234"]
assert wfile.getvalue() == b"O\n"
rfile = io.BytesIO(
b"\n".join([b"HELLO", b"Lpriv/0123/devicetoken\tuser@example.org"])
)
wfile = io.BytesIO()
handle_dovecot_protocol(rfile, wfile, notifier)
handle_dovecot_protocol(rfile, wfile, notifier, metadata)
assert wfile.getvalue() == b"O01234\n"
def test_handle_dovecot_protocol_iterate(notifier):
def test_handle_dovecot_protocol_iterate(metadata, notifier):
rfile = io.BytesIO(
b"\n".join(
[
@ -166,45 +180,35 @@ def test_handle_dovecot_protocol_iterate(notifier):
)
)
wfile = io.BytesIO()
handle_dovecot_protocol(rfile, wfile, notifier)
handle_dovecot_protocol(rfile, wfile, notifier, metadata)
assert wfile.getvalue() == b"\n"
def test_notifier_thread_firstrun(notifier, testaddr):
def test_notifier_thread_deletes_persistent_file(metadata, notifier, testaddr):
reqmock = get_mocked_requests([200])
notifier.add_token_to_addr(testaddr, "01234")
notifier.new_message_for_addr(testaddr)
NotifyThread(notifier, retry_num=0).retry_one(reqmock)
metadata.add_token_to_addr(testaddr, "01234")
notifier.new_message_for_addr(testaddr, metadata)
NotifyThread(notifier, 0, None).retry_one(reqmock)
url, data, timeout = reqmock.requests[0]
assert data == "01234"
assert notifier.get_tokens_for_addr(testaddr) == ["01234"]
notifier.requeue_persistent_pending_tokens()
assert notifier.retry_queues[0].qsize() == 0
def test_notifier_thread_run(notifier, testaddr):
notifier.add_token_to_addr(testaddr, "01234")
notifier.new_message_for_addr(testaddr)
reqmock = get_mocked_requests([200])
NotifyThread(notifier, retry_num=0).retry_one(reqmock)
url, data, timeout = reqmock.requests[0]
assert data == "01234"
assert notifier.get_tokens_for_addr(testaddr) == ["01234"]
assert metadata.get_tokens_for_addr(testaddr) == ["01234"]
notifier.requeue_persistent_pending_tokens()
assert notifier.retry_queues[0].qsize() == 0
@pytest.mark.parametrize("status", [requests.exceptions.RequestException(), 404, 500])
def test_notifier_thread_connection_failures(notifier, testaddr, status, caplog):
def test_notifier_thread_connection_failures(
metadata, notifier, testaddr, status, caplog
):
"""test that tokens keep getting retried until they are given up."""
notifier.add_token_to_addr(testaddr, "01234")
notifier.new_message_for_addr(testaddr)
metadata.add_token_to_addr(testaddr, "01234")
notifier.new_message_for_addr(testaddr, metadata)
notifier.NOTIFICATION_RETRY_DELAY = 5
for i in range(notifier.MAX_NUMBER_OF_TRIES):
caplog.clear()
reqmock = get_mocked_requests([status])
sleep_calls = []
NotifyThread(notifier, retry_num=i).retry_one(reqmock, sleep=sleep_calls.append)
NotifyThread(notifier, i, None).retry_one(reqmock, sleep=sleep_calls.append)
assert notifier.retry_queues[i].qsize() == 0
assert "request failed" in caplog.records[0].msg
if i > 0:
@ -220,41 +224,41 @@ def test_notifier_thread_connection_failures(notifier, testaddr, status, caplog)
def test_start_and_stop_notification_threads(notifier, testaddr):
threads = notifier.start_notification_threads()
threads = notifier.start_notification_threads(None)
for retry_num, threadlist in threads.items():
for t in threadlist:
t.stop()
t.join()
def test_multi_device_notifier(notifier, testaddr):
notifier.add_token_to_addr(testaddr, "01234")
notifier.add_token_to_addr(testaddr, "56789")
notifier.new_message_for_addr(testaddr)
def test_multi_device_notifier(metadata, notifier, testaddr):
metadata.add_token_to_addr(testaddr, "01234")
metadata.add_token_to_addr(testaddr, "56789")
notifier.new_message_for_addr(testaddr, metadata)
reqmock = get_mocked_requests([200, 200])
NotifyThread(notifier, retry_num=0).retry_one(reqmock)
NotifyThread(notifier, retry_num=0).retry_one(reqmock)
NotifyThread(notifier, 0, None).retry_one(reqmock)
NotifyThread(notifier, 0, None).retry_one(reqmock)
assert notifier.retry_queues[0].qsize() == 0
assert notifier.retry_queues[1].qsize() == 0
url, data, timeout = reqmock.requests[0]
assert data == "01234"
url, data, timeout = reqmock.requests[1]
assert data == "56789"
assert notifier.get_tokens_for_addr(testaddr) == ["01234", "56789"]
assert metadata.get_tokens_for_addr(testaddr) == ["01234", "56789"]
def test_notifier_thread_run_gone_removes_token(notifier, testaddr):
notifier.add_token_to_addr(testaddr, "01234")
notifier.add_token_to_addr(testaddr, "45678")
notifier.new_message_for_addr(testaddr)
def test_notifier_thread_run_gone_removes_token(metadata, notifier, testaddr):
metadata.add_token_to_addr(testaddr, "01234")
metadata.add_token_to_addr(testaddr, "45678")
notifier.new_message_for_addr(testaddr, metadata)
reqmock = get_mocked_requests([410, 200])
NotifyThread(notifier, retry_num=0).retry_one(reqmock)
NotifyThread(notifier, retry_num=0).retry_one(reqmock)
NotifyThread(notifier, 0, metadata.remove_token_from_addr).retry_one(reqmock)
NotifyThread(notifier, 0, None).retry_one(reqmock)
url, data, timeout = reqmock.requests[0]
assert data == "01234"
url, data, timeout = reqmock.requests[1]
assert data == "45678"
assert notifier.get_tokens_for_addr(testaddr) == ["45678"]
assert metadata.get_tokens_for_addr(testaddr) == ["45678"]
assert notifier.retry_queues[0].qsize() == 0
assert notifier.retry_queues[1].qsize() == 0