Update worker.py

worker.py

@@ -1,50 +1,73 @@
 import asyncio
-import db
+import random
+from state import state
 import scraper
-import traceback
+from logger import get_logger
 
-async def start_worker(worker_id: int):
-    print(f"⚙️ Worker {worker_id} initialized.")
-    r = await db.get_redis()
+log = get_logger()
+workers_active = []
+
+async def worker_task(worker_id: int):
+    log.info(f"Worker {worker_id} started.")
+    await state.init_queue()
     
     while True:
         try:
-            state = await db.get_state()
-            if state.get("running") != "1" or state.get("paused") == "1":
-                await asyncio.sleep(2)
+            # Handle STOP state
+            if state.status == "STOPPED":
+                break
+                
+            # Handle PAUSE state
+            if state.status == "PAUSED":
+                await asyncio.sleep(1)
                 continue
 
-            word = await r.rpop("frag:queue")
-            if not word:
-                await asyncio.sleep(2)
-                continue
+            # Try to get a word from the in-memory queue
+            try:
+                word = state.queue.get_nowait()
+            except asyncio.QueueEmpty:
+                log.info(f"Worker {worker_id} finished. Queue empty.")
+                break
 
-            proxy = await db.get_random_proxy()
+            # Pick a proxy if available
+            proxy = random.choice(state.proxies) if state.proxies else None
+            
+            # Execute scrape
             status = await scraper.check_fragment(word, proxy_url=proxy)
             
             if status == "ERROR":
-                await r.lpush("frag:queue", word)
+                # Put it back to try again later
+                await state.queue.put(word)
                 await asyncio.sleep(1.5)
-                continue
-
-            if status:
-                await db.mark_done(word, status)
+            elif status:
+                await state.save_result(word, status)
+                
+            state.queue.task_done()
 
         except Exception as e:
-            print(f"❌ Worker {worker_id} Error: {str(e)}")
-            traceback.print_exc()
+            log.exception(f"Worker {worker_id} crashed:")
             await asyncio.sleep(2)
 
-async def run_worker():
-    print("🚀 Sniper Engine Starting...")
-    concurrency = await db.get_concurrency()
+async def manage_workers():
+    """Monitors state and spins up workers as needed."""
+    global workers_active
     
-    tasks = []
-    for i in range(concurrency):
-        task = asyncio.create_task(start_worker(i))
-        tasks.append(task)
+    while True:
+        if state.status == "RUNNING" and len(workers_active) < state.concurrency:
+            # Spin up missing workers
+            needed = state.concurrency - len(workers_active)
+            for i in range(needed):
+                t = asyncio.create_task(worker_task(len(workers_active) + i))
+                workers_active.append(t)
         
-    await asyncio.gather(*tasks, return_exceptions=True)
-
-if __name__ == "__main__":
-    asyncio.run(run_worker())
+        elif state.status == "STOPPED" and workers_active:
+            # Cancel all workers if stopped
+            for t in workers_active:
+                t.cancel()
+            workers_active.clear()
+            log.info("All workers forcefully stopped.")
+            
+        # Clean up finished tasks from the active list
+        workers_active = [t for t in workers_active if not t.done()]
+        
+        await asyncio.sleep(1)