Lifecycle Management

Flokkit Flow provides comprehensive lifecycle hooks for managing resources, handling errors, and ensuring clean shutdown of your pipelines.

Lifecycle Hooks Overview

Every node in Flokkit Flow can implement four lifecycle methods:

class MyNode:
    async def on_start(self) -> None:
        """Called when the node starts."""
        pass
    
    async def on_stop(self) -> None:
        """Called when the node stops."""
        pass
    
    async def on_error(self, error: Exception) -> None:
        """Called when an error occurs."""
        pass
    
    async def on_complete(self) -> None:
        """Called when processing completes normally."""
        pass

Resource Management

Database Connections

class DatabaseProcessor:
    def __init__(self, connection_string: str):
        self.connection_string = connection_string
        self.pool = None
    
    async def on_start(self):
        """Initialize connection pool on startup."""
        self.pool = await asyncpg.create_pool(
            self.connection_string,
            min_size=5,
            max_size=20
        )
        logger.info("Database pool created")
    
    async def on_stop(self):
        """Close connection pool on shutdown."""
        if self.pool:
            await self.pool.close()
            logger.info("Database pool closed")
    
    async def process(self, record: dict) -> dict:
        async with self.pool.acquire() as conn:
            result = await conn.fetchrow(
                "SELECT * FROM process_record($1)",
                record['id']
            )
            return dict(result)

# Usage
async def run_db_pipeline():
    builder = flow("Database Pipeline")
    processor = DatabaseProcessor("postgresql://localhost/db")
    
    (builder
        .source(generate_records)
        .to(processor.process)
        .to(save_results))
    
    await builder.run()

File Handles

class FileWriter:
    def __init__(self, filename: str):
        self.filename = filename
        self.file = None
        self.buffer = []
        self.buffer_size = 100
    
    async def on_start(self):
        """Open file for writing."""
        self.file = open(self.filename, 'w')
        logger.info(f"Opened file: {self.filename}")
    
    async def on_stop(self):
        """Flush buffer and close file."""
        if self.buffer:
            self._flush_buffer()
        
        if self.file:
            self.file.close()
            logger.info(f"Closed file: {self.filename}")
    
    async def on_complete(self):
        """Called when all data processed successfully."""
        logger.info(f"Successfully wrote all data to {self.filename}")
    
    def write(self, data: str) -> None:
        self.buffer.append(data)
        if len(self.buffer) >= self.buffer_size:
            self._flush_buffer()
    
    def _flush_buffer(self):
        if self.file and self.buffer:
            self.file.write('\n'.join(self.buffer) + '\n')
            self.file.flush()
            self.buffer.clear()

External Services

class APIClient:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.session = None
        self.rate_limiter = None
    
    async def on_start(self):
        """Initialize HTTP session and rate limiter."""
        self.session = aiohttp.ClientSession(
            headers={"Authorization": f"Bearer {self.api_key}"}
        )
        self.rate_limiter = RateLimiter(calls_per_second=10)
        
        # Test connection
        async with self.session.get("/health") as response:
            if response.status != 200:
                raise ConnectionError("API health check failed")
    
    async def on_stop(self):
        """Close HTTP session."""
        if self.session:
            await self.session.close()
    
    async def on_error(self, error: Exception):
        """Handle API errors with exponential backoff."""
        if isinstance(error, aiohttp.ClientError):
            logger.error(f"API error: {error}")
            # Could implement retry logic here
    
    async def fetch_data(self, item_id: str) -> dict:
        await self.rate_limiter.acquire()
        
        async with self.session.get(f"/items/{item_id}") as response:
            response.raise_for_status()
            return await response.json()

Error Handling

Graceful Degradation

