jd-tb-yaofang-2/commons/scheduler.py

import time
import requests
from commons.Logger import get_spider_logger
import threading
from commons.config import (
    SCHEDULE_API_BASE,
    CRAWLER_TOKEN,
    SCHEDULE_HEARTBEAT_TIMEOUT,
    SCHEDULE_PULL_TIMEOUT,
    SCHEDULE_REPORT_TIMEOUT,
    SCHEDULE_HEARTBEAT_INTERVAL,
)

logger = get_spider_logger('scheduler')


class CrawlerScheduler:
    """爬虫任务调度器"""

    def __init__(self, DEVICE_ID, platform, heartbeat_interval=30):
        """
        初始化调度器

        Args:
            platform: 平台名称
            heartbeat_url: 心跳上报URL
            heartbeat_interval: 心跳间隔时间（秒）
        """

        self.username = DEVICE_ID
        self.platform = platform

        self._lock = threading.Lock()

        self.heartbeat_url = f'{SCHEDULE_API_BASE}/api/collect_task/heartbeat'
        self.heartbeat_interval = heartbeat_interval or SCHEDULE_HEARTBEAT_INTERVAL
        self.end = False
        self.heartbeat_thread = None

    def _heartbeat_reporter(self):
        """守护线程：只负责上报心跳"""
        headers = {'X-Crawler-Token': CRAWLER_TOKEN}
        while True:
            try:

                response = requests.post(
                    self.heartbeat_url,
                    json={
                        "platform": self.platform,
                        "username": self.username
                    },
                    headers=headers,
                    timeout=SCHEDULE_HEARTBEAT_TIMEOUT
                )

                print(f"[心跳] 发送成功: {response.status_code}")
                result = response.json()

                if self.end or result.get('code') != 'success':
                    logger.error(f'心跳回传：{result}')
                    self.set_flag(True)
                    # 写日志
                    break

                logger.info(f'心跳回传：{result}')
                self.set_flag(False)

                time.sleep(self.heartbeat_interval)

            except Exception as e:
                print(e)
                logger.error(e)
                time.sleep(5)

    def start(self):
        """启动调度器"""
        self.set_flag(False)
        self.heartbeat_thread = threading.Thread(
            target=self._heartbeat_reporter,
            daemon=True
        )
        self.heartbeat_thread.start()

    def stop(self):
        """停止调度器"""
        self.set_flag(True)
        logger.info('心跳停止')

    def get_task(self):
        try:
            task_api = f"{SCHEDULE_API_BASE}/api/collect_task/pull"
            headers = {'X-Crawler-Token': CRAWLER_TOKEN}
            params = {
                'platform': self.platform,
                'username': self.username
            }
            response = requests.get(task_api, params=params, headers=headers, timeout=SCHEDULE_PULL_TIMEOUT)
            result = response.json()
            logger.info(f'拉取任务：{result}')
            if result.get('code') == 'success':
                return result.get('data').get('task')

            print('拉取任务返回', result)

        except Exception as e:
            logger.error('获取任务报错', e)

    def post_report(self, data):
        try:
            url = f"{SCHEDULE_API_BASE}/api/collect_task/report"
            print('传给返回接口的数据', data)
            logger.info(f'report上传数据：{data}')
            headers = {'X-Crawler-Token': CRAWLER_TOKEN}
            response = requests.post(url, json=data, headers=headers, timeout=SCHEDULE_REPORT_TIMEOUT)
            # 记录日志
            result = response.json()

            logger.info(result)

            if (result.get('code') != 'success'):
                logger.error(f'翻页回传结果不成功：{result}')
                self.stop()
            print(f'任务进度上传 {result}')
        except Exception as e:
            logger.error(e)

    def set_flag(self, value):
        with self._lock:
            self.end = value