怎样实现任务分发防重复脚本

本文目录导读：

1. 基于Redis的防重复方案（最推荐）
2. 基于数据库的防重复方案
3. 基于文件系统的防重复方案
4. 分布式环境下的复杂防重复方案
5. 使用建议

我来提供一个完整的任务分发防重复脚本实现方案,包含多种策略：

基于Redis的防重复方案（最推荐）

import redis
import hashlib
import json
import time
from functools import wraps
class TaskDistributor:
    def __init__(self, redis_host='localhost', redis_port=6379, db=0):
        self.redis = redis.Redis(
            host=redis_host, 
            port=redis_port, 
            db=db,
            decode_responses=True
        )
        self.lock_prefix = "task:lock:"
        self.task_prefix = "task:data:"
        self.lock_timeout = 3600  # 1小时锁超时
    def _generate_task_id(self, task_data):
        """生成任务唯一ID"""
        task_str = json.dumps(task_data, sort_keys=True)
        return hashlib.md5(task_str.encode()).hexdigest()
    def try_acquire_lock(self, task_id, ttl=None):
        """尝试获取任务锁"""
        if ttl is None:
            ttl = self.lock_timeout
        lock_key = f"{self.lock_prefix}{task_id}"
        # 使用setnx实现原子操作
        acquired = self.redis.setnx(lock_key, str(time.time()))
        if acquired:
            # 设置过期时间防止死锁
            self.redis.expire(lock_key, ttl)
            return True
        return False
    def distribute_task(self, task_data):
        """分发任务（带防重复）"""
        task_id = self._generate_task_id(task_data)
        # 尝试获取锁
        if not self.try_acquire_lock(task_id):
            return {
                'status': 'duplicate',
                'task_id': task_id,
                'message': '任务已存在'
            }
        # 存储任务数据
        task_key = f"{self.task_prefix}{task_id}"
        self.redis.setex(task_key, self.lock_timeout, json.dumps(task_data))
        # 添加到任务队列
        self.redis.lpush('task_queue', task_id)
        return {
            'status': 'success',
            'task_id': task_id,
            'message': '任务分发成功'
        }
    def release_lock(self, task_id):
        """释放任务锁"""
        lock_key = f"{self.lock_prefix}{task_id}"
        # 使用Lua脚本保证原子性
        script = """
        if redis.call('get', KEYS[1]) == ARGV[1] then
            return redis.call('del', KEYS[1])
        else
            return 0
        end
        """
        self.redis.eval(script, 1, lock_key, str(time.time()))
    def mark_completed(self, task_id):
        """标记任务完成"""
        completed_key = f"task:completed:{task_id}"
        self.redis.setex(completed_key, 86400 * 7, time.time())  # 保存7天
        self.release_lock(task_id)
# 使用示例
distributor = TaskDistributor()
# 分发任务
task = {'type': 'data_sync', 'params': {'table': 'users', 'date': '2024-01-15'}}
result = distributor.distribute_task(task)
print(result)  # {'status': 'success', 'task_id': 'xxx', ...}
# 重复分发相同任务
result = distributor.distribute_task(task)
print(result)  # {'status': 'duplicate', 'task_id': 'xxx', ...}

基于数据库的防重复方案

