File size: 3,763 Bytes
3a7eb07
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
"""
Serverless Task Queue
Lightweight in-memory queue for document processing.
"""

import asyncio
import logging
from typing import Optional

logger = logging.getLogger(__name__)

# Single global queue for the server process
_task_queue: Optional[asyncio.Queue] = None


def get_queue() -> asyncio.Queue:
    global _task_queue
    if _task_queue is None:
        _task_queue = asyncio.Queue()
    return _task_queue


async def document_worker():
    """Background worker that processes documents sequentially/concurrently."""
    from app.services.document_service import process_document_async

    queue = get_queue()
    logger.info("Document worker started.")

    while True:
        try:
            document_id = await queue.get()
            logger.info(f"Worker picked up document: {document_id}")

            try:
                # We await the document processing.
                # Concurrency is handled internally by Orchestrator or we can spawn tasks here.
                # Since DocumentService is async, we can just await it directly or create a task.
                asyncio.create_task(process_document_async(document_id))
            except Exception as e:
                logger.error(
                    f"Worker failed dispatching document {document_id}: {e}",
                    exc_info=True,
                )
            finally:
                queue.task_done()

        except asyncio.CancelledError:
            logger.info("Document worker cancelled.")
            break
        except Exception as e:
            logger.error(f"Error in document worker loop: {e}", exc_info=True)
            await asyncio.sleep(1)


def enqueue_document_task(document_id: str):
    """Adds a document to the queue without blocking."""
    queue = get_queue()
    try:
        queue.put_nowait(document_id)
        logger.info(f"Document {document_id} enqueued.")
    except asyncio.QueueFull:
        logger.error(f"Failed to enqueue document {document_id}: Queue is full")


# ── Compliance Audit Queue ─────────────────────────────────────────────────────

_compliance_queue: Optional[asyncio.Queue] = None


def get_compliance_queue() -> asyncio.Queue:
    global _compliance_queue
    if _compliance_queue is None:
        _compliance_queue = asyncio.Queue()
    return _compliance_queue


async def compliance_worker():
    """Background worker that processes compliance audits."""
    from app.services.compliance_service import run_compliance_audit_async

    queue = get_compliance_queue()
    logger.info("Compliance worker started.")

    while True:
        try:
            audit_id = await queue.get()
            logger.info(f"Compliance worker picked up audit: {audit_id}")

            try:
                asyncio.create_task(run_compliance_audit_async(audit_id))
            except Exception as e:
                logger.error(
                    f"Worker failed dispatching audit {audit_id}: {e}", exc_info=True
                )
            finally:
                queue.task_done()

        except asyncio.CancelledError:
            logger.info("Compliance worker cancelled.")
            break
        except Exception as e:
            logger.error(f"Error in compliance worker loop: {e}", exc_info=True)
            await asyncio.sleep(1)


async def enqueue_compliance_task(audit_id: str):
    """Adds a compliance audit to the queue without blocking."""
    queue = get_compliance_queue()
    try:
        queue.put_nowait(audit_id)
        logger.info(f"Compliance audit {audit_id} enqueued.")
    except asyncio.QueueFull:
        logger.error(f"Failed to enqueue audit {audit_id}: Queue is full")