Zichao Lin
1be3aadb08
- main.py spawns 3 daemon threads: email poller, email processor, TG poller - Each thread runs independently so retries don't block other operations - _pending_uids set + lock prevents duplicate queueing across poll cycles - summarizer.py split into poll_accounts() (generator) and process_email() - Graceful shutdown via shared _running flag
107 lines
2.9 KiB
Python
107 lines
2.9 KiB
Python
import logging
|
||
import queue
|
||
import signal
|
||
import sys
|
||
import threading
|
||
import time
|
||
from src.config import load_config
|
||
from src.summarizer import poll_accounts, process_email
|
||
from src.tg_bot import poll_telegram
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||
datefmt="%Y-%m-%d %H:%M:%S",
|
||
)
|
||
logger = logging.getLogger("main")
|
||
|
||
_running = True
|
||
_pending_lock = threading.Lock()
|
||
_pending_uids: set[bytes] = set()
|
||
|
||
|
||
def _signal_handler(signum, frame):
|
||
global _running
|
||
logger.info("收到退出信号，正在停止...")
|
||
_running = False
|
||
|
||
|
||
def _email_poller(cfg, processing_queue):
|
||
logger.info("邮件轮询线程已启动")
|
||
while _running:
|
||
for acct_idx, mail in poll_accounts(cfg):
|
||
if not _running:
|
||
return
|
||
with _pending_lock:
|
||
if mail.uid not in _pending_uids:
|
||
_pending_uids.add(mail.uid)
|
||
processing_queue.put((acct_idx, mail))
|
||
if _running:
|
||
for _ in range(cfg.polling.interval_seconds):
|
||
if not _running:
|
||
return
|
||
time.sleep(1)
|
||
|
||
|
||
def _email_processor(cfg, processing_queue):
|
||
logger.info("邮件处理线程已启动")
|
||
while _running:
|
||
try:
|
||
acct_idx, mail = processing_queue.get(timeout=1)
|
||
except queue.Empty:
|
||
continue
|
||
try:
|
||
process_email(cfg, acct_idx, mail)
|
||
except Exception as e:
|
||
logger.error(f"处理邮件失败: {e}", exc_info=True)
|
||
finally:
|
||
with _pending_lock:
|
||
_pending_uids.discard(mail.uid)
|
||
|
||
|
||
def _tg_poller(cfg):
|
||
logger.info("TG 轮询线程已启动")
|
||
last_update_id = 0
|
||
while _running:
|
||
try:
|
||
last_update_id = poll_telegram(cfg.telegram, cfg, last_update_id)
|
||
except Exception as e:
|
||
logger.error(f"TG 轮询错误: {e}", exc_info=True)
|
||
if _running:
|
||
time.sleep(1)
|
||
|
||
|
||
def main():
|
||
global _running
|
||
signal.signal(signal.SIGINT, _signal_handler)
|
||
signal.signal(signal.SIGTERM, _signal_handler)
|
||
|
||
cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
|
||
cfg = load_config(cfg_path)
|
||
logger.info(f"AI邮件摘要机器人已启动，轮询间隔: {cfg.polling.interval_seconds}s")
|
||
|
||
processing_queue: queue.Queue = queue.Queue()
|
||
|
||
threads = [
|
||
threading.Thread(target=_email_poller, args=(cfg, processing_queue), daemon=True),
|
||
threading.Thread(target=_email_processor, args=(cfg, processing_queue), daemon=True),
|
||
threading.Thread(target=_tg_poller, args=(cfg,), daemon=True),
|
||
]
|
||
|
||
for t in threads:
|
||
t.start()
|
||
|
||
try:
|
||
while _running:
|
||
time.sleep(1)
|
||
except KeyboardInterrupt:
|
||
_running = False
|
||
|
||
for t in threads:
|
||
t.join(timeout=5)
|
||
logger.info("机器人已停止")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|