Server: Add file lock for blocklist loading
- Only one worker downloads blocklist - Other workers wait and read from cache - Prevents duplicate downloads on startup
This commit is contained in:
+35
-10
@@ -53,29 +53,54 @@ async def lifespan(app: FastAPI):
|
||||
# Initialize logging
|
||||
init_logging()
|
||||
|
||||
# Load public blocklists (use cache if available)
|
||||
# Load public blocklists (single worker loads, others wait for cache)
|
||||
USE_PUBLIC_BLOCKLIST = os.environ.get("PUBLIC_BLOCKLIST", "true").lower() == "true"
|
||||
all_blocked = set(MANUAL_BLOCKED_IPS)
|
||||
|
||||
if USE_PUBLIC_BLOCKLIST:
|
||||
cached_ips = set()
|
||||
|
||||
# Try to load from cache first
|
||||
if BLOCKLIST_CACHE_FILE.exists():
|
||||
try:
|
||||
cached_ips = set(BLOCKLIST_CACHE_FILE.read_text().strip().splitlines())
|
||||
logger.info(f"Loaded {len(cached_ips)} IPs from blocklist cache")
|
||||
if cached_ips:
|
||||
logger.info(f"Loaded {len(cached_ips)} IPs from blocklist cache")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load blocklist cache: {e}")
|
||||
cached_ips = set()
|
||||
|
||||
# If no cache, download (only one worker will do this)
|
||||
if not cached_ips:
|
||||
cached_ips = mw.load_public_blocklists()
|
||||
if cached_ips:
|
||||
DATA_DIR.mkdir(exist_ok=True)
|
||||
lock_file = DATA_DIR / "blocklist.lock"
|
||||
|
||||
try:
|
||||
# Try to acquire lock (non-blocking)
|
||||
import fcntl
|
||||
lock_fd = open(lock_file, 'w')
|
||||
try:
|
||||
DATA_DIR.mkdir(exist_ok=True)
|
||||
BLOCKLIST_CACHE_FILE.write_text("\n".join(cached_ips))
|
||||
logger.info(f"Saved {len(cached_ips)} IPs to blocklist cache")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save blocklist cache: {e}")
|
||||
fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
# We got the lock - download
|
||||
cached_ips = mw.load_public_blocklists()
|
||||
if cached_ips:
|
||||
BLOCKLIST_CACHE_FILE.write_text("\n".join(cached_ips))
|
||||
logger.info(f"Downloaded and saved {len(cached_ips)} IPs to blocklist cache")
|
||||
except BlockingIOError:
|
||||
# Another process is downloading - wait for cache
|
||||
pass
|
||||
finally:
|
||||
lock_fd.close()
|
||||
except Exception as e:
|
||||
logger.warning(f"Lock error: {e}")
|
||||
|
||||
# Re-read cache after download
|
||||
if BLOCKLIST_CACHE_FILE.exists() and not cached_ips:
|
||||
try:
|
||||
cached_ips = set(BLOCKLIST_CACHE_FILE.read_text().strip().splitlines())
|
||||
if cached_ips:
|
||||
logger.info(f"Loaded {len(cached_ips)} IPs from blocklist cache (after wait)")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
all_blocked.update(cached_ips)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user