perf: cap query time, bound discussion recursion, isolate health checks

hive/server/bridge_api/thread.py

@@ -11,6 +11,15 @@
 
 log = logging.getLogger(__name__)
 
+# Hard caps to prevent connection-pool exhaustion on pathological threads.
+# _load_discussion walks the comment tree level by level, issuing one
+# _child_ids query per depth level. Without bounds, a 1000+ comment / 50+
+# depth thread can hold a connection for an unbounded time and starve the
+# pool. MAX_DEPTH bounds the number of sequential queries; MAX_THREAD_POSTS
+# bounds the total number of posts loaded into memory.
+MAX_THREAD_POSTS = 500
+MAX_DEPTH = 50
+
 @return_error_info
 async def get_discussion(context, author, permlink):
     """Modified `get_state` thread implementation."""
@@ -43,16 +52,21 @@ async def _get_post_id(db, author, permlink):
 
 async def _get_author_hide_id(db, author):
     """Given an author, retrieve the id from db."""
+    # Hide status changes rarely; cache to avoid spending a connection on every
+    # get_discussion request. The db layer caches the "not found" case too.
     sql = ("SELECT id FROM hive_posts_status WHERE list_type = '3'"
            "AND author = :a LIMIT 1")
-    return await db.query_one(sql, a=author)
+    return await db.query_one(sql, a=author, cache_key='author_hide_id_' + author,
+                              cache_ttl=300)
 
 
 async def _check_posts_hide_id(db, post_id):
     """Given an post_id, retrieve the id from db."""
     sql = ("SELECT id FROM hive_posts_status WHERE list_type = '1'"
            "AND post_id = :post_id LIMIT 1")
-    return await db.query_one(sql, post_id=post_id)
+    return await db.query_one(sql, post_id=post_id,
+                              cache_key='post_hide_id_' + str(post_id),
+                              cache_ttl=300)
 
 def _ref(post):
     return post['author'] + '/' + post['permlink']
@@ -80,7 +94,24 @@ async def _load_discussion(db, root_id):
     ids = []
     tree = {}
     todo = [root_id]
+    depth = 0
+    truncated = False
     while todo:
+        # Bound the number of sequential _child_ids queries (one per depth).
+        if depth >= MAX_DEPTH:
+            truncated = True
+            break
+        # Bound total posts collected so a wide thread cannot exhaust memory
+        # or hold connections for too long. If this batch would exceed the cap,
+        # take only what fits, resolve their (empty) child mapping, and stop.
+        if len(ids) + len(todo) > MAX_THREAD_POSTS:
+            todo = todo[:MAX_THREAD_POSTS - len(ids)]
+            ids.extend(todo)
+            rows = await _child_ids(db, todo)
+            for pid, _cids in rows:
+                tree[pid] = []
+            truncated = True
+            break
         ids.extend(todo)
         rows = await _child_ids(db, todo)
         todo = []
@@ -93,6 +124,11 @@ async def _load_discussion(db, root_id):
 
             tree[pid] = cids
             todo.extend(cids)
+        depth += 1
+
+    if truncated:
+        log.warning("discussion %s truncated at depth=%d posts=%d",
+                    root_id, depth, len(ids))
 
     # load all post objects, build ref-map
     posts = await load_posts_keyed(db, ids)

hive/server/db.py

@@ -16,6 +16,14 @@
 
 CACHE_NAMESPACE = "hivemind"
 
+# Per-query execution timeout (milliseconds). PostgreSQL cancels any statement
+# running longer than this, immediately releasing the connection back to the
+# pool. This is the safety net that prevents a single pathological query (e.g.
+# 160s get_discussion lookups) from exhausting the aiopg pool. The pool's
+# acquire `timeout` only bounds how long we wait for a free connection, NOT
+# query execution time. See connection-pool-exhaustion incident.
+STATEMENT_TIMEOUT_MS = 30000  # 30 seconds
+
 # Sentinel value to represent 'record not found' in cache.
 # Using a string marker that can be easily serialized/deserialized.
 # This allows us to distinguish between:
@@ -87,33 +95,77 @@ async def create(cls, url, redis_url=None, pool_size=20):
 
     def __init__(self):
         self.db = None
+        # Dedicated single-connection engine for health checks, isolated from
+        # the main pool so that a saturated main pool cannot starve /health and
+        # /head_age (which would cause the ELB to mark the instance unhealthy).
+        self.health_db = None
         self.redis_cache = None
         self._prep_sql = {}
 
     async def init(self, url, redis_url, pool_size=20):
         """Initialize the aiopg.sa engine."""
         conf = make_url(url)