class ResilientProcessor:
    def __init__(self):
        self.primary_service = None
        self.fallback_service = None
        self.use_fallback = False
        self.error_count = 0
        self.error_threshold = 5
    
    async def on_start(self):
        """Initialize both primary and fallback services."""
        self.primary_service = await PrimaryService.connect()
        self.fallback_service = await FallbackService.connect()
    
    async def on_error(self, error: Exception):
        """Switch to fallback service on repeated errors."""
        self.error_count += 1
        logger.error(f"Error #{self.error_count}: {error}")
        
        if self.error_count >= self.error_threshold:
            logger.warning("Switching to fallback service")
            self.use_fallback = True
            self.error_count = 0  # Reset for fallback monitoring
    
    async def process(self, data: dict) -> dict:
        service = self.fallback_service if self.use_fallback else self.primary_service
        
        try:
            result = await service.process(data)
            # Reset error count on success
            if not self.use_fallback:
                self.error_count = 0
            return result
        except Exception as e:
            await self.on_error(e)
            if self.use_fallback:
                # Even fallback failed
                raise
            # Retry with same data
            return await self.process(data)

Error Recovery

class RecoverableProcessor:
    def __init__(self):
        self.checkpoint_file = "checkpoint.json"
        self.processed_ids = set()
    
    async def on_start(self):
        """Load checkpoint on startup."""
        if os.path.exists(self.checkpoint_file):
            with open(self.checkpoint_file, 'r') as f:
                data = json.load(f)
                self.processed_ids = set(data['processed_ids'])
                logger.info(f"Resumed from checkpoint: {len(self.processed_ids)} items")
    
    async def on_stop(self):
        """Save checkpoint on shutdown."""
        await self._save_checkpoint()
    
    async def on_error(self, error: Exception):
        """Save checkpoint on error for recovery."""
        logger.error(f"Error occurred: {error}")
        await self._save_checkpoint()
    
    async def _save_checkpoint(self):
        with open(self.checkpoint_file, 'w') as f:
            json.dump({
                'processed_ids': list(self.processed_ids),
                'timestamp': datetime.now().isoformat()
            }, f)
    
    async def process(self, item: dict) -> Optional[dict]:
        if item['id'] in self.processed_ids:
            logger.info(f"Skipping already processed: {item['id']}")
            return None
        
        result = await self._process_item(item)
        self.processed_ids.add(item['id'])
        
        # Periodic checkpoint
        if len(self.processed_ids) % 100 == 0:
            await self._save_checkpoint()
        
        return result

Shutdown Strategies

Graceful Shutdown

class GracefulPipeline:
    def __init__(self):
        self.shutdown_event = asyncio.Event()
        self.active_tasks = set()
    
    async def on_start(self):
        """Set up signal handlers."""
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
    
    def _signal_handler(self, signum, frame):
        logger.info(f"Received signal {signum}, initiating graceful shutdown")
        self.shutdown_event.set()
    
    async def process_with_tracking(self, item: dict) -> dict:
        task_id = id(asyncio.current_task())
        self.active_tasks.add(task_id)
        
        try:
            # Check for shutdown before processing
            if self.shutdown_event.is_set():
                logger.info("Skipping new item due to shutdown")
                return None
            
            result = await self._long_running_process(item)
            return result
        finally:
            self.active_tasks.discard(task_id)
    
    async def on_stop(self):
        """Wait for active tasks to complete."""
        if self.active_tasks:
            logger.info(f"Waiting for {len(self.active_tasks)} active tasks")
            # Wait up to 30 seconds for tasks to complete
            for _ in range(30):
                if not self.active_tasks:
                    break
                await asyncio.sleep(1)
            
            if self.active_tasks:
                logger.warning(f"{len(self.active_tasks)} tasks still active after timeout")

State Persistence