import sqlite3
import hashlib
import json
from datetime import datetime
from contextlib import contextmanager
class DatabaseTaskDistributor:
    def __init__(self, db_path='tasks.db'):
        self.db_path = db_path
        self.init_database()
    def init_database(self):
        """初始化数据库表"""
        with self.get_connection() as conn:
            conn.execute('''
                CREATE TABLE IF NOT EXISTS tasks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    task_hash TEXT UNIQUE NOT NULL,
                    task_data TEXT NOT NULL,
                    status TEXT DEFAULT 'pending',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    completed_at TIMESTAMP,
                    worker_id TEXT
                )
            ''')
            conn.execute('''
                CREATE INDEX IF NOT EXISTS idx_task_hash 
                ON tasks(task_hash)
            ''')
            conn.execute('''
                CREATE INDEX IF NOT EXISTS idx_status 
                ON tasks(status)
            ''')
    @contextmanager
    def get_connection(self):
        """获取数据库连接"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        except Exception as e:
            conn.rollback()
            raise e
        finally:
            conn.close()
    def _hash_task(self, task_data):
        """生成任务哈希"""
        task_str = json.dumps(task_data, sort_keys=True)
        return hashlib.sha256(task_str.encode()).hexdigest()
    def distribute_task(self, task_data, worker_id=None):
        """分发任务（带防重复）"""
        task_hash = self._hash_task(task_data)
        with self.get_connection() as conn:
            try:
                cursor = conn.execute('''
                    INSERT INTO tasks (task_hash, task_data, worker_id)
                    VALUES (?, ?, ?)
                ''', (task_hash, json.dumps(task_data), worker_id))
                return {
                    'status': 'success',
                    'task_id': cursor.lastrowid,
                    'task_hash': task_hash
                }
            except sqlite3.IntegrityError:
                # 任务已存在
                cursor = conn.execute('''
                    SELECT id, status, completed_at 
                    FROM tasks 
                    WHERE task_hash = ?
                ''', (task_hash,))
                existing_task = cursor.fetchone()
                return {
                    'status': 'duplicate',
                    'task_id': existing_task['id'],
                    'existing_status': existing_task['status'],
                    'completed_at': existing_task['completed_at']
                }
    def get_pending_tasks(self, limit=10):
        """获取待处理任务"""
        with self.get_connection() as conn:
            cursor = conn.execute('''
                SELECT id, task_data, created_at 
                FROM tasks 
                WHERE status = 'pending' 
                ORDER BY created_at ASC 
                LIMIT ?
            ''', (limit,))
            tasks = []
            for row in cursor.fetchall():
                tasks.append({
                    'id': row['id'],
                    'task_data': json.loads(row['task_data']),
                    'created_at': row['created_at']
                })
            return tasks
    def mark_in_progress(self, task_id, worker_id):
        """标记任务进行中"""
        with self.get_connection() as conn:
            conn.execute('''
                UPDATE tasks 
                SET status = 'in_progress', worker_id = ? 
                WHERE id = ? AND status = 'pending'
            ''', (worker_id, task_id))
            return conn.rowcount > 0
    def mark_completed(self, task_id):
        """标记任务完成"""
        with self.get_connection() as conn:
            conn.execute('''
                UPDATE tasks 
                SET status = 'completed', completed_at = CURRENT_TIMESTAMP 
                WHERE id = ?
            ''', (task_id,))
# 使用示例
db_distributor = DatabaseTaskDistributor()
# 分发任务
task = {'type': 'email_send', 'to': 'user@example.com', 'template': 'welcome'}
result = db_distributor.distribute_task(task)
print(f"Task ID: {result['task_id']}, Status: {result['status']}")
# 获取待处理任务
pending_tasks = db_distributor.get_pending_tasks()
for task in pending_tasks:
    print(f"Processing task {task['id']}")
    # 处理任务...
    db_distributor.mark_completed(task['id'])

基于文件系统的防重复方案

import os
import json
import hashlib
import threading
import time
from pathlib import Path
class FileBasedTaskDistributor:
    def __init__(self, task_dir='./tasks'):
        self.task_dir = Path(task_dir)
        self.task_dir.mkdir(parents=True, exist_ok=True)
        self.lock_dir = self.task_dir / 'locks'
        self.lock_dir.mkdir(exist_ok=True)
        self.completed_dir = self.task_dir / 'completed'
        self.completed_dir.mkdir(exist_ok=True)
    def _generate_task_name(self, task_data):
        """生成任务文件名"""
        task_str = json.dumps(task_data, sort_keys=True)
        task_hash = hashlib.md5(task_str.encode()).hexdigest()
        return f"task_{task_hash}.json"
    def _acquire_file_lock(self, task_name):
        """获取文件锁"""
        lock_file = self.lock_dir / f"{task_name}.lock"
        try:
            # 尝试创建锁文件
            fd = os.open(lock_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            os.close(fd)
            return True
        except FileExistsError:
            return False
    def _release_file_lock(self, task_name):
        """释放文件锁"""
        lock_file = self.lock_dir / f"{task_name}.lock"
        try:
            os.remove(lock_file)
        except FileNotFoundError:
            pass
    def distribute_task(self, task_data):
        """分发任务"""
        task_name = self._generate_task_name(task_data)
        task_file = self.task_dir / task_name
        # 检查是否已存在
        if task_file.exists():
            return {
                'status': 'duplicate',
                'task_name': task_name,
                'message': '任务已存在'
            }
        # 检查是否已完成
        completed_file = self.completed_dir / task_name
        if completed_file.exists():
            return {
                'status': 'completed',
                'task_name': task_name,
                'message': '任务已完成'
            }
        # 尝试获取锁
        if not self._acquire_file_lock(task_name):
            return {
                'status': 'locked',
                'task_name': task_name,
                'message': '任务正在处理中'
            }
        try:
            # 写入任务文件
            with open(task_file, 'w') as f:
                json.dump({
                    'task_data': task_data,
                    'created_at': time.time(),
                    'status': 'pending'
                }, f)
            return {
                'status': 'success',
                'task_name': task_name,
                'message': '任务分发成功'
            }
        finally:
            self._release_file_lock(task_name)
    def get_pending_tasks(self):
        """获取待处理任务"""
        tasks = []
        for task_file in self.task_dir.glob('task_*.json'):
            with open(task_file, 'r') as f:
                task_data = json.load(f)
                if task_data.get('status') == 'pending':
                    tasks.append({
                        'file': task_file,
                        'data': task_data
                    })
        return tasks
    def mark_completed(self, task_name):
        """标记任务完成"""
        task_file = self.task_dir / task_name
        completed_file = self.completed_dir / task_name
        if task_file.exists():
            # 移动到completed目录
            task_file.rename(completed_file)
            return True
        return False
# 使用示例
file_distributor = FileBasedTaskDistributor()
# 分发任务
task = {'type': 'report_generate', 'report_id': 'R123', 'format': 'pdf'}
result = file_distributor.distribute_task(task)
print(f"Status: {result['status']}")
# 获取待处理任务
pending = file_distributor.get_pending_tasks()
for task_info in pending:
    print(f"Processing: {task_info['file']}")
    # 处理任务...
    file_distributor.mark_completed(task_info['file'].name)

分布式环境下的复杂防重复方案

import hashlib
import json
import time
from datetime import datetime, timedelta
class DistributedTaskManager:
    def __init__(self, redis_client=None):
        self.redis = redis_client
        self.bloom_filter_key = "task:bloom_filter"
        self.task_hash_key = "task:hash:{task_hash}"
        self.task_timeout = 3600  # 1小时
    def create_bloom_filter(self):
        """创建布隆过滤器"""
        # 使用Redis的布隆过滤器模块
        try:
            self.redis.bf().create(self.bloom_filter_key, 0.01, 1000000)
        except:
            # 如果已存在，忽略
            pass
    def check_and_add_task(self, task_data, timeout=None):
        """检查并添加任务（防重复）"""
        task_hash = self._generate_hash(task_data)
        if timeout is None:
            timeout = self.task_timeout
        # 检查布隆过滤器
        if self.redis.bf().exists(self.bloom_filter_key, task_hash):
            # 进一步精确检查
            task_key = self.task_hash_key.format(task_hash=task_hash)
            if self.redis.exists(task_key):
                return {
                    'status': 'duplicate',
                    'task_hash': task_hash,
                    'message': '任务已存在'
                }
        # 添加任务
        pipeline = self.redis.pipeline()
        pipeline.bf().add(self.bloom_filter_key, task_hash)
        task_key = self.task_hash_key.format(task_hash=task_hash)
        pipeline.setex(task_key, timeout, json.dumps(task_data))
        pipeline.execute()
        return {
            'status': 'success',
            'task_hash': task_hash,
            'message': '任务添加成功'
        }
    def _generate_hash(self, task_data):
        """生成任务哈希"""
        task_str = json.dumps(task_data, sort_keys=True)
        return hashlib.sha256(task_str.encode()).hexdigest()
# 使用示例
manager = DistributedTaskManager(redis_client=redis.Redis())
manager.create_bloom_filter()
# 批量分发任务
tasks = [
    {'type': 'sync', 'source': 'db1', 'target': 'db2'},
    {'type': 'sync', 'source': 'db1', 'target': 'db2'},  # 重复
    {'type': 'backup', 'source': 'db1'},
]
for task in tasks:
    result = manager.check_and_add_task(task)
    print(f"Task: {task}, Result: {result['status']}")

使用建议

选择策略

1. Redis方案：适合高并发、分布式场景，性能最优
2. 数据库方案：适合中小规模、需要持久化存储
3. 文件方案：适合单机、简单部署场景
4. 布隆过滤器方案：适合海量任务、内存敏感场景

最佳实践

# 组合使用多种防重复机制
class HybridTaskDistributor:
    def __init__(self):
        self.redis_distributor = TaskDistributor()
        self.db_distributor = DatabaseTaskDistributor()
    def distribute_with_guarantee(self, task_data):
        """带多重保障的任务分发"""
        # 1. 首先尝试Redis（高性能）
        result = self.redis_distributor.distribute_task(task_data)
        if result['status'] == 'success':
            # 2. 额外写入数据库（持久化保障）
            db_result = self.db_distributor.distribute_task(task_data)
            result['db_task_id'] = db_result.get('task_id')
        return result

这个方案提供了多种防重复实现,你可以根据实际需求选择合适的方案。