File size: 3,131 Bytes
9d416ac
54a4bb3
9d416ac
c0e26bd
 
 
 
 
 
 
 
 
 
 
 
 
54a4bb3
 
c0e26bd
ec847b0
 
 
 
 
 
 
 
 
 
54a4bb3
c0e26bd
 
ec847b0
c0e26bd
 
54a4bb3
ec847b0
c0e26bd
54a4bb3
 
ec847b0
c0e26bd
9d416ac
ec847b0
 
 
c0e26bd
 
 
 
 
 
ec847b0
c0e26bd
 
 
 
ec847b0
 
c0e26bd
ec847b0
 
 
9d416ac
54a4bb3
c0e26bd
54a4bb3
 
c0e26bd
 
ec847b0
c0e26bd
 
54a4bb3
ec847b0
c0e26bd
54a4bb3
ec847b0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
54a4bb3
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
import asyncio
from datetime import datetime, timezone

from app.config import (
    ROOT_FOLDER_ID,
    POLL_SECONDS,
    HEARTBEAT_FILENAME,
    GOOGLE_TOKEN_JSON,
    SERVICE_NAME,
    VERSION,
)

from app.drive_client import build_drive_service, list_children, upsert_json_file


def _utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _make_hb_payload() -> dict:
    return {
        "service": SERVICE_NAME,
        "status": "alive",
        "last_seen": _utc_now_iso(),
        "version": VERSION,
        "root_folder_id": ROOT_FOLDER_ID,
    }


async def worker_loop():
    print("[STARTUP] Worker loop starting")

    if not GOOGLE_TOKEN_JSON or not GOOGLE_TOKEN_JSON.strip():
        print("[ERROR] Missing GOOGLE_TOKEN_JSON (HF Secret).")
        return

    if not ROOT_FOLDER_ID or not ROOT_FOLDER_ID.strip():
        print("[ERROR] Missing ROOT_FOLDER_ID (HF Variable).")
        return

    # Build Drive client once; we’ll rebuild on transient failures
    service = build_drive_service(GOOGLE_TOKEN_JSON)

    # ----------------------------
    # 1) RUN ONCE IMMEDIATELY
    # ----------------------------
    try:
        files = list_children(service, ROOT_FOLDER_ID, page_size=20)
        print(f"[DRIVE] Found {len(files)} files")
        if files:
            print("[DRIVE] Sample:", ", ".join([f["name"] for f in files[:5]]))

        upsert_json_file(service, ROOT_FOLDER_ID, HEARTBEAT_FILENAME, _make_hb_payload())
        print(f"[HEARTBEAT] Wrote {HEARTBEAT_FILENAME} to Drive root (folder={ROOT_FOLDER_ID})")

    except Exception as e:
        print(f"[ERROR] Startup drive/heartbeat failed: {e}")
        # backoff so we don't thrash on startup
        await asyncio.sleep(10)

    # ----------------------------
    # 2) POLL FOREVER
    # ----------------------------
    while True:
        try:
            files = list_children(service, ROOT_FOLDER_ID, page_size=20)
            print(f"[DRIVE] Found {len(files)} files")
            if files:
                print("[DRIVE] Sample:", ", ".join([f["name"] for f in files[:5]]))

            upsert_json_file(service, ROOT_FOLDER_ID, HEARTBEAT_FILENAME, _make_hb_payload())
            print(f"[HEARTBEAT] Updated {HEARTBEAT_FILENAME} (every {POLL_SECONDS}s)")

        except Exception as e:
            # Transient HF/network issues show up here (Broken pipe, EOF, etc.)
            print(f"[ERROR] Drive/heartbeat poll failed: {e}")

            # Retry strategy:
            # 1) Rebuild drive client (fresh TLS session)
            # 2) Retry heartbeat write once
            try:
                service = build_drive_service(GOOGLE_TOKEN_JSON)
                upsert_json_file(service, ROOT_FOLDER_ID, HEARTBEAT_FILENAME, _make_hb_payload())
                print("[HEARTBEAT] Recovered after rebuild/retry")

            except Exception as e2:
                print(f"[ERROR] Recovery retry failed: {e2}")

            # small backoff so we don't spam on repeated failures
            await asyncio.sleep(10)

        # normal poll interval
        await asyncio.sleep(POLL_SECONDS)