...

pkg/system/stats/manager.go

@@ -0,0 +1,595 @@
+package stats
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+// StatsManager is a factory for managing system statistics with caching
+type StatsManager struct {
+	// Redis client for caching
+	redisClient *redis.Client
+
+	// Expiration times for different types of stats in seconds
+	Expiration map[string]time.Duration
+
+	// Debug mode - if true, requests are direct without caching
+	Debug bool
+
+	// Queue for requesting stats updates
+	updateQueue chan string
+
+	// Mutex for thread-safe operations
+	mu sync.Mutex
+
+	// Context for controlling the background goroutine
+	ctx    context.Context
+	cancel context.CancelFunc
+	
+	// Default timeout for waiting for stats
+	defaultTimeout time.Duration
+	
+	// Logger for StatsManager operations
+	logger *log.Logger
+}
+
+// NewStatsManager creates a new StatsManager with Redis connection
+func NewStatsManager(config *Config) (*StatsManager, error) {
+	// Use default config if nil is provided
+	if config == nil {
+		config = DefaultConfig()
+	}
+
+	// Create Redis client
+	client := redis.NewClient(&redis.Options{
+		Addr:     config.RedisAddr,
+		Password: config.RedisPassword,
+		DB:       config.RedisDB,
+	})
+
+	// Test connection
+	ctx := context.Background()
+	_, err := client.Ping(ctx).Result()
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to Redis: %w", err)
+	}
+
+	// Create context with cancel for the background goroutine
+	ctx, cancel := context.WithCancel(ctx)
+
+	// Create logger
+	logger := log.New(os.Stdout, "[StatsManager] ", log.LstdFlags)
+
+	// Create the manager
+	manager := &StatsManager{
+		redisClient:    client,
+		Expiration:     config.ExpirationTimes,
+		Debug:          config.Debug,
+		updateQueue:    make(chan string, config.QueueSize),
+		ctx:            ctx,
+		cancel:         cancel,
+		defaultTimeout: config.DefaultTimeout,
+		logger:         logger,
+	}
+
+	// Start the background goroutine for updates
+	go manager.updateWorker()
+
+	// Initialize cache with first fetch
+	manager.initializeCache()
+
+	return manager, nil
+}
+
+// NewStatsManagerWithDefaults creates a new StatsManager with default settings
+func NewStatsManagerWithDefaults() (*StatsManager, error) {
+	return NewStatsManager(DefaultConfig())
+}
+
+// Close closes the StatsManager and its connections
+func (sm *StatsManager) Close() error {
+	// Stop the background goroutine
+	sm.cancel()
+
+	// Close Redis connection
+	return sm.redisClient.Close()
+}
+
+// updateWorker is a background goroutine that processes update requests
+func (sm *StatsManager) updateWorker() {
+	sm.logger.Println("Starting stats update worker goroutine")
+	for {
+		select {
+		case <-sm.ctx.Done():
+			// Context cancelled, exit the goroutine
+			sm.logger.Println("Stopping stats update worker goroutine")
+			return
+		case statsType := <-sm.updateQueue:
+			// Process the update request
+			sm.logger.Printf("Processing update request for %s stats", statsType)
+			sm.fetchAndCacheStats(statsType)
+		}
+	}
+}
+
+// fetchAndCacheStats fetches stats and caches them in Redis
+func (sm *StatsManager) fetchAndCacheStats(statsType string) {
+	var data interface{}
+	var err error
+	
+	sm.logger.Printf("Fetching %s stats", statsType)
+	startTime := time.Now()
+	
+	// Fetch the requested stats
+	switch statsType {
+	case "system":
+		data, err = GetSystemInfo()
+	case "disk":
+		data, err = GetDiskStats()
+	case "process":
+		data, err = GetProcessStats(0) // Get all processes
+	case "root_disk":
+		data, err = GetRootDiskInfo()
+	case "network":
+		data = GetNetworkSpeedResult()
+	case "hardware":
+		data = GetHardwareStatsJSON()
+	default:
+		sm.logger.Printf("Unknown stats type: %s", statsType)
+		return // Unknown stats type
+	}
+	
+	if err != nil {
+		// Log error but continue
+		sm.logger.Printf("Error fetching %s stats: %v", statsType, err)
+		return
+	}
+	
+	// Marshal to JSON
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		sm.logger.Printf("Error marshaling %s stats: %v", statsType, err)
+		return
+	}
+	
+	// Cache in Redis
+	key := fmt.Sprintf("stats:%s", statsType)
+	err = sm.redisClient.Set(sm.ctx, key, jsonData, sm.Expiration[statsType]).Err()
+	if err != nil {
+		sm.logger.Printf("Error caching %s stats: %v", statsType, err)
+		return
+	}
+	
+	// Set last update time
+	lastUpdateKey := fmt.Sprintf("stats:%s:last_update", statsType)
+	sm.redisClient.Set(sm.ctx, lastUpdateKey, time.Now().Unix(), 0)
+	
+	sm.logger.Printf("Successfully cached %s stats in %v", statsType, time.Since(startTime))
+}
+
+// initializeCache initializes the cache with initial values
+func (sm *StatsManager) initializeCache() {
+	sm.logger.Println("Initializing stats cache")
+	
+	// Queue initial fetches for all stats types
+	statsTypes := []string{"system", "disk", "process", "root_disk", "network", "hardware"}
+	for _, statsType := range statsTypes {
+		sm.logger.Printf("Queueing initial fetch for %s stats", statsType)
+		sm.updateQueue <- statsType
+	}
+}
+
+// getFromCache gets stats from cache or triggers an update if expired
+func (sm *StatsManager) getFromCache(statsType string, result interface{}) error {
+	// In debug mode, fetch directly without caching
+	if sm.Debug {
+		sm.logger.Printf("Debug mode enabled, fetching %s stats directly", statsType)
+		return sm.fetchDirect(statsType, result)
+	}
+	
+	key := fmt.Sprintf("stats:%s", statsType)
+	sm.logger.Printf("Getting %s stats from cache", statsType)
+	
+	// Get from Redis
+	jsonData, err := sm.redisClient.Get(sm.ctx, key).Bytes()
+	if err == redis.Nil {
+		// Not in cache, fetch directly and wait
+		sm.logger.Printf("%s stats not found in cache, fetching directly", statsType)
+		return sm.fetchDirectAndCache(statsType, result)
+	} else if err != nil {
+		sm.logger.Printf("Redis error when getting %s stats: %v", statsType, err)
+		return fmt.Errorf("redis error: %w", err)
+	}
+	
+	// Unmarshal the data
+	if err := json.Unmarshal(jsonData, result); err != nil {
+		sm.logger.Printf("Error unmarshaling %s stats: %v", statsType, err)
+		return fmt.Errorf("error unmarshaling data: %w", err)
+	}
+	
+	// Check if data is expired
+	lastUpdateKey := fmt.Sprintf("stats:%s:last_update", statsType)
+	lastUpdateStr, err := sm.redisClient.Get(sm.ctx, lastUpdateKey).Result()
+	if err == nil {
+		var lastUpdate int64
+		fmt.Sscanf(lastUpdateStr, "%d", &lastUpdate)
+		
+		// If expired, queue an update for next time
+		expiration := sm.Expiration[statsType]
+		updateTime := time.Unix(lastUpdate, 0)
+		age := time.Since(updateTime)
+		
+		sm.logger.Printf("%s stats age: %v (expiration: %v)", statsType, age, expiration)
+		
+		if age > expiration {
+			sm.logger.Printf("%s stats expired, queueing update for next request", statsType)
+			// Queue update for next request
+			select {
+			case sm.updateQueue <- statsType:
+				// Successfully queued
+				sm.logger.Printf("Successfully queued %s stats update", statsType)
+			default:
+				// Queue is full, skip update
+				sm.logger.Printf("Update queue full, skipping %s stats update", statsType)
+			}
+		}
+	}
+	
+	return nil
+}
+
+// fetchDirect fetches stats directly without caching
+func (sm *StatsManager) fetchDirect(statsType string, result interface{}) error {
+	var data interface{}
+	var err error
+
+	// Fetch the requested stats
+	switch statsType {
+	case "system":
+		data, err = GetSystemInfo()
+	case "disk":
+		data, err = GetDiskStats()
+	case "process":
+		data, err = GetProcessStats(0) // Get all processes
+	case "root_disk":
+		data, err = GetRootDiskInfo()
+	case "network":
+		data = GetNetworkSpeedResult()
+	case "hardware":
+		data = GetHardwareStatsJSON()
+	default:
+		return fmt.Errorf("unknown stats type: %s", statsType)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	// Convert data to the expected type
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		return err
+	}
+
+	return json.Unmarshal(jsonData, result)
+}
+
+// fetchDirectAndCache fetches stats directly and caches them
+func (sm *StatsManager) fetchDirectAndCache(statsType string, result interface{}) error {
+	var data interface{}
+	var err error
+
+	// Fetch the requested stats
+	switch statsType {
+	case "system":
+		data, err = GetSystemInfo()
+	case "disk":
+		data, err = GetDiskStats()
+	case "process":
+		data, err = GetProcessStats(0) // Get all processes
+	case "root_disk":
+		data, err = GetRootDiskInfo()
+	case "network":
+		data = GetNetworkSpeedResult()
+	case "hardware":
+		data = GetHardwareStatsJSON()
+	default:
+		return fmt.Errorf("unknown stats type: %s", statsType)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	// Convert data to the expected type
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		return err
+	}
+
+	// Cache in Redis
+	key := fmt.Sprintf("stats:%s", statsType)
+	err = sm.redisClient.Set(sm.ctx, key, jsonData, sm.Expiration[statsType]).Err()
+	if err != nil {
+		return err
+	}
+
+	// Set last update time
+	lastUpdateKey := fmt.Sprintf("stats:%s:last_update", statsType)
+	sm.redisClient.Set(sm.ctx, lastUpdateKey, time.Now().Unix(), 0)
+
+	return json.Unmarshal(jsonData, result)
+}
+
+// waitForCachedData waits for data to be available in cache with timeout
+func (sm *StatsManager) waitForCachedData(statsType string, timeout time.Duration) bool {
+	key := fmt.Sprintf("stats:%s", statsType)
+	startTime := time.Now()
+
+	sm.logger.Printf("Waiting for %s stats to be available in cache (timeout: %v)", statsType, timeout)
+
+	for {
+		// Check if data exists
+		exists, err := sm.redisClient.Exists(sm.ctx, key).Result()
+		if err == nil && exists > 0 {
+			sm.logger.Printf("%s stats found in cache after %v", statsType, time.Since(startTime))
+			return true
+		}
+
+		// Check timeout
+		if time.Since(startTime) > timeout {
+			sm.logger.Printf("Timeout waiting for %s stats in cache", statsType)
+			return false
+		}
+
+		// Wait a bit before checking again
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
+// ClearCache clears all cached stats or a specific stats type
+func (sm *StatsManager) ClearCache(statsType string) error {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	if statsType == "" {
+		// Clear all stats
+		sm.logger.Println("Clearing all cached stats")
+		statsTypes := []string{"system", "disk", "process", "root_disk", "network", "hardware"}
+		for _, t := range statsTypes {
+			key := fmt.Sprintf("stats:%s", t)
+			lastUpdateKey := fmt.Sprintf("stats:%s:last_update", t)
+
+			sm.redisClient.Del(sm.ctx, key)
+			sm.redisClient.Del(sm.ctx, lastUpdateKey)
+		}
+	} else {
+		// Clear specific stats type
+		sm.logger.Printf("Clearing cached %s stats", statsType)
+		key := fmt.Sprintf("stats:%s", statsType)
+		lastUpdateKey := fmt.Sprintf("stats:%s:last_update", statsType)
+
+		sm.redisClient.Del(sm.ctx, key)
+		sm.redisClient.Del(sm.ctx, lastUpdateKey)
+	}
+
+	return nil
+}
+
+// ForceUpdate forces an immediate update of stats
+func (sm *StatsManager) ForceUpdate(statsType string) error {
+	sm.logger.Printf("Forcing immediate update of %s stats", statsType)
+
+	// Clear the cache for this stats type
+	err := sm.ClearCache(statsType)
+	if err != nil {
+		return err
+	}
+
+	// Fetch and cache directly
+	switch statsType {
+	case "system", "disk", "process", "root_disk", "network", "hardware":
+		sm.fetchAndCacheStats(statsType)
+		return nil
+	default:
+		return fmt.Errorf("unknown stats type: %s", statsType)
+	}
+}
+
+// GetSystemInfo gets system information with caching
+func (sm *StatsManager) GetSystemInfo() (*SystemInfo, error) {
+	var result SystemInfo
+
+	// Try to get from cache
+	err := sm.getFromCache("system", &result)
+	if err != nil {
+		return nil, err
+	}
+
+	return &result, nil
+}
+
+// GetDiskStats gets disk statistics with caching
+func (sm *StatsManager) GetDiskStats() (*DiskStats, error) {
+	var result DiskStats
+
+	// Try to get from cache
+	err := sm.getFromCache("disk", &result)
+	if err != nil {
+		return nil, err
+	}
+
+	return &result, nil
+}
+
+// GetRootDiskInfo gets root disk information with caching
+func (sm *StatsManager) GetRootDiskInfo() (*DiskInfo, error) {
+	var result DiskInfo
+
+	// Try to get from cache
+	err := sm.getFromCache("root_disk", &result)
+	if err != nil {
+		return nil, err
+	}
+
+	return &result, nil
+}
+
+// GetProcessStats gets process statistics with caching
+func (sm *StatsManager) GetProcessStats(limit int) (*ProcessStats, error) {
+	var result ProcessStats
+
+	// Try to get from cache
+	err := sm.getFromCache("process", &result)
+	if err != nil {
+		return nil, err
+	}
+
+	// Apply limit if needed
+	if limit > 0 && len(result.Processes) > limit {
+		result.Processes = result.Processes[:limit]
+	}
+
+	return &result, nil
+}
+
+// GetProcessStatsFresh gets fresh process statistics bypassing the cache
+func (sm *StatsManager) GetProcessStatsFresh(limit int) (*ProcessStats, error) {
+	var result ProcessStats
+
+	// Get fresh data and update cache
+	err := sm.fetchDirectAndCache("process", &result)
+	if err != nil {
+		return nil, err
+	}
+
+	// Apply limit if needed
+	if limit > 0 && len(result.Processes) > limit {
+		result.Processes = result.Processes[:limit]
+	}
+
+	// Log that we're bypassing cache
+	sm.logger.Printf("Bypassing cache for process stats with manual refresh")
+	
+	return &result, nil
+}
+
+// GetTopProcesses gets top processes by CPU usage with caching
+func (sm *StatsManager) GetTopProcesses(n int) ([]ProcessInfo, error) {
+	stats, err := sm.GetProcessStats(n)
+	if err != nil {
+		return nil, err
+	}
+
+	return stats.Processes, nil
+}
+
+// GetNetworkSpeedResult gets network speed with caching
+func (sm *StatsManager) GetNetworkSpeedResult() NetworkSpeedResult {
+	var result NetworkSpeedResult
+
+	// Try to get from cache
+	err := sm.getFromCache("network", &result)
+	if err != nil {
+		// Fallback to direct fetch on error
+		uploadSpeed, downloadSpeed := GetNetworkSpeed()
+		return NetworkSpeedResult{
+			UploadSpeed:   uploadSpeed,
+			DownloadSpeed: downloadSpeed,
+		}
+	}
+
+	return result
+}
+
+// GetHardwareStats gets hardware statistics with caching
+func (sm *StatsManager) GetHardwareStats() map[string]interface{} {
+	var result map[string]interface{}
+
+	// Try to get from cache
+	err := sm.getFromCache("hardware", &result)
+	if err != nil {
+		// Fallback to direct fetch on error
+		return GetHardwareStats()
+	}
+
+	return result
+}
+
+// GetHardwareStatsJSON gets hardware statistics in JSON format with caching
+func (sm *StatsManager) GetHardwareStatsJSON() map[string]interface{} {
+	var result map[string]interface{}
+
+	// Try to get from cache
+	err := sm.getFromCache("hardware", &result)
+	if err != nil {
+		// Fallback to direct fetch on error
+		return GetHardwareStatsJSON()
+	}
+
+	return result
+}
+
+// GetFormattedCPUInfo gets formatted CPU info with caching
+func (sm *StatsManager) GetFormattedCPUInfo() string {
+	sysInfo, err := sm.GetSystemInfo()
+	if err != nil {
+		return "Unknown"
+	}
+
+	return fmt.Sprintf("%d cores (%s)", sysInfo.CPU.Cores, sysInfo.CPU.ModelName)
+}
+
+// GetFormattedMemoryInfo gets formatted memory info with caching
+func (sm *StatsManager) GetFormattedMemoryInfo() string {
+	sysInfo, err := sm.GetSystemInfo()
+	if err != nil {
+		return "Unknown"
+	}
+
+	return fmt.Sprintf("%.1fGB (%.1fGB used)", sysInfo.Memory.Total, sysInfo.Memory.Used)
+}
+
+// GetFormattedDiskInfo gets formatted disk info with caching
+func (sm *StatsManager) GetFormattedDiskInfo() string {
+	diskInfo, err := sm.GetRootDiskInfo()
+	if err != nil {
+		return "Unknown"
+	}
+
+	return fmt.Sprintf("%.0fGB (%.0fGB free)", diskInfo.Total, diskInfo.Free)
+}
+
+// GetFormattedNetworkInfo gets formatted network info with caching
+func (sm *StatsManager) GetFormattedNetworkInfo() string {
+	netSpeed := sm.GetNetworkSpeedResult()
+
+	return fmt.Sprintf("Up: %s\nDown: %s", netSpeed.UploadSpeed, netSpeed.DownloadSpeed)
+}
+
+// GetProcessStatsJSON gets process statistics in JSON format with caching
+func (sm *StatsManager) GetProcessStatsJSON(limit int) map[string]interface{} {
+	// Get process stats
+	processStats, err := sm.GetProcessStats(limit)
+	if err != nil {
+		return map[string]interface{}{
+			"processes": []interface{}{},
+			"total":     0,
+			"filtered":  0,
+		}
+	}
+
+	// Convert to JSON-friendly format
+	return map[string]interface{}{
+		"processes": processStats.Processes,
+		"total":     processStats.Total,
+		"filtered":  processStats.Filtered,
+	}
+}