Skip to content

Commit 377049d

Browse files
ChrisNolanChrisNolan
authored
fix(utils): improve async DB connection handling (#1120)
Phase-1 fix for the tracing-persistence connection leak/runaway (#1119): shared bounded executor for run_async, close_old_connections() before and connections.close_all() after each task, atexit shutdown. Phase-2 Huey refactor tracked separately (PRDs/TRACING_PERSISTENCE_PHASE2_HUEY_REFACTOR.md). Co-authored-by: Chris Nolan <chrisnolan.ca+github@gmail.com> refs #1119
1 parent 4a426a7 commit 377049d

3 files changed

Lines changed: 259 additions & 1 deletion

File tree

‎PRDs/TRACING_PERSISTENCE_PHASE2_HUEY_REFACTOR.md‎

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
# Tracing Persistence Phase 2: Move API Thread Async to Huey DB Task
2+
3+
| Field | Value |
4+
| --------- | ---------------- |
5+
| Project | Karrio |
6+
| Version | 1.0 |
7+
| Date | 2026-06-09 |
8+
| Status | Planning |
9+
| Owner | Server/Core Team |
10+
| Type | Refactoring |
11+
| Reference | AGENTS.md |
12+
13+
---
14+
15+
## Executive Summary
16+
17+
Phase 1 mitigates idle PostgreSQL session buildup by cleaning DB connections in API-side async threads.
18+
Phase 2 removes this risk class entirely by migrating tracing persistence from request-local thread execution to the existing Huey `db_task` pattern already used for server background ORM work.
19+
20+
### Key Decisions
21+
22+
1. Persist tracing records via Huey `db_task` instead of `@utils.async_wrapper`.
23+
2. Keep existing tracing record schema and payload shape unchanged.
24+
3. Keep deduplication by `request_log_id` in worker task logic.
25+
4. Roll out behind existing `PERSIST_SDK_TRACING` flag with no API contract changes.
26+
27+
### Scope
28+
29+
| In Scope | Out of Scope |
30+
| --------------------------------------------- | -------------------------------------- |
31+
| Move persistence execution path to Huey task | Changes to tracing record model fields |
32+
| Keep existing dedupe and org-link behavior | Tracing UI/query redesign |
33+
| Add task-focused tests and integration checks | Broad logging/telemetry redesign |
34+
35+
---
36+
37+
## Problem Statement
38+
39+
### Current State
40+
41+
Tracing persistence for API requests is triggered in middleware and currently uses an async wrapper backed by a thread executor in server core utils.
42+
43+
Relevant files:
44+
45+
- `modules/core/karrio/server/core/middleware.py`
46+
- `modules/core/karrio/server/tracing/utils.py`
47+
- `modules/core/karrio/server/core/utils.py`
48+
49+
### Why Change
50+
51+
Even with Phase 1 cleanup, API request handling still depends on ad hoc thread async for ORM persistence.
52+
Karrio already has a standard and safer background ORM pattern via Huey `db_task` wrappers in events tasks.
53+
54+
### Desired State
55+
56+
API middleware enqueues a Huey tracing persistence task. Worker-side execution handles ORM writes and org-linking under known task lifecycle semantics.
57+
58+
---
59+
60+
## Existing Code Analysis
61+
62+
| Component | Location | Reuse Strategy |
63+
| ----------------------- | ----------------------------------------------------------------------- | ---------------------------------------------------- |
64+
| Tracing save entrypoint | `modules/core/karrio/server/tracing/utils.py` | Split into enqueue + worker task body |
65+
| Request middleware hook | `modules/core/karrio/server/core/middleware.py` | Keep call site stable, change implementation beneath |
66+
| Huey task pattern | `modules/events/karrio/server/events/task_definitions/base/__init__.py` | Mirror `@db_task` + tenant-aware style |
67+
| Worker settings | `apps/api/karrio/server/settings/workers.py` | Reuse existing queue execution model |
68+
69+
---
70+
71+
## Architecture Overview
72+
73+
```text
74+
Before (current)
75+
76+
HTTP Request
77+
|
78+
v
79+
SessionContext middleware
80+
|
81+
v
82+
save_tracing_records()
83+
|
84+
v
85+
ThreadPoolExecutor (API process)
86+
|
87+
v
88+
ORM bulk_create("tracing-record")
89+
90+
91+
After (phase 2)
92+
93+
HTTP Request
94+
|
95+
v
96+
SessionContext middleware
97+
|
98+
v
99+
enqueue_tracing_records_task(...)
100+
|
101+
v
102+
Huey queue
103+
|
104+
v
105+
Huey worker db_task
106+
|
107+
v
108+
ORM bulk_create("tracing-record")
109+
```
110+
111+
## Sequence
112+
113+
```text
114+
Client -> API: request
115+
API -> Middleware: complete response
116+
Middleware -> Tracing Utils: enqueue(payload)
117+
Tracing Utils -> Huey: task.delay(payload)
118+
Huey Worker -> DB: dedupe check + bulk_create
119+
Huey Worker -> DB: bulk_link_org (if org)
120+
```
121+
122+
---
123+
124+
## Technical Design
125+
126+
1. Create a dedicated tracing persistence task function in server-side tasks module.
127+
2. Move ORM write logic from nested async closure into task body.
128+
3. Keep payload minimal and serializable:
129+
- actor_id
130+
- org_id
131+
- schema
132+
- tracer context values (`request_id`, `request_log_id`, `object_id`)
133+
- flattened tracing records list (key, timestamp, record, connection metadata)
134+
4. Preserve behavior:
135+
- skip when `PERSIST_SDK_TRACING` is false
136+
- skip when no records or no actor
137+
- preserve request_log_id dedupe check
138+
- preserve org linking
139+
5. Keep middleware call shape unchanged to minimize blast radius.
140+
141+
### Compatibility Notes
142+
143+
- No API schema changes.
144+
- No migration required.
145+
- Existing tracing readers remain unchanged.
146+
147+
---
148+
149+
## Implementation Plan
150+
151+
| Step | Change | Files |
152+
| ---- | -------------------------------------------------------- | ---------------------------------------------------------------------------------------- |
153+
| 1 | Add Huey db_task for tracing persistence | `modules/events/karrio/server/events/task_definitions/base/*.py` |
154+
| 2 | Refactor tracing util to enqueue task payload | `modules/core/karrio/server/tracing/utils.py` |
155+
| 3 | Keep middleware integration stable | `modules/core/karrio/server/core/middleware.py` |
156+
| 4 | Add tests for enqueue + worker persistence behavior | `modules/core/karrio/server/core/tests/*`, `modules/events/karrio/server/events/tests/*` |
157+
| 5 | Validate under load and compare pg_stat_activity profile | ops verification |
158+
159+
---
160+
161+
## Testing Strategy
162+
163+
1. Unit tests
164+
- enqueue is called once per request context with expected payload
165+
- worker task no-ops on empty records or missing actor
166+
- worker task dedupe by `request_log_id`
167+
2. Integration tests
168+
- middleware path still results in saved tracing records
169+
- org links are created correctly when org exists
170+
3. Non-functional validation
171+
- run load test with tracing enabled
172+
- compare idle `karrio.api` DB sessions before/after
173+
174+
---
175+
176+
## Risks and Mitigations
177+
178+
| Risk | Impact | Mitigation |
179+
| ---------------------------------- | --------------------------------- | ----------------------------------------------------------------------------------- |
180+
| Task payload missing context field | Missing metadata in trace records | Contract test asserting payload keys |
181+
| Queue lag delays trace visibility | Delayed debugging data | Document eventual consistency; keep synchronous fallback toggle for troubleshooting |
182+
| Duplicate writes in retries | Data noise | Keep request_log_id dedupe guard in task |
183+
184+
---
185+
186+
## Migration and Rollback
187+
188+
### Migration
189+
190+
- Deploy code with task path enabled.
191+
- Keep `PERSIST_SDK_TRACING` configurable for staged rollout.
192+
193+
### Rollback
194+
195+
- Revert to previous tracing utils implementation.
196+
- Disable tracing persistence (`PERSIST_SDK_TRACING=False`) if immediate operational relief is needed.
197+
198+
---
199+
200+
## Definition of Done
201+
202+
- [ ] Tracing persistence no longer uses API-side thread async path.
203+
- [ ] Tracing writes run through Huey `db_task` worker path.
204+
- [ ] Existing tracing metadata and dedupe behavior preserved.
205+
- [ ] Tests added and passing for enqueue and persistence logic.
206+
- [ ] Load validation shows stable/expected idle DB session count.

‎modules/core/karrio/server/core/tests/test_async_db_cleanup.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
from unittest import TestCase
2+
from unittest.mock import patch
3+
4+
from karrio.server.core import utils
5+
6+
7+
class TestAsyncDbCleanup(TestCase):
8+
def test_run_async_cleans_up_db_connections_on_success(self):
9+
with patch("karrio.server.core.utils.close_old_connections") as close_old, patch.object(
10+
utils.connections, "close_all"
11+
) as close_all:
12+
result = utils.run_async(lambda: "ok").result(timeout=2)
13+
14+
self.assertEqual(result, "ok")
15+
close_old.assert_called_once_with()
16+
close_all.assert_called_once_with()
17+
18+
def test_run_async_cleans_up_db_connections_on_error(self):
19+
with patch("karrio.server.core.utils.close_old_connections") as close_old, patch.object(
20+
utils.connections, "close_all"
21+
) as close_all:
22+
future = utils.run_async(lambda: (_ for _ in ()).throw(ValueError("boom")))
23+
24+
with self.assertRaises(ValueError):
25+
future.result(timeout=2)
26+
27+
close_old.assert_called_once_with()
28+
close_all.assert_called_once_with()

‎modules/core/karrio/server/core/utils.py‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,13 @@
22
import typing
33
import inspect
44
import functools
5+
import atexit
56
from concurrent import futures
67
from datetime import timedelta, datetime, timezone
78
from typing import TypeVar, Union, Callable, Any, List, Optional
89

910
from django.conf import settings
11+
from django.db import close_old_connections, connections
1012
from django.utils.translation import gettext_lazy as _
1113
import django_email_verification.confirm as confirm
1214
import rest_framework_simplejwt.tokens as jwt
@@ -20,6 +22,17 @@
2022
T = TypeVar("T")
2123

2224

25+
# Reuse a bounded executor for server-side fire-and-forget operations to avoid
26+
# creating one thread pool per request and leaking thread-local DB connections.
27+
_ASYNC_EXECUTOR = futures.ThreadPoolExecutor(max_workers=4)
28+
29+
30+
@atexit.register
31+
def _shutdown_async_executor():
32+
# Registered with atexit: invoked automatically when the process exits.
33+
_ASYNC_EXECUTOR.shutdown(wait=False, cancel_futures=True)
34+
35+
2336
def identity(value: Union[Any, Callable]) -> Any:
2437
"""
2538
:param value: function or value desired to be wrapped
@@ -144,7 +157,18 @@ def run_async(callable: Callable[[], Any]) -> futures.Future:
144157
of a callable in a non-blocking thread and return a
145158
handle for a future response.
146159
"""
147-
return futures.ThreadPoolExecutor(max_workers=1).submit(callable)
160+
161+
def _wrapped_call():
162+
# Ensure this worker thread starts from a clean DB connection state.
163+
close_old_connections()
164+
165+
try:
166+
return callable()
167+
finally:
168+
# Close thread-local DB connections opened during background ORM work.
169+
connections.close_all()
170+
171+
return _ASYNC_EXECUTOR.submit(_wrapped_call)
148172

149173

150174
def error_wrapper(func):

0 commit comments

Comments
 (0)