class StatefulProcessor:
    def __init__(self, state_file: str = "processor_state.pkl"):
        self.state_file = state_file
        self.state = {
            'counter': 0,
            'last_processed': None,
            'metrics': defaultdict(int)
        }
    
    async def on_start(self):
        """Load state from disk."""
        if os.path.exists(self.state_file):
            with open(self.state_file, 'rb') as f:
                self.state = pickle.load(f)
                logger.info(f"Loaded state: counter={self.state['counter']}")
    
    async def on_stop(self):
        """Persist state to disk."""
        await self._save_state()
        logger.info("State persisted successfully")
    
    async def on_complete(self):
        """Final state save and cleanup."""
        self.state['completed_at'] = datetime.now().isoformat()
        await self._save_state()
        
        # Archive state file
        archive_name = f"{self.state_file}.{datetime.now():%Y%m%d_%H%M%S}"
        shutil.copy2(self.state_file, archive_name)
    
    async def _save_state(self):
        temp_file = f"{self.state_file}.tmp"
        with open(temp_file, 'wb') as f:
            pickle.dump(self.state, f)
        # Atomic rename
        os.rename(temp_file, self.state_file)
    
    async def process(self, item: dict) -> dict:
        self.state['counter'] += 1
        self.state['last_processed'] = item['id']
        self.state['metrics'][item['type']] += 1
        
        # Periodic state save
        if self.state['counter'] % 1000 == 0:
            await self._save_state()
        
        return {'item': item, 'count': self.state['counter']}

Pipeline vs Server Mode

Pipeline Mode (auto_stop=True)

class BatchProcessor:
    def __init__(self):
        self.start_time = None
        self.items_processed = 0
    
    async def on_start(self):
        """Initialize batch processing."""
        self.start_time = time.time()
        logger.info("Starting batch processing")
    
    async def on_complete(self):
        """Report batch statistics."""
        duration = time.time() - self.start_time
        rate = self.items_processed / duration
        
        logger.info(f"Batch complete: {self.items_processed} items in {duration:.1f}s ({rate:.1f} items/s)")
        
        # Send completion notification
        await notify_completion({
            'items': self.items_processed,
            'duration': duration,
            'rate': rate
        })
    
    async def process(self, item: dict) -> dict:
        self.items_processed += 1
        return await transform_item(item)

# Usage
builder = flow("Batch Job", auto_stop=True)  # Stops when complete

Server Mode (auto_stop=False)

class StreamProcessor:
    def __init__(self):
        self.health_check_interval = 60
        self.last_health_check = None
        self.is_healthy = True
    
    async def on_start(self):
        """Start health check task."""
        asyncio.create_task(self._health_check_loop())
        logger.info("Stream processor started")
    
    async def _health_check_loop(self):
        """Periodic health checks."""
        while self.is_healthy:
            await asyncio.sleep(self.health_check_interval)
            
            # Check system health
            if not await self._check_health():
                logger.error("Health check failed")
                self.is_healthy = False
    
    async def on_error(self, error: Exception):
        """Handle streaming errors."""
        if isinstance(error, ConnectionError):
            # Try to reconnect
            await self._reconnect()
        else:
            # Log and continue
            logger.error(f"Stream error: {error}")
    
    async def process(self, event: dict) -> dict:
        if not self.is_healthy:
            raise RuntimeError("Processor unhealthy")
        
        return await handle_event(event)

# Usage
builder = flow("Stream Server", auto_stop=False)  # Runs forever

Best Practices

1. Always Clean Up Resources

async def on_stop(self):
    """Clean up in reverse order of creation."""
    if hasattr(self, 'writer'):
        await self.writer.close()
    if hasattr(self, 'reader'):
        await self.reader.close()
    if hasattr(self, 'connection'):
        await self.connection.close()

2. Idempotent Lifecycle Methods

async def on_start(self):
    """Make sure this can be called multiple times safely."""
    if self.connection is None:
        self.connection = await create_connection()

3. Fail Fast on Startup

async def on_start(self):
    """Validate configuration and test connections."""
    if not self.api_key:
        raise ValueError("API key required")
    
    # Test connection
    if not await self.test_connection():
        raise ConnectionError("Failed to connect to service")

4. Log Lifecycle Events

async def on_start(self):
    logger.info(f"Starting {self.__class__.__name__}")

async def on_stop(self):
    logger.info(f"Stopping {self.__class__.__name__}")

5. Handle Partial Completion

async def on_stop(self):
    """Handle both normal and abnormal shutdown."""
    if self.buffer:
        logger.warning(f"Flushing {len(self.buffer)} buffered items")
        await self.flush_buffer()

Next Steps