xrpc_sync.subscribeRepos: wait up to 60s for out-of-order sequence numbers

for #34

Author: snarfed

arroba/tests/test_xrpc_sync.py

@@ -1,6 +1,8 @@
 """Unit tests for xrpc_sync.py."""
+from datetime import timedelta
 from io import BytesIO
 from threading import Semaphore, Thread
+import time
 from unittest import skip
 from unittest.mock import patch
 
@@ -18,7 +20,7 @@
 from .. import server
 from ..storage import Action, Storage, SUBSCRIBE_REPOS_NSID
 from .. import util
-from ..util import dag_cbor_cid, int_to_tid, next_tid
+from ..util import dag_cbor_cid, int_to_tid, next_tid, tid_to_int
 from .. import xrpc_sync
 
 from . import testutil
@@ -522,7 +524,7 @@ def test_get_blocks_repo_tombstoned(self):
     #     with self.assertRaises(ValueError):
     #         sync.loadCheckout(syncStorage, checkoutCar, repoDid, keypair.did())
 
-    # based atproto/packages/pds/tests/sync/list.test.ts
+    # based on atproto/packages/pds/tests/sync/list.test.ts
     # def test_paginates_listed_hosted_repos(self):
     #     full = xrpc_sync.list_repos({})
     #     pt1 = xrpc_sync.list_repos({}, limit=2)
@@ -562,7 +564,8 @@ def subscribe(self, received, delivered=None, limit=None, cursor=None):
                 return
 
     def assertCommitMessage(self, commit_msg, record=None, write=None,
-                            repo=None, cur=None, prev=None, seq=None):
+                            repo=None, cur=None, prev=None, seq=None,
+                            check_commit=True):
         if not repo:
             repo = self.repo
         if not cur:
