Firehose server occasionally loses ndb context

[snarfed]

When the firehose server uses datastore storage, individual subscription threads occasionally lose their ndb context and crash. 😕 Stack trace below; originally reported as snarfed/bridgy-fed#2150, example in error tracking. No longer takes down the whole firehose server like it did originally, but still not great.

Only happens a few times a day right now, and the symptom is just that it drops a subscriber, who will then reconnect, so it's not that bad.

Current root cause theory is some tasklet magic for swapping the active context that I don't fully understand:

...single-thread, contextvars-frame mismatch:

• _state.context is a ContextVar.
• ndb's tasklet/eventloop machinery uses contextvars internally; some callbacks run under a captured/copied context (contextvars.Context.run() semantics).
• If a callback was scheduled in a frame where _state.context was None, and then runs after our @ndb_context decorator entered with ctx.use(), the callback's view of the ContextVar is the captured-earlier value (None), not our just-set one.
• Key.new → _project_from_app → get_context running inside such a callback would see None even though our with is still "active" on the synchronous stack.

...

The correct fix, if the theory holds, is to establish a fresh context at the right level so the storage calls don't depend on tasklet-frame magic to see one. Concretely, in the firehose path, wrap each event's processing in its own with self.ndb_client.context(): — e.g. inside process_event or at the top of the per-event work in subscribe(). That way:

• The context is entered on the current contextvars frame, synchronously, right before any tasklet/eventloop work begins. No earlier-captured-frame can shadow it.
• @ndb.non_transactional() on read() becomes a no-op (we're not in a transaction) but stays as defense-in-depth.
• You don't depend on read_blocks_by_seq's suspended-generator context still being the active one — which is the fragile thing here.

The trade-off vs. one-context-per-websocket: per-event contexts mean no cache reuse across events, but the existing code is already doing per-call contexts via @ndb_context anyway, so you're not losing anything.

Traceback (most recent call last):
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/arroba/firehose.py", line 238, in run
    return self._collect()
           ^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/arroba/firehose.py", line 298, in _collect
    header, payload = process_event(event)
                      ^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/arroba/firehose.py", line 351, in process_event
    tree.add_covering_proofs(event, blocks=event.blocks)
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/arroba/mst.py", line 990, in add_covering_proofs
    right_block = load(right_cid, 'R')
                  ^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/arroba/mst.py", line 953, in load
    blocks[cid] = self.storage.read(cid)
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/arroba/datastore_storage.py", line 295, in decorated
    ret = fn(self, *args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/_transaction.py", line 468, in non_transactional_inner_wrapper
    return wrapped(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/arroba/datastore_storage.py", line 772, in read
    block = AtpBlock.get_by_id(cid.encode('base32'))
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/_options.py", line 102, in wrapper
    return wrapped(*pass_args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/utils.py", line 150, in positional_wrapper
    return wrapped(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/model.py", line 5846, in _get_by_id
    ).result()
      ^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/tasklets.py", line 210, in result
    self.check_success()
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/tasklets.py", line 157, in check_success
    raise self._exception
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/tasklets.py", line 319, in _advance_tasklet
    yielded = self.generator.throw(type(error), error, traceback)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/key.py", line 943, in get
    entity_pb = yield _datastore_api.lookup(self._key, _options)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/tasklets.py", line 323, in _advance_tasklet
    yielded = self.generator.send(send_value)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/_datastore_api.py", line 145, in lookup
    result = yield _cache.global_get(cache_key)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/tasklets.py", line 489, in tasklet_wrapper
    context = context_module.get_context()
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/layers/google.python.pip/pip/lib/python3.12/site-packages/google/cloud/ndb/context.py", line 117, in get_context
    raise exceptions.ContextError()
google.cloud.ndb.exceptions.ContextError: No current context. NDB calls must be made in context established by google.cloud.ndb.Client.context.