+        # statement_timeout (ms) cancels runaway queries server-side so a single
+        # slow query cannot hold a pool connection indefinitely. See note on
+        # STATEMENT_TIMEOUT_MS above. Passed via the standard libpq `options`
+        # string (every psycopg2/libpq version accepts this); the newer
+        # `server_settings` dict is rejected as an invalid DSN option by the
+        # older psycopg2 shipped in the runtime image. Merge into conf.query
+        # (rather than a separate kwarg) so a DATABASE_URL that already carries
+        # an `options=...` query param does not cause a duplicate-keyword error;
+        # any pre-existing options string is preserved and appended to.
+        query = dict(conf.query)
+        timeout_opt = '-c statement_timeout=%d' % STATEMENT_TIMEOUT_MS
+        if 'options' in query and query['options']:
+            query['options'] = query['options'] + ' ' + timeout_opt
+        else:
+            query['options'] = timeout_opt
         self.db = await create_engine(user=conf.username,
                                       database=conf.database,
                                       password=conf.password,
                                       host=conf.host,
                                       port=conf.port,
                                       maxsize=pool_size,
                                       timeout=10,
-                                      **conf.query)
+                                      **query)
+        # Lightweight isolated engine (1 connection) for health checks. A short
+        # acquire timeout keeps health checks responsive instead of blocking.
+        self.health_db = await create_engine(user=conf.username,
+                                             database=conf.database,
+                                             password=conf.password,
+                                             host=conf.host,
+                                             port=conf.port,
+                                             maxsize=1,
+                                             timeout=2,
+                                             **query)
         if redis_url is not None:
             self.redis_cache = Cache.from_url(redis_url)
             self.redis_cache.serializer = SafeUniversalSerializer()
 
     def close(self):
         """Close pool."""
         self.db.close()
+        if self.health_db is not None:
+            self.health_db.close()
         if self.redis_cache is not None:
             self.redis_cache.close()
 
     async def wait_closed(self):
         """Wait for releasing and closing all acquired connections."""
         await self.db.wait_closed()
+        if self.health_db is not None:
+            await self.health_db.wait_closed()
+
+    async def query_row_health(self, sql, **kwargs):
+        """Run a `SELECT 1*m` on the isolated health engine.
+
+        Bypasses the cache decorator and the main pool so that health checks
+        remain responsive even when the main pool is fully saturated by slow
+        API queries.
+        """
+        async with self.health_db.acquire() as conn:
+            cur = await self._query(conn, sql, **kwargs)
+            return await cur.first()
 
     @sqltimer
     @cacher

hive/server/serve.py

@@ -35,7 +35,10 @@ async def db_head_state(context):
     db = context['db']
     sql = ("SELECT num, created_at, extract(epoch from created_at) ts "
            "FROM hive_blocks ORDER BY num DESC LIMIT 1")
-    row = await db.query_row(sql)
+    # Use the isolated health engine so a saturated main pool (e.g. many slow
+    # get_discussion queries) cannot cause /health to time out and get the
+    # instance killed by the ELB.
+    row = await db.query_row_health(sql)
     return dict(db_head_block=row['num'],
                 db_head_time=str(row['created_at']),
                 db_head_age=int(time.time() - float(row['ts'])))

