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
|
# Initialize logging
|
||||||
init_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"
|
USE_PUBLIC_BLOCKLIST = os.environ.get("PUBLIC_BLOCKLIST", "true").lower() == "true"
|
||||||
all_blocked = set(MANUAL_BLOCKED_IPS)
|
all_blocked = set(MANUAL_BLOCKED_IPS)
|
||||||
|
|
||||||
if USE_PUBLIC_BLOCKLIST:
|
if USE_PUBLIC_BLOCKLIST:
|
||||||
cached_ips = set()
|
cached_ips = set()
|
||||||
|
|
||||||
|
# Try to load from cache first
|
||||||
if BLOCKLIST_CACHE_FILE.exists():
|
if BLOCKLIST_CACHE_FILE.exists():
|
||||||
try:
|
try:
|
||||||
cached_ips = set(BLOCKLIST_CACHE_FILE.read_text().strip().splitlines())
|
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:
|
except Exception as e:
|
||||||
logger.warning(f"Failed to load blocklist cache: {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:
|
if not cached_ips:
|
||||||
cached_ips = mw.load_public_blocklists()
|
DATA_DIR.mkdir(exist_ok=True)
|
||||||
if cached_ips:
|
lock_file = DATA_DIR / "blocklist.lock"
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Try to acquire lock (non-blocking)
|
||||||
|
import fcntl
|
||||||
|
lock_fd = open(lock_file, 'w')
|
||||||
try:
|
try:
|
||||||
DATA_DIR.mkdir(exist_ok=True)
|
fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||||
BLOCKLIST_CACHE_FILE.write_text("\n".join(cached_ips))
|
# We got the lock - download
|
||||||
logger.info(f"Saved {len(cached_ips)} IPs to blocklist cache")
|
cached_ips = mw.load_public_blocklists()
|
||||||
except Exception as e:
|
if cached_ips:
|
||||||
logger.warning(f"Failed to save blocklist cache: {e}")
|
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)
|
all_blocked.update(cached_ips)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user