4936e5cd73
- Introduced `BaseDataCollector` and `CollectorManager` classes for standardized data collection and centralized management. - Added health monitoring features, including auto-restart capabilities and detailed status reporting for collectors. - Updated `env.template` to include new logging and health check configurations. - Enhanced documentation in `docs/data_collectors.md` to provide comprehensive guidance on the new data collection system. - Added unit tests for `BaseDataCollector` and `CollectorManager` to ensure reliability and functionality.
309 lines
10 KiB
Python
309 lines
10 KiB
Python
"""
|
||
Demonstration of the enhanced data collector system with health monitoring and auto-restart.
|
||
|
||
This example shows how to:
|
||
1. Create data collectors with health monitoring
|
||
2. Use the collector manager for coordinated management
|
||
3. Monitor collector health and handle failures
|
||
4. Enable/disable collectors dynamically
|
||
"""
|
||
|
||
import asyncio
|
||
from datetime import datetime, timezone
|
||
from typing import Any, Optional
|
||
|
||
from data import (
|
||
BaseDataCollector, DataType, CollectorStatus, MarketDataPoint,
|
||
CollectorManager, CollectorConfig
|
||
)
|
||
|
||
|
||
class DemoDataCollector(BaseDataCollector):
|
||
"""
|
||
Demo implementation of a data collector for demonstration purposes.
|
||
|
||
This collector simulates receiving market data and can be configured
|
||
to fail periodically to demonstrate auto-restart functionality.
|
||
"""
|
||
|
||
def __init__(self,
|
||
exchange_name: str,
|
||
symbols: list,
|
||
fail_every_n_messages: int = 0,
|
||
connection_delay: float = 0.1):
|
||
"""
|
||
Initialize demo collector.
|
||
|
||
Args:
|
||
exchange_name: Name of the exchange
|
||
symbols: Trading symbols to collect
|
||
fail_every_n_messages: Simulate failure every N messages (0 = no failures)
|
||
connection_delay: Simulated connection delay
|
||
"""
|
||
super().__init__(exchange_name, symbols, [DataType.TICKER])
|
||
self.fail_every_n_messages = fail_every_n_messages
|
||
self.connection_delay = connection_delay
|
||
self.message_count = 0
|
||
self.connected = False
|
||
self.subscribed = False
|
||
|
||
async def connect(self) -> bool:
|
||
"""Simulate connection to exchange."""
|
||
print(f"[{self.exchange_name}] Connecting...")
|
||
await asyncio.sleep(self.connection_delay)
|
||
self.connected = True
|
||
print(f"[{self.exchange_name}] Connected successfully")
|
||
return True
|
||
|
||
async def disconnect(self) -> None:
|
||
"""Simulate disconnection from exchange."""
|
||
print(f"[{self.exchange_name}] Disconnecting...")
|
||
await asyncio.sleep(self.connection_delay / 2)
|
||
self.connected = False
|
||
self.subscribed = False
|
||
print(f"[{self.exchange_name}] Disconnected")
|
||
|
||
async def subscribe_to_data(self, symbols: list, data_types: list) -> bool:
|
||
"""Simulate subscription to data streams."""
|
||
if not self.connected:
|
||
return False
|
||
|
||
print(f"[{self.exchange_name}] Subscribing to {len(symbols)} symbols: {', '.join(symbols)}")
|
||
await asyncio.sleep(0.05)
|
||
self.subscribed = True
|
||
return True
|
||
|
||
async def unsubscribe_from_data(self, symbols: list, data_types: list) -> bool:
|
||
"""Simulate unsubscription from data streams."""
|
||
print(f"[{self.exchange_name}] Unsubscribing from data streams")
|
||
self.subscribed = False
|
||
return True
|
||
|
||
async def _process_message(self, message: Any) -> Optional[MarketDataPoint]:
|
||
"""Process simulated market data message."""
|
||
self.message_count += 1
|
||
|
||
# Simulate periodic failures if configured
|
||
if (self.fail_every_n_messages > 0 and
|
||
self.message_count % self.fail_every_n_messages == 0):
|
||
raise Exception(f"Simulated failure after {self.message_count} messages")
|
||
|
||
# Create mock market data
|
||
data_point = MarketDataPoint(
|
||
exchange=self.exchange_name,
|
||
symbol=message['symbol'],
|
||
timestamp=datetime.now(timezone.utc),
|
||
data_type=DataType.TICKER,
|
||
data={
|
||
'price': message['price'],
|
||
'volume': message.get('volume', 100),
|
||
'timestamp': datetime.now(timezone.utc).isoformat()
|
||
}
|
||
)
|
||
|
||
return data_point
|
||
|
||
async def _handle_messages(self) -> None:
|
||
"""Simulate receiving and processing messages."""
|
||
if not self.connected or not self.subscribed:
|
||
await asyncio.sleep(0.1)
|
||
return
|
||
|
||
# Simulate receiving data for each symbol
|
||
for symbol in self.symbols:
|
||
try:
|
||
# Create simulated message
|
||
simulated_message = {
|
||
'symbol': symbol,
|
||
'price': 50000 + (self.message_count % 1000), # Fake price that changes
|
||
'volume': 1.5
|
||
}
|
||
|
||
# Process the message
|
||
data_point = await self._process_message(simulated_message)
|
||
if data_point:
|
||
self._stats['messages_processed'] += 1
|
||
await self._notify_callbacks(data_point)
|
||
|
||
except Exception as e:
|
||
# This will trigger reconnection logic
|
||
raise e
|
||
|
||
# Simulate processing delay
|
||
await asyncio.sleep(1.0)
|
||
|
||
|
||
async def data_callback(data_point: MarketDataPoint):
|
||
"""Callback function to handle received data."""
|
||
print(f"📊 Data received: {data_point.exchange} - {data_point.symbol} - "
|
||
f"Price: {data_point.data.get('price')} at {data_point.timestamp.strftime('%H:%M:%S')}")
|
||
|
||
|
||
async def monitor_collectors(manager: CollectorManager, duration: int = 30):
|
||
"""Monitor collector status and print updates."""
|
||
print(f"\n🔍 Starting monitoring for {duration} seconds...")
|
||
|
||
for i in range(duration):
|
||
await asyncio.sleep(1)
|
||
|
||
status = manager.get_status()
|
||
running = len(manager.get_running_collectors())
|
||
failed = len(manager.get_failed_collectors())
|
||
|
||
if i % 5 == 0: # Print status every 5 seconds
|
||
print(f"⏰ Status at {i+1}s: {running} running, {failed} failed, "
|
||
f"{status['statistics']['restarts_performed']} restarts")
|
||
|
||
print("🏁 Monitoring complete")
|
||
|
||
|
||
async def demo_basic_usage():
|
||
"""Demonstrate basic collector usage."""
|
||
print("=" * 60)
|
||
print("🚀 Demo 1: Basic Data Collector Usage")
|
||
print("=" * 60)
|
||
|
||
# Create a stable collector
|
||
collector = DemoDataCollector("demo_exchange", ["BTC-USDT", "ETH-USDT"])
|
||
|
||
# Add data callback
|
||
collector.add_data_callback(DataType.TICKER, data_callback)
|
||
|
||
# Start the collector
|
||
print("Starting collector...")
|
||
success = await collector.start()
|
||
if success:
|
||
print("✅ Collector started successfully")
|
||
|
||
# Let it run for a few seconds
|
||
await asyncio.sleep(5)
|
||
|
||
# Show status
|
||
status = collector.get_status()
|
||
print(f"📈 Messages processed: {status['statistics']['messages_processed']}")
|
||
print(f"⏱️ Uptime: {status['statistics']['uptime_seconds']:.1f}s")
|
||
|
||
# Stop the collector
|
||
await collector.stop()
|
||
print("✅ Collector stopped")
|
||
else:
|
||
print("❌ Failed to start collector")
|
||
|
||
|
||
async def demo_manager_usage():
|
||
"""Demonstrate collector manager usage."""
|
||
print("\n" + "=" * 60)
|
||
print("🎛️ Demo 2: Collector Manager Usage")
|
||
print("=" * 60)
|
||
|
||
# Create manager
|
||
manager = CollectorManager("demo_manager", global_health_check_interval=3.0)
|
||
|
||
# Create multiple collectors
|
||
stable_collector = DemoDataCollector("stable_exchange", ["BTC-USDT"])
|
||
failing_collector = DemoDataCollector("failing_exchange", ["ETH-USDT"],
|
||
fail_every_n_messages=5) # Fails every 5 messages
|
||
|
||
# Add data callbacks
|
||
stable_collector.add_data_callback(DataType.TICKER, data_callback)
|
||
failing_collector.add_data_callback(DataType.TICKER, data_callback)
|
||
|
||
# Add collectors to manager
|
||
manager.add_collector(stable_collector)
|
||
manager.add_collector(failing_collector)
|
||
|
||
print(f"📝 Added {len(manager.list_collectors())} collectors to manager")
|
||
|
||
# Start manager
|
||
success = await manager.start()
|
||
if success:
|
||
print("✅ Manager started successfully")
|
||
|
||
# Monitor for a while
|
||
await monitor_collectors(manager, duration=15)
|
||
|
||
# Show final status
|
||
status = manager.get_status()
|
||
print(f"\n📊 Final Statistics:")
|
||
print(f" - Total restarts: {status['statistics']['restarts_performed']}")
|
||
print(f" - Running collectors: {len(manager.get_running_collectors())}")
|
||
print(f" - Failed collectors: {len(manager.get_failed_collectors())}")
|
||
|
||
# Stop manager
|
||
await manager.stop()
|
||
print("✅ Manager stopped")
|
||
else:
|
||
print("❌ Failed to start manager")
|
||
|
||
|
||
async def demo_dynamic_management():
|
||
"""Demonstrate dynamic collector management."""
|
||
print("\n" + "=" * 60)
|
||
print("🔄 Demo 3: Dynamic Collector Management")
|
||
print("=" * 60)
|
||
|
||
# Create manager
|
||
manager = CollectorManager("dynamic_manager", global_health_check_interval=2.0)
|
||
|
||
# Start with one collector
|
||
collector1 = DemoDataCollector("exchange_1", ["BTC-USDT"])
|
||
collector1.add_data_callback(DataType.TICKER, data_callback)
|
||
|
||
manager.add_collector(collector1)
|
||
await manager.start()
|
||
|
||
print("✅ Started with 1 collector")
|
||
await asyncio.sleep(3)
|
||
|
||
# Add second collector
|
||
collector2 = DemoDataCollector("exchange_2", ["ETH-USDT"])
|
||
collector2.add_data_callback(DataType.TICKER, data_callback)
|
||
manager.add_collector(collector2)
|
||
|
||
print("➕ Added second collector")
|
||
await asyncio.sleep(3)
|
||
|
||
# Disable first collector
|
||
collector_names = manager.list_collectors()
|
||
manager.disable_collector(collector_names[0])
|
||
|
||
print("⏸️ Disabled first collector")
|
||
await asyncio.sleep(3)
|
||
|
||
# Re-enable first collector
|
||
manager.enable_collector(collector_names[0])
|
||
|
||
print("▶️ Re-enabled first collector")
|
||
await asyncio.sleep(3)
|
||
|
||
# Show final status
|
||
status = manager.get_status()
|
||
print(f"📊 Final state: {len(manager.get_running_collectors())} running collectors")
|
||
|
||
await manager.stop()
|
||
print("✅ Dynamic demo complete")
|
||
|
||
|
||
async def main():
|
||
"""Run all demonstrations."""
|
||
print("🎯 Data Collector System Demonstration")
|
||
print("This demo shows health monitoring and auto-restart capabilities\n")
|
||
|
||
try:
|
||
# Run demonstrations
|
||
await demo_basic_usage()
|
||
await demo_manager_usage()
|
||
await demo_dynamic_management()
|
||
|
||
print("\n" + "=" * 60)
|
||
print("🎉 All demonstrations completed successfully!")
|
||
print("=" * 60)
|
||
|
||
except Exception as e:
|
||
print(f"❌ Demo failed with error: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main()) |