tests/bridge_thread/test_bridge_thread.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+"""
+Unit tests for hive.server.bridge_api.thread.
+
+These are pure-logic tests that do NOT require a live database. They live
+under tests/bridge_thread/ (rather than tests/server/) on purpose: the
+tests/server/__init__.py eagerly opens a real DB connection, which would
+prevent these no-DB tests from running here.
+
+They exercise the connection-pool-exhaustion safeguards added to
+`_load_discussion` (MAX_DEPTH / MAX_THREAD_POSTS caps) and verify the
+hide-id lookups forward their cache_key/cache_ttl to the db layer.
+
+A fake async db records the calls it receives so assertions can inspect how
+many `_child_ids`-equivalent queries ran and what cache params were passed.
+"""
+
+# pylint: disable=protected-access,missing-docstring
+
+import asyncio
+import os
+import sys
+
+# Allow running directly (python test_bridge_thread.py) without pytest's rootdir.
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
+
+from hive.server.bridge_api import thread  # noqa: E402
+from hive.server.bridge_api.thread import (  # noqa: E402
+    MAX_DEPTH, MAX_THREAD_POSTS, _check_posts_hide_id, _get_author_hide_id,
+    _get_post_id, _load_discussion,
+)
+
+import pytest  # noqa: E402
+
+
+class FakeAsyncDb:
+    """Records calls and returns canned results for the methods thread.py uses.
+
+    `query_one` results are configured per-cache_key. `query_all` results are
+    configured per-call-count to simulate walking a comment tree level by level.
+    """
+
+    def __init__(self, query_all_seq=None, query_one_map=None):
+        self.query_all_calls = []
+        self.query_one_calls = []
+        self._query_all_seq = query_all_seq or []
+        self._query_one_map = query_one_map or {}
+
+    async def query_one(self, sql, **kwargs):
+        cache_key = kwargs.get('cache_key')
+        self.query_one_calls.append({'sql': sql, 'kwargs': kwargs})
+        return self._query_one_map.get(cache_key)
+
+    async def query_all(self, sql, **kwargs):
+        self.query_all_calls.append({'sql': sql, 'kwargs': kwargs})
+        idx = len(self.query_all_calls) - 1
+        if idx < len(self._query_all_seq):
+            return self._query_all_seq[idx]
+        return []
+
+
+def _run(coro):
+    """Run a coroutine to completion in a fresh event loop."""
+    loop = asyncio.new_event_loop()
+    try:
+        return loop.run_until_complete(coro)
+    finally:
+        loop.close()
+
+
+def _install_load_posts_keyed(monkeypatch, posts=None):
+    """Stub out load_posts_keyed so _load_discussion needs no real post data."""
+
+    async def _stub(_db, _ids, _truncate_body=0):
+        return posts or {}
+
+    monkeypatch.setattr(thread, 'load_posts_keyed', _stub)
+
+
+def _install_hide_pids_by_ids(monkeypatch, hidden=None):
+    """Stub out hide_pids_by_ids so no real filtering happens."""
+    hidden = hidden or set()
+
+    async def _stub(_db, cids):
+        return [pid for pid in cids if pid in hidden]
+
+    monkeypatch.setattr(thread, 'hide_pids_by_ids', _stub)
+
+
+def _child_ids_rows(parent_to_children):
+    """Build a query_all result list (one batch) from a {parent: [children]} map."""
+    return [{
+        'parent_id': pid,
+        'child_ids': cids
+    } for pid, cids in parent_to_children.items()]
+
+
+def test_load_discussion_respects_max_depth(monkeypatch):
+    """A chain deeper than MAX_DEPTH must stop after MAX_DEPTH _child_ids calls.
+
+    Build an infinitely deep single-child chain (1 -> 2 -> 3 -> ...). Each
+    query_all call returns the next level. Without the cap this loops forever.
+    """
+    _install_load_posts_keyed(monkeypatch)
+    _install_hide_pids_by_ids(monkeypatch)
+
+    call_count = {'n': 0}
+
+    def next_level():
+        call_count['n'] += 1
+        parent = call_count['n']  # level N's parent is N
+        return _child_ids_rows({parent: [parent + 1]})
+
+    db = FakeAsyncDb()
+    # Provide MAX_DEPTH+5 levels so we can prove it stops at the cap rather
+    # than running out of data.
+    db._query_all_seq = [next_level() for _ in range(MAX_DEPTH + 5)]
+
+    _run(_load_discussion(db, 1))
+
+    # One _child_ids query per depth level, capped at MAX_DEPTH.
+    assert len(db.query_all_calls) == MAX_DEPTH
+
+
+def test_load_discussion_respects_max_thread_posts(monkeypatch):
+    """A very wide thread must be truncated at MAX_THREAD_POSTS total posts."""
+    _install_load_posts_keyed(monkeypatch)
+    _install_hide_pids_by_ids(monkeypatch)
+
+    # Level 1: root has 600 children (already > MAX_THREAD_POSTS=500).
+    wide_children = list(range(1000, 1600))
+    db = FakeAsyncDb(query_all_seq=[_child_ids_rows({1: wide_children})])
+
+    _run(_load_discussion(db, 1))
+
+    # Root level resolves, then the over-cap batch is truncated to fit.
+    assert len(db.query_all_calls) == 2
+    # The second call's parent_ids length is what got added: 500 - 1 = 499.
+    second_call_ids = db.query_all_calls[1]['kwargs']['ids']
+    assert len(second_call_ids) == MAX_THREAD_POSTS - 1
+
+
+def test_load_discussion_terminates_on_leaf(monkeypatch):
+    """A root with no children completes in a single _child_ids call."""
+    _install_load_posts_keyed(monkeypatch)
+    _install_hide_pids_by_ids(monkeypatch)
+
+    db = FakeAsyncDb(query_all_seq=[_child_ids_rows({1: []})])
+    _run(_load_discussion(db, 1))
+    assert len(db.query_all_calls) == 1
+
+
+def test_get_post_id_forwards_cache_params():
+    """_get_post_id must pass a long-TTL cache_key so the lookup is cached."""
+    db = FakeAsyncDb(query_one_map={'post_id_a_p': 42})
+    result = _run(_get_post_id(db, 'a', 'p'))
+    assert result == 42
+    call = db.query_one_calls[0]
+    assert call['kwargs']['cache_key'] == 'post_id_a_p'
+    assert call['kwargs']['cache_ttl'] == 3600
+
+
+def test_get_author_hide_id_forwards_cache_params():
+    """_get_author_hide_id must now be cached (was uncached before the fix)."""
+    db = FakeAsyncDb()
+    _run(_get_author_hide_id(db, 'alice'))
+    call = db.query_one_calls[0]
+    assert call['kwargs']['cache_key'] == 'author_hide_id_alice'
+    assert call['kwargs']['cache_ttl'] == 300
+
+
+def test_check_posts_hide_id_forwards_cache_params():
+    """_check_posts_hide_id must now be cached (was uncached before the fix)."""
+    db = FakeAsyncDb()
+    _run(_check_posts_hide_id(db, 99))
+    call = db.query_one_calls[0]
+    assert call['kwargs']['cache_key'] == 'post_hide_id_99'
+    assert call['kwargs']['cache_ttl'] == 300
+
+
+if __name__ == '__main__':
+    # Allow `python test_bridge_thread.py` style execution without pytest.
+    sys.exit(pytest.main([__file__, '-v']))