mirror of
https://github.com/taylanbakircioglu/flowfish.git
synced 2026-09-17 08:05:09 +00:00
taylanbakircioglu
d7ca50b387
Multi-cluster dependency mapping, real-time network monitoring, impact analysis, and CI/CD integration capabilities. Made-with: Cursor
390 lines
15 KiB
Python
390 lines
15 KiB
Python
"""
|
|
Scheduled Simulation Worker - Background task for automatic simulation execution
|
|
|
|
This worker runs as a background task and periodically checks for scheduled
|
|
simulations that are due, executing them automatically.
|
|
|
|
Features:
|
|
- Periodic check for due simulations (every 30 seconds)
|
|
- Automatic execution when scheduled_time <= NOW()
|
|
- Support for recurring schedules (once, daily, weekly)
|
|
- Notification before execution (notify_before_minutes)
|
|
- Status updates after execution
|
|
- WebSocket notifications for simulation events
|
|
|
|
Configuration (via environment variables):
|
|
- SCHEDULED_SIMULATION_ENABLED: Enable/disable worker (default: true)
|
|
- SCHEDULED_SIMULATION_CHECK_INTERVAL: Check interval in seconds (default: 30)
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
from typing import Optional, List, Dict, Any
|
|
from datetime import datetime, timedelta
|
|
import structlog
|
|
|
|
from database.postgresql import database
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
|
|
class ScheduledSimulationWorker:
|
|
"""
|
|
Background worker that automatically executes scheduled simulations.
|
|
|
|
Features:
|
|
- Checks for due simulations every 30 seconds
|
|
- Executes simulations when scheduled_time <= NOW()
|
|
- Updates status after execution
|
|
- Handles recurring schedules (daily, weekly)
|
|
- Sends notifications before execution
|
|
"""
|
|
|
|
def __init__(self):
|
|
# Configuration from environment
|
|
self.CHECK_INTERVAL = int(os.getenv("SCHEDULED_SIMULATION_CHECK_INTERVAL", "30"))
|
|
self.ENABLED = os.getenv("SCHEDULED_SIMULATION_ENABLED", "true").lower() == "true"
|
|
|
|
# Internal state
|
|
self._running = False
|
|
self._task: Optional[asyncio.Task] = None
|
|
self._notification_task: Optional[asyncio.Task] = None
|
|
self._notified_simulations: set = set() # Track already notified simulations
|
|
|
|
async def start(self) -> None:
|
|
"""Start the background scheduler task"""
|
|
if not self.ENABLED:
|
|
logger.info("Scheduled simulation worker is disabled")
|
|
return
|
|
|
|
if self._running:
|
|
logger.warning("Scheduled simulation worker already running")
|
|
return
|
|
|
|
self._running = True
|
|
self._task = asyncio.create_task(self._run_scheduler_loop())
|
|
self._notification_task = asyncio.create_task(self._run_notification_loop())
|
|
|
|
logger.info(
|
|
"Scheduled simulation worker started",
|
|
check_interval=self.CHECK_INTERVAL
|
|
)
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the background scheduler task"""
|
|
self._running = False
|
|
|
|
for task in [self._task, self._notification_task]:
|
|
if task:
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
self._task = None
|
|
self._notification_task = None
|
|
self._notified_simulations.clear()
|
|
|
|
logger.info("Scheduled simulation worker stopped")
|
|
|
|
async def _run_scheduler_loop(self) -> None:
|
|
"""Main scheduler loop - checks for due simulations"""
|
|
# Initial delay to let the system stabilize
|
|
await asyncio.sleep(15)
|
|
|
|
while self._running:
|
|
try:
|
|
await self._check_and_execute_due_simulations()
|
|
except Exception as e:
|
|
logger.error("Scheduled simulation check failed", error=str(e))
|
|
|
|
await asyncio.sleep(self.CHECK_INTERVAL)
|
|
|
|
async def _run_notification_loop(self) -> None:
|
|
"""Notification loop - sends notifications before scheduled time"""
|
|
# Initial delay
|
|
await asyncio.sleep(20)
|
|
|
|
while self._running:
|
|
try:
|
|
await self._send_upcoming_notifications()
|
|
except Exception as e:
|
|
logger.error("Notification check failed", error=str(e))
|
|
|
|
# Check every minute for upcoming notifications
|
|
await asyncio.sleep(60)
|
|
|
|
async def _check_and_execute_due_simulations(self) -> None:
|
|
"""Check for simulations that are due and execute them"""
|
|
try:
|
|
# Find simulations where scheduled_time <= NOW() and status = 'scheduled'
|
|
query = """
|
|
SELECT * FROM scheduled_simulations
|
|
WHERE status = 'scheduled'
|
|
AND scheduled_time <= NOW()
|
|
ORDER BY scheduled_time ASC
|
|
LIMIT 10
|
|
"""
|
|
|
|
due_simulations = await database.fetch_all(query, {})
|
|
|
|
if not due_simulations:
|
|
return
|
|
|
|
logger.info(f"Found {len(due_simulations)} due simulations to execute")
|
|
|
|
for simulation in due_simulations:
|
|
await self._execute_simulation(simulation)
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to check due simulations", error=str(e))
|
|
|
|
async def _execute_simulation(self, simulation: dict) -> None:
|
|
"""Execute a single scheduled simulation"""
|
|
simulation_id = simulation['id']
|
|
|
|
try:
|
|
logger.info(
|
|
"Executing scheduled simulation",
|
|
simulation_id=simulation_id,
|
|
name=simulation['name'],
|
|
target=f"{simulation['target_namespace']}/{simulation['target_name']}"
|
|
)
|
|
|
|
# Import here to avoid circular imports
|
|
from schemas.simulation import ImpactSimulationRequest, ChangeType
|
|
from routers.simulation import run_impact_simulation
|
|
from services.network_policy_service import NetworkPolicyService
|
|
|
|
# Create the simulation request
|
|
target_id = simulation['target_name']
|
|
|
|
request = ImpactSimulationRequest(
|
|
cluster_id=simulation['cluster_id'],
|
|
analysis_id=simulation['analysis_id'],
|
|
target_id=target_id,
|
|
target_name=simulation['target_name'],
|
|
target_namespace=simulation['target_namespace'],
|
|
target_kind=simulation['target_kind'],
|
|
change_type=ChangeType(simulation['change_type'])
|
|
)
|
|
|
|
# Create a mock user for the automated execution
|
|
system_user = {
|
|
'user_id': 0,
|
|
'username': 'scheduler',
|
|
'roles': ['admin']
|
|
}
|
|
|
|
# Get service instance
|
|
service = NetworkPolicyService()
|
|
|
|
# Execute the simulation
|
|
result = await run_impact_simulation(request, system_user, service)
|
|
|
|
# Update simulation status based on schedule type
|
|
if simulation['schedule_type'] == 'once':
|
|
new_status = 'completed'
|
|
next_run = None
|
|
elif simulation['schedule_type'] == 'daily':
|
|
new_status = 'scheduled'
|
|
next_run = datetime.utcnow() + timedelta(days=1)
|
|
elif simulation['schedule_type'] == 'weekly':
|
|
new_status = 'scheduled'
|
|
next_run = datetime.utcnow() + timedelta(weeks=1)
|
|
else:
|
|
new_status = 'completed'
|
|
next_run = None
|
|
|
|
# Update the simulation record
|
|
if next_run:
|
|
update_query = """
|
|
UPDATE scheduled_simulations
|
|
SET last_run_at = NOW(),
|
|
last_run_result = :result,
|
|
status = :status,
|
|
scheduled_time = :next_run
|
|
WHERE id = :id
|
|
"""
|
|
await database.execute(update_query, {
|
|
'id': simulation_id,
|
|
'result': 'success' if result.success else 'failed',
|
|
'status': new_status,
|
|
'next_run': next_run
|
|
})
|
|
else:
|
|
update_query = """
|
|
UPDATE scheduled_simulations
|
|
SET last_run_at = NOW(),
|
|
last_run_result = :result,
|
|
status = :status
|
|
WHERE id = :id
|
|
"""
|
|
await database.execute(update_query, {
|
|
'id': simulation_id,
|
|
'result': 'success' if result.success else 'failed',
|
|
'status': new_status
|
|
})
|
|
|
|
logger.info(
|
|
"Scheduled simulation executed successfully",
|
|
simulation_id=simulation_id,
|
|
name=simulation['name'],
|
|
success=result.success,
|
|
affected_count=result.summary.total_affected if result.summary else 0,
|
|
new_status=new_status
|
|
)
|
|
|
|
# Send WebSocket notification
|
|
await self._send_execution_notification(simulation, result, 'success')
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
"Failed to execute scheduled simulation",
|
|
simulation_id=simulation_id,
|
|
error=str(e)
|
|
)
|
|
|
|
# Update status to failed
|
|
update_query = """
|
|
UPDATE scheduled_simulations
|
|
SET last_run_at = NOW(),
|
|
last_run_result = :result,
|
|
status = CASE
|
|
WHEN schedule_type = 'once' THEN 'failed'
|
|
ELSE status
|
|
END
|
|
WHERE id = :id
|
|
"""
|
|
await database.execute(update_query, {
|
|
'id': simulation_id,
|
|
'result': f'error: {str(e)}'
|
|
})
|
|
|
|
# Send failure notification
|
|
await self._send_execution_notification(simulation, None, 'failed', str(e))
|
|
|
|
async def _send_upcoming_notifications(self) -> None:
|
|
"""Send notifications for simulations that will run soon"""
|
|
try:
|
|
# Find simulations that will run within their notify_before_minutes window
|
|
query = """
|
|
SELECT * FROM scheduled_simulations
|
|
WHERE status = 'scheduled'
|
|
AND scheduled_time > NOW()
|
|
AND scheduled_time <= NOW() + (notify_before_minutes || ' minutes')::INTERVAL
|
|
"""
|
|
|
|
upcoming = await database.fetch_all(query, {})
|
|
|
|
for simulation in upcoming:
|
|
sim_id = simulation['id']
|
|
|
|
# Skip if already notified
|
|
if sim_id in self._notified_simulations:
|
|
continue
|
|
|
|
minutes_until = int((simulation['scheduled_time'] - datetime.utcnow()).total_seconds() / 60)
|
|
|
|
logger.info(
|
|
"Sending upcoming simulation notification",
|
|
simulation_id=sim_id,
|
|
name=simulation['name'],
|
|
minutes_until=minutes_until
|
|
)
|
|
|
|
await self._send_upcoming_notification(simulation, minutes_until)
|
|
self._notified_simulations.add(sim_id)
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to check upcoming notifications", error=str(e))
|
|
|
|
async def _send_execution_notification(
|
|
self,
|
|
simulation: dict,
|
|
result: Any,
|
|
status: str,
|
|
error: str = None
|
|
) -> None:
|
|
"""Send WebSocket notification about simulation execution"""
|
|
try:
|
|
from routers.websocket import broadcast_event
|
|
|
|
event_data = {
|
|
"type": "SCHEDULED_SIMULATION_EXECUTED",
|
|
"simulation_id": simulation['id'],
|
|
"name": simulation['name'],
|
|
"target": f"{simulation['target_namespace']}/{simulation['target_name']}",
|
|
"status": status,
|
|
"timestamp": datetime.utcnow().isoformat()
|
|
}
|
|
|
|
if result and hasattr(result, 'summary') and result.summary:
|
|
event_data["affected_count"] = result.summary.total_affected
|
|
event_data["high_impact"] = result.summary.high_impact
|
|
|
|
if error:
|
|
event_data["error"] = error
|
|
|
|
await broadcast_event("simulation", event_data)
|
|
|
|
except Exception as e:
|
|
logger.warning("Failed to send execution notification", error=str(e))
|
|
|
|
async def _send_upcoming_notification(self, simulation: dict, minutes_until: int) -> None:
|
|
"""Send WebSocket notification about upcoming simulation"""
|
|
try:
|
|
from routers.websocket import broadcast_event
|
|
|
|
event_data = {
|
|
"type": "SCHEDULED_SIMULATION_UPCOMING",
|
|
"simulation_id": simulation['id'],
|
|
"name": simulation['name'],
|
|
"target": f"{simulation['target_namespace']}/{simulation['target_name']}",
|
|
"change_type": simulation['change_type'],
|
|
"minutes_until": minutes_until,
|
|
"scheduled_time": simulation['scheduled_time'].isoformat() if simulation['scheduled_time'] else None,
|
|
"timestamp": datetime.utcnow().isoformat()
|
|
}
|
|
|
|
await broadcast_event("simulation", event_data)
|
|
|
|
except Exception as e:
|
|
logger.warning("Failed to send upcoming notification", error=str(e))
|
|
|
|
async def get_status(self) -> Dict[str, Any]:
|
|
"""Get worker status for health checks"""
|
|
try:
|
|
# Count pending simulations
|
|
query = "SELECT COUNT(*) as count FROM scheduled_simulations WHERE status = 'scheduled'"
|
|
result = await database.fetch_one(query, {})
|
|
pending_count = result['count'] if result else 0
|
|
|
|
# Count simulations due soon (next hour)
|
|
query_due = """
|
|
SELECT COUNT(*) as count FROM scheduled_simulations
|
|
WHERE status = 'scheduled'
|
|
AND scheduled_time <= NOW() + INTERVAL '1 hour'
|
|
"""
|
|
result_due = await database.fetch_one(query_due, {})
|
|
due_soon_count = result_due['count'] if result_due else 0
|
|
|
|
return {
|
|
"enabled": self.ENABLED,
|
|
"running": self._running,
|
|
"check_interval": self.CHECK_INTERVAL,
|
|
"pending_simulations": pending_count,
|
|
"due_within_hour": due_soon_count
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"enabled": self.ENABLED,
|
|
"running": self._running,
|
|
"error": str(e)
|
|
}
|
|
|
|
|
|
# Global worker instance
|
|
scheduled_simulation_worker = ScheduledSimulationWorker()
|