Function cleanup_task_memory

dispatcharr/celery.py:85–145

Clean up memory and database connections after each task completes

(**kwargs)

# Add memory cleanup after task completion
@task_postrun.connect  # Use the imported signal
def cleanup_task_memory(**kwargs):
    """Clean up memory and database connections after each task completes"""
    from django.db import close_old_connections

    # Get task name from kwargs
    task_name = kwargs.get('task').name if kwargs.get('task') else ''

    # Return all DB connections to the pool in a clean state
    try:
        close_old_connections()
    except Exception:
        pass

    # Only run memory cleanup for memory-intensive tasks
    memory_intensive_tasks = [
        'apps.m3u.tasks.refresh_single_m3u_account',
        'apps.m3u.tasks.refresh_m3u_accounts',
        'apps.m3u.tasks.process_m3u_batch',
        'apps.m3u.tasks.process_xc_category',
        'apps.m3u.tasks.sync_auto_channels',
        'apps.epg.tasks.refresh_epg_data',
        'apps.epg.tasks.refresh_all_epg_data',
        'apps.epg.tasks.parse_programs_for_source',
        'apps.epg.tasks.parse_programs_for_tvg_id',
        'apps.epg.tasks.build_programme_index_task',
        'apps.channels.tasks.match_epg_channels',
        'apps.channels.tasks.match_selected_channels_epg',
        'apps.channels.tasks.match_single_channel_epg',
        'core.tasks.rehash_streams',
        'apps.vod.tasks.refresh_vod_content',
        'apps.vod.tasks.batch_refresh_series_episodes',
    ]

    # Check if this is a memory-intensive task
    if task_name in memory_intensive_tasks:
        # Import cleanup_memory function
        from core.utils import cleanup_memory

        # Use the comprehensive cleanup function
        cleanup_memory(log_usage=True, force_collection=True)

        # Log memory usage if psutil is installed
        try:
            import psutil
            process = psutil.Process()
            if hasattr(process, 'memory_info'):
                mem = process.memory_info().rss / (1024 * 1024)
                print(f"Memory usage after {task_name}: {mem:.2f} MB")
        except (ImportError, Exception):
            pass
    else:
        # For non-intensive tasks, just log but don't force cleanup
        try:
            import psutil
            process = psutil.Process()
            if hasattr(process, 'memory_info'):
                mem = process.memory_info().rss / (1024 * 1024)
                if mem > 500:  # Only log if using more than 500MB