| 14 | from botasaurus.dontcache import is_dont_cache |
| 15 | |
| 16 | class TaskExecutor: |
| 17 | |
| 18 | def load(self): |
| 19 | self.current_capacity = {"browser": 0, "request": 0, "task": 0} |
| 20 | self.lock = Lock() |
| 21 | |
| 22 | def start(self): |
| 23 | self.fix_in_progress_tasks() |
| 24 | Thread(target=self.task_worker, daemon=True).start() |
| 25 | |
| 26 | @retry_on_db_error |
| 27 | def fix_in_progress_tasks(self): |
| 28 | with Session() as session: |
| 29 | # Delete tasks with is_sync=True and status as either pending or in progress |
| 30 | session.query(Task).filter( |
| 31 | and_( |
| 32 | Task.is_sync == True, |
| 33 | Task.status.in_([TaskStatus.PENDING, TaskStatus.IN_PROGRESS]), |
| 34 | ) |
| 35 | ).delete() |
| 36 | # Update in progress tasks to pending |
| 37 | session.query(Task).filter(Task.status == TaskStatus.IN_PROGRESS).update( |
| 38 | {"status": TaskStatus.PENDING, "started_at": None, "finished_at": None,}, |
| 39 | ) |
| 40 | session.commit() |
| 41 | |
| 42 | def task_worker(self): |
| 43 | browser_scrapers = len(Server.get_browser_scrapers()) > 0 |
| 44 | request_scrapers = len(Server.get_request_scrapers()) > 0 |
| 45 | task_scrapers = len(Server.get_task_scrapers()) > 0 |
| 46 | |
| 47 | while True: |
| 48 | try: |
| 49 | if browser_scrapers: |
| 50 | self.process_tasks(ScraperType.BROWSER) |
| 51 | |
| 52 | if request_scrapers: |
| 53 | self.process_tasks(ScraperType.REQUEST) |
| 54 | |
| 55 | if task_scrapers: |
| 56 | self.process_tasks(ScraperType.TASK) |
| 57 | except: |
| 58 | traceback.print_exc() |
| 59 | |
| 60 | time.sleep(1) |
| 61 | |
| 62 | def process_tasks(self, scraper_type): |
| 63 | rate_limit = self.get_max_running_count(scraper_type) |
| 64 | current_count = self.get_current_running_count(scraper_type) |
| 65 | if current_count < rate_limit: |
| 66 | self.execute_pending_tasks( |
| 67 | and_( |
| 68 | Task.status == TaskStatus.PENDING, |
| 69 | Task.scraper_type == scraper_type, |
| 70 | Task.is_all_task == False, |
| 71 | ), |
| 72 | current_count, |
| 73 | rate_limit, |