114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
SMTP Connection Pool with robust connection handling
|
|
"""
|
|
|
|
import smtplib
|
|
from queue import Queue, Empty
|
|
from typing import Optional
|
|
|
|
from logger import log
|
|
from config import config
|
|
|
|
|
|
class SMTPPool:
|
|
"""Thread-safe SMTP Connection Pool"""
|
|
|
|
def __init__(self, host: str, port: int, pool_size: int = 5):
|
|
self.host = host
|
|
self.port = port
|
|
self.pool_size = pool_size
|
|
self._pool: Queue = Queue(maxsize=pool_size)
|
|
self._initialized = False
|
|
|
|
def _create_connection(self) -> Optional[smtplib.SMTP]:
|
|
"""Create new SMTP connection"""
|
|
try:
|
|
conn = smtplib.SMTP(self.host, self.port, timeout=30)
|
|
conn.ehlo()
|
|
if config.smtp_use_tls:
|
|
conn.starttls()
|
|
conn.ehlo()
|
|
if config.smtp_user and config.smtp_pass:
|
|
conn.login(config.smtp_user, config.smtp_pass)
|
|
log(f" 📡 New SMTP connection created to {self.host}:{self.port}")
|
|
return conn
|
|
except Exception as e:
|
|
log(f"Failed to create SMTP connection: {e}", 'ERROR')
|
|
return None
|
|
|
|
def _test_connection(self, conn: smtplib.SMTP) -> bool:
|
|
"""Test if connection is still alive"""
|
|
try:
|
|
status = conn.noop()[0]
|
|
return status == 250
|
|
except Exception:
|
|
return False
|
|
|
|
def initialize(self):
|
|
"""Pre-create connections"""
|
|
if self._initialized:
|
|
return
|
|
|
|
# Only 1-2 connections initially, rest on-demand
|
|
for _ in range(min(2, self.pool_size)):
|
|
conn = self._create_connection()
|
|
if conn:
|
|
self._pool.put(conn)
|
|
|
|
self._initialized = True
|
|
log(f"SMTP pool initialized with {self._pool.qsize()} connections (max: {self.pool_size})")
|
|
|
|
def get_connection(self, timeout: float = 5.0) -> Optional[smtplib.SMTP]:
|
|
"""Get a valid connection from pool or create new one"""
|
|
# Try to get from pool
|
|
try:
|
|
conn = self._pool.get(block=False)
|
|
# Test if connection is still alive
|
|
if self._test_connection(conn):
|
|
return conn
|
|
else:
|
|
# Connection is dead, close and create new one
|
|
try:
|
|
conn.quit()
|
|
except:
|
|
pass
|
|
log(f" ♻ Recycled stale SMTP connection")
|
|
return self._create_connection()
|
|
except Empty:
|
|
# Pool empty, create new connection
|
|
return self._create_connection()
|
|
|
|
def return_connection(self, conn: smtplib.SMTP):
|
|
"""Return connection to pool if still valid"""
|
|
if conn is None:
|
|
return
|
|
|
|
# Check if connection is still good
|
|
if not self._test_connection(conn):
|
|
try:
|
|
conn.quit()
|
|
except:
|
|
pass
|
|
log(f" 🗑 Discarded broken SMTP connection")
|
|
return
|
|
|
|
# Try to return to pool
|
|
try:
|
|
self._pool.put_nowait(conn)
|
|
except:
|
|
# Pool full, close connection
|
|
try:
|
|
conn.quit()
|
|
except:
|
|
pass
|
|
|
|
def close_all(self):
|
|
"""Close all connections"""
|
|
while not self._pool.empty():
|
|
try:
|
|
conn = self._pool.get_nowait()
|
|
conn.quit()
|
|
except:
|
|
pass
|