@@ -597,7 +600,7 @@ def assertCommitMessage(self, commit_msg, record=None, write=None,
             record_cid = dag_cbor_cid(record)
             mst_entry = {
                 'e': [{
-                    'k': f'co.ll/{int_to_tid(util._tid_ts_last)}'.encode(),
+                    'k': f'co.ll/{write.rkey}'.encode(),
                     'v': record_cid,
                     'p': 0,
                     't': None,
@@ -613,20 +616,21 @@ def assertCommitMessage(self, commit_msg, record=None, write=None,
                 'l': None,
             }
 
-        commit_record = {
-            'version': 3,
-            'did': repo.did,
-            'data': dag_cbor_cid(mst_entry),
-            'rev': int_to_tid(seq, clock_id=0),
-            'prev': prev,
-        }
-
         msg_records = [b.decoded for b in msg_blocks]
         # TODO: if I util.sign(commit_record), the sig doesn't match. why?
         for msg_record in msg_records:
             msg_record.pop('sig', None)
 
-        self.assertIn(commit_record, msg_records)
+        if check_commit:
+            commit_record = {
+                'version': 3,
+                'did': repo.did,
+                'data': dag_cbor_cid(mst_entry),
+                'rev': int_to_tid(seq, clock_id=0),
+                'prev': prev,
+            }
+            self.assertIn(commit_record, msg_records)
+
         if record:
             self.assertIn(record, msg_records)
 
@@ -766,7 +770,7 @@ def test_include_preexisting_record_block(self, *_):
 
         subscriber.join()
 
-    def test_tombstone(self, *_):
+    def test_tombstoned(self, *_):
         # second repo: bob
         bob_repo = Repo.create(server.storage, 'did:bob',
                                handle='bo.bb', signing_key=self.key)
@@ -826,6 +830,56 @@ def test_tombstone(self, *_):
             'time': testutil.NOW.isoformat(),
         }, payload)
 
+    @patch('arroba.xrpc_sync.NEW_EVENTS_TIMEOUT', timedelta(seconds=2))
+    def test_subscribe_repos_skipped_seq(self, *_):
+        # https://github.com/snarfed/arroba/issues/34
+        received = []
+        delivered = Semaphore(value=0)
+        subscriber = Thread(target=self.subscribe,
+                              args=[received, delivered, 2])
+        subscriber.start()
+
+        # prepare two writes with seqs 4 and 5
+        write_4 = Write(Action.CREATE, 'co.ll', next_tid(), {'a': 'b'})
+        commit_4 = Repo.format_commit(repo=self.repo, writes=[write_4])
+        self.assertEqual(4, tid_to_int(commit_4.commit.decoded['rev']))
+
+        write_5 = Write(Action.CREATE, 'co.ll', next_tid(), {'x': 'y'})
+        commit_5 = Repo.format_commit(repo=self.repo, writes=[write_5])
+        self.assertEqual(5, tid_to_int(commit_5.commit.decoded['rev']))
+
+        prev = self.repo.head.cid
+
+        with self.assertLogs() as logs:
+            # first write, skip seq 4, write with seq 5 instead
+            self.repo.apply_commit(commit_5)
+            head_5 = self.repo.head.cid
+
+            # there's a small chance that this could be flaky, if >.2s elapses
+            # between starting the subscriber above and receiving the second
+            # write below
+            time.sleep(.1)
+
+            # shouldn't receive the event yet
+            self.assertEqual(0, len(received))
+
+            # second write, use seq 4 that we skipped above
+            self.repo.apply_commit(commit_4)
+
+            delivered.acquire()
+            delivered.acquire()
+
+        self.assertIn('WARNING:arroba.xrpc_sync:Waiting for seq 4', logs.output)
+
+        # should receive both commits
+        self.assertEqual(2, len(received))
+        self.assertCommitMessage(received[0], {'a': 'b'}, write=write_4,
+                                 cur=self.repo.head.cid, prev=prev, seq=4)
+        self.assertCommitMessage(received[1], {'x': 'y'}, write=write_5,
+                                 cur=head_5, prev=prev, seq=5, check_commit=False)
+
+        subscriber.join()
+
 
 class DatastoreXrpcSyncTest(XrpcSyncTest, testutil.DatastoreTest):
     STORAGE_CLS = DatastoreStorage

arroba/tests/testutil.py

@@ -1,6 +1,8 @@
 """Common test utility code."""
+import contextlib
 from datetime import datetime, timezone
 import json
+import logging
 import random
 import os
 import unittest
@@ -137,6 +139,25 @@ def random_keys_and_cids(num):
     def random_objects(num):
         return {next_tid(): {'foo': random.randint(1, 999999999)} for i in range(num)}
 
+    @contextlib.contextmanager
+    def assertLogs(self):
+        """Wraps :meth:`unittest.TestCase.assertLogs` and enables/disables logs.
+
+        Copied from bridgy-fed/tests/testutil.py
+        """
+        orig_disable_level = logging.root.manager.disable
+        logging.disable(logging.NOTSET)
+
+        try:
+            with super().assertLogs() as logs:
+                yield logs
+        finally:
+            # emit logs that were captured
+            for record in logs.records:
+                if record.levelno >= orig_disable_level:
+                    logging.root.handle(record)
+            logging.disable(orig_disable_level)
+
 
 class DatastoreTest(TestCase):
     ndb_client = ndb.Client(project='app', credentials=AnonymousCredentials())

arroba/xrpc_sync.py

@@ -4,6 +4,7 @@
 import logging
 import os
 from threading import Condition
+import time
 
 from carbox import car
 import dag_cbor
@@ -127,13 +128,13 @@ def subscribe_repos(cursor=None):
     Returns:
       (dict, dict) tuple: (header, payload)
     """
-    last_seq = server.storage.last_seq(SUBSCRIBE_REPOS_NSID)
+    cur_seq = server.storage.last_seq(SUBSCRIBE_REPOS_NSID)
 
     def handle(event):
-        nonlocal last_seq
+        nonlocal cur_seq
 
         if isinstance(event, dict):  # non-commit event
-            last_seq = event['seq']
+            cur_seq = event['seq']
             type = event.pop('$type')
             type_fragment = type.removeprefix('com.atproto.sync.subscribeRepos')
             assert type_fragment != type, type
@@ -142,14 +143,14 @@ def handle(event):
         assert isinstance(event, CommitData), \
             f'unexpected event type {event.__class__} {event}'
 
-        last_seq = event.commit.seq
+        cur_seq = event.commit.seq
         commit = event.commit.decoded
         car_blocks = [car.Block(cid=block.cid, data=block.encoded,
                                 decoded=block.decoded)
                       for block in event.blocks.values()]
         return ({  # header
-          'op': 1,
-          't': '#commit',
+            'op': 1,
+            't': '#commit',
         }, {  # payload
             'repo': commit['did'],
             'ops': [{
@@ -173,14 +174,14 @@ def handle(event):
 
     if cursor is not None:
         # validate cursor
-        if cursor > last_seq:
-            msg = f'Cursor {cursor} is past our current sequence number {last_seq}'
+        if cursor > cur_seq:
+            msg = f'Cursor {cursor} is past our current sequence number {cur_seq}'
             logger.warning(msg)
             yield ({'op': -1}, {'error': 'FutureCursor', 'message': msg})
             return
 
         if ROLLBACK_WINDOW is not None:
-            rollback_start = max(last_seq - ROLLBACK_WINDOW - 1, 0)
+            rollback_start = max(cur_seq - ROLLBACK_WINDOW - 1, 0)
             if cursor < rollback_start:
                 logger.warning(f'Cursor {cursor} is before our rollback window; starting at {rollback_start}')
                 yield ({'op': 1, 't': '#info'}, {'name': 'OutdatedCursor'})
@@ -190,14 +191,32 @@ def handle(event):
         for event in server.storage.read_events_by_seq(start=cursor):
             yield handle(event)
 
-    # serve new events as they happen
+    # serve new events as they happen. if we see a sequence number skipped, wait
+    # for it up to NEW_EVENTS_TIMEOUT before giving up on it and moving on
     logger.info(f'serving new events')
+    timeout_s = NEW_EVENTS_TIMEOUT.total_seconds()
+    last_yield = time.time()
+
     while True:
         with new_events:
-            new_events.wait(NEW_EVENTS_TIMEOUT.total_seconds())
+            new_events.wait(timeout_s)
+
+        for commit_data in server.storage.read_events_by_seq(start=cur_seq + 1):
+            last_seq = cur_seq
+            event = handle(commit_data)
+
+            waited_enough = time.time() - last_yield > timeout_s
+            if cur_seq == last_seq + 1 or waited_enough:
+                if cur_seq > last_seq + 1:
+                    logger.warning(f'Gave up waiting for seqs {last_seq + 1} to {cur_seq - 1}!')
+
+                last_yield = time.time()
+                yield event
+            else:
+                logger.warning(f'Waiting for seq {last_seq + 1}')
+                cur_seq = last_seq
+                break
 
-        for commit_data in server.storage.read_events_by_seq(start=last_seq + 1):
-            yield handle(commit_data)
 
 
 @server.server.method('com.atproto.sync.getBlocks')