futriix/internal/storage/runtime_limits.go

/*
 * Copyright 2026 Safronov Grigorii
 *
 * Licensed under the CDDL, Version 1.0 (the "License");
 * you may not use this file except in compliance with the License.
 *
 * You may obtain a copy of the License at
 * https://opensource.org/licenses/CDDL-1.0
 */

// Файл: internal/storage/runtime_limits.go
// Назначение: Ограничения на размер коллекции/документа в рантайме

package storage

import (
    "fmt"
    "runtime"
    "sync"
    "sync/atomic"
    "time"
)

// LoggerInterface определяет интерфейс для логирования
type LoggerInterface interface {
    Debug(msg string)
    Info(msg string)
    Error(msg string)
    Warn(msg string)
}

// EvictionPolicy определяет политику вытеснения
type EvictionPolicy int

const (
    EvictionNone     EvictionPolicy = iota // Без вытеснения
    EvictionLRU                            // Least Recently Used
    EvictionTTL                            // Time To Live
    EvictionOldest                         // Самые старые
)

// RuntimeLimitsManager управляет runtime-ограничениями
// ИСПРАВЛЕНО: Добавлен механизм eviction
type RuntimeLimitsManager struct {
    mu                   sync.RWMutex
    globalMaxDocSize     int64              // Максимальный размер документа (байт)
    globalMaxCollSize    int64              // Максимальный размер коллекции (байт)
    globalMaxDocsPerColl int64              // Максимальное количество документов в коллекции
    globalMaxMemory      int64              // Максимальный размер памяти (байт)
    collectionOverrides  map[string]*CollectionLimits // Переопределения для коллекций
    metrics              *LimitMetrics
    logger               LoggerInterface
    enabled              bool
    evictionPolicy       EvictionPolicy
    memoryThreshold      float64            // Порог использования памяти (0.0 - 1.0)
    // ИСПРАВЛЕНО: Каналы для eviction
    evictionChan         chan string        // Канал запросов на eviction
    stopChan             chan struct{}
    wg                   sync.WaitGroup
}

// CollectionLimits содержит лимиты для конкретной коллекции
type CollectionLimits struct {
    MaxDocSize        int64
    MaxCollectionSize int64
    MaxDocuments      int64
    EvictionPolicy    EvictionPolicy
    LastUpdated       int64
}

// LimitMetrics хранит метрики ограничений
// ИСПРАВЛЕНО: Добавлены метрики eviction
type LimitMetrics struct {
    RejectedBySize      atomic.Uint64
    RejectedByDocCount  atomic.Uint64
    RejectedByCollSize  atomic.Uint64
    RejectedByMemory    atomic.Uint64
    EvictedDocuments    atomic.Uint64
    EvictedBytes        atomic.Uint64
    LastCheckTime       atomic.Int64
    LastEvictionTime    atomic.Int64
}

// RuntimeLimitsConfig содержит конфигурацию ограничений
// ИСПРАВЛЕНО: Добавлены параметры eviction
type RuntimeLimitsConfig struct {
    Enabled              bool           `json:"enabled"`
    GlobalMaxDocSizeMB   int            `json:"global_max_doc_size_mb"`
    GlobalMaxCollSizeMB  int64          `json:"global_max_coll_size_mb"`
    GlobalMaxDocsPerColl int64          `json:"global_max_docs_per_coll"`
    GlobalMaxMemoryMB    int64          `json:"global_max_memory_mb"`
    EvictionPolicy       EvictionPolicy `json:"eviction_policy"`
    MemoryThreshold      float64        `json:"memory_threshold"` // 0.0 - 1.0, при каком % памяти начинать eviction
}

// DefaultRuntimeLimitsConfig возвращает конфигурацию по умолчанию
// ИСПРАВЛЕНО: Добавлены параметры eviction
func DefaultRuntimeLimitsConfig() *RuntimeLimitsConfig {
    return &RuntimeLimitsConfig{
        Enabled:              true,
        GlobalMaxDocSizeMB:   16,       // 16 MB на документ
        GlobalMaxCollSizeMB:  10240,    // 10 GB на коллекцию
        GlobalMaxDocsPerColl: 10000000, // 10 млн документов
        GlobalMaxMemoryMB:    0,        // 0 = автоопределение
        EvictionPolicy:       EvictionLRU,
        MemoryThreshold:      0.85,     // Начинать eviction при 85% использования
    }
}

// NewRuntimeLimitsManager создаёт новый менеджер ограничений
// ИСПРАВЛЕНО: Инициализация eviction механизма
func NewRuntimeLimitsManager(cfg *RuntimeLimitsConfig, logger LoggerInterface) *RuntimeLimitsManager {
    if cfg == nil {
        cfg = DefaultRuntimeLimitsConfig()
    }
    
    // Автоопределение максимального размера памяти
    maxMemory := cfg.GlobalMaxMemoryMB * 1024 * 1024
    if maxMemory <= 0 {
        // Используем 80% доступной памяти
        var memStats runtime.MemStats
        runtime.ReadMemStats(&memStats)
        maxMemory = int64(float64(memStats.Sys) * 0.8)
    }
    
    rlm := &RuntimeLimitsManager{
        globalMaxDocSize:     int64(cfg.GlobalMaxDocSizeMB) * 1024 * 1024,
        globalMaxCollSize:    cfg.GlobalMaxCollSizeMB * 1024 * 1024,
        globalMaxDocsPerColl: cfg.GlobalMaxDocsPerColl,
        globalMaxMemory:      maxMemory,
        collectionOverrides:  make(map[string]*CollectionLimits),
        metrics:              &LimitMetrics{},
        logger:               logger,
        enabled:              cfg.Enabled,
        evictionPolicy:       cfg.EvictionPolicy,
        memoryThreshold:      cfg.MemoryThreshold,
        evictionChan:         make(chan string, 100),
        stopChan:             make(chan struct{}),
    }
    
    // Запускаем фоновый мониторинг памяти
    rlm.wg.Add(1)
    go rlm.memoryMonitorLoop()
    
    if logger != nil {
        logger.Debug(fmt.Sprintf("Runtime limits manager initialized: maxDoc=%dMB, maxColl=%dMB, maxDocs=%d, maxMemory=%dMB, eviction=%d",
            cfg.GlobalMaxDocSizeMB, cfg.GlobalMaxCollSizeMB, cfg.GlobalMaxDocsPerColl, maxMemory/(1024*1024), cfg.EvictionPolicy))
    }
    
    return rlm
}

// memoryMonitorLoop мониторит использование памяти и запускает eviction
// ИСПРАВЛЕНО: Добавлен мониторинг памяти для предотвращения OOM
func (rlm *RuntimeLimitsManager) memoryMonitorLoop() {
    defer rlm.wg.Done()
    
    ticker := time.NewTicker(10 * time.Second)
    defer ticker.Stop()
    
    for {
        select {
        case <-ticker.C:
            if !rlm.enabled {
                continue
            }
            
            var memStats runtime.MemStats
            runtime.ReadMemStats(&memStats)
            
            // Проверяем использование памяти
            currentUsage := int64(memStats.Alloc)
            usageRatio := float64(currentUsage) / float64(rlm.globalMaxMemory)
            
            if usageRatio > rlm.memoryThreshold {
                if rlm.logger != nil {
                    rlm.logger.Warn(fmt.Sprintf("Memory usage %.2f%% exceeds threshold %.2f%%, triggering eviction",
                        usageRatio*100, rlm.memoryThreshold*100))
                }
                
                // Запускаем eviction
                rlm.triggerEviction()
            }
            
        case <-rlm.stopChan:
            return
        }
    }
}

// triggerEviction запускает процесс вытеснения
// ИСПРАВЛЕНО: Добавлен метод для запуска eviction
func (rlm *RuntimeLimitsManager) triggerEviction() {
    select {
    case rlm.evictionChan <- "global":
        // Eviction запрошен
    default:
        // Канал заполнен, eviction уже в процессе
    }
}

// EvictFromCollection выполняет вытеснение документов из коллекции
// ИСПРАВЛЕНО: Добавлен метод для eviction
func (rlm *RuntimeLimitsManager) EvictFromCollection(coll *Collection, targetBytes int64) (int64, error) {
    if !rlm.enabled {
        return 0, nil
    }
    
    evictedBytes := int64(0)
    evictedCount := int64(0)
    
    // Получаем все документы
    docs := coll.GetAllDocumentsIncludingDeleted()
    
    // Сортируем по времени последнего доступа (LRU)
    // В реальной реализации нужно отслеживать время доступа
    // Для простоты используем UpdatedAt
    
    // Ищем удалённые документы для вытеснения в первую очередь
    for _, doc := range docs {
        if evictedBytes >= targetBytes {
            break
        }
        
        if doc.IsDeleted() {
            // Вытесняем удалённые документы
            size := doc.OriginalSize
            if size == 0 {
                size = 1024 // Примерная оценка
            }
            
            if err := coll.PermanentDelete(doc.ID); err == nil {
                evictedBytes += size
                evictedCount++
            }
        }
    }
    
    // Если нужно больше, вытесняем самые старые
    if evictedBytes < targetBytes && rlm.evictionPolicy == EvictionLRU {
        for _, doc := range docs {
            if evictedBytes >= targetBytes {
                break
            }
            
            if !doc.IsDeleted() {
                size := doc.OriginalSize
                if size == 0 {
                    size = 1024
                }
                
                // Soft delete вместо permanent delete для возможности восстановления
                if coll.metadata.Settings.SoftDelete {
                    if err := coll.Delete(doc.ID); err == nil {
                        evictedBytes += size
                        evictedCount++
                    }
                } else {
                    if err := coll.PermanentDelete(doc.ID); err == nil {
                        evictedBytes += size
                        evictedCount++
                    }
                }
            }
        }
    }
    
    rlm.metrics.EvictedDocuments.Add(uint64(evictedCount))
    rlm.metrics.EvictedBytes.Add(uint64(evictedBytes))
    rlm.metrics.LastEvictionTime.Store(time.Now().UnixMilli())
    
    if rlm.logger != nil {
        rlm.logger.Info(fmt.Sprintf("Evicted %d documents (%d bytes) from collection %s.%s",
            evictedCount, evictedBytes, coll.DBName(), coll.Name()))
    }
    
    return evictedBytes, nil
}

// ValidateDocumentSize проверяет размер документа
func (rlm *RuntimeLimitsManager) ValidateDocumentSize(dbName, collName string, docSize int64) error {
    if !rlm.enabled {
        return nil
    }
    
    // Проверяем переопределение для коллекции
    limit := rlm.getCollectionLimit(dbName, collName)
    
    maxSize := rlm.globalMaxDocSize
    if limit != nil && limit.MaxDocSize > 0 {
        maxSize = limit.MaxDocSize
    }
    
    if docSize > maxSize {
        rlm.metrics.RejectedBySize.Add(1)
        return fmt.Errorf("document size %d bytes exceeds limit %d bytes", docSize, maxSize)
    }
    
    return nil
}

// ValidateCollectionSize проверяет размер коллекции
// ИСПРАВЛЕНО: Добавлена проверка памяти и eviction при необходимости
func (rlm *RuntimeLimitsManager) ValidateCollectionSize(coll *Collection, newDocSize int64) error {
    if !rlm.enabled {
        return nil
    }
    
    limit := rlm.getCollectionLimit(coll.DBName(), coll.Name())
    
    maxSize := rlm.globalMaxCollSize
    if limit != nil && limit.MaxCollectionSize > 0 {
        maxSize = limit.MaxCollectionSize
    }
    
    currentSize := coll.Size()
    if currentSize+newDocSize > maxSize {
        // Пытаемся вытеснить данные
        neededBytes := currentSize + newDocSize - maxSize
        if evicted, err := rlm.EvictFromCollection(coll, neededBytes); err == nil && evicted >= neededBytes {
            // Eviction успешен, повторяем проверку
            return nil
        }
        
        rlm.metrics.RejectedByCollSize.Add(1)
        return fmt.Errorf("collection size would exceed limit %d bytes (current: %d, new: %d)",
            maxSize, currentSize, newDocSize)
    }
    
    return nil
}

// ValidateDocumentCount проверяет количество документов в коллекции
// ИСПРАВЛЕНО: Добавлен eviction при превышении лимита
func (rlm *RuntimeLimitsManager) ValidateDocumentCount(coll *Collection) error {
    if !rlm.enabled {
        return nil
    }
    
    limit := rlm.getCollectionLimit(coll.DBName(), coll.Name())
    
    maxDocs := rlm.globalMaxDocsPerColl
    if limit != nil && limit.MaxDocuments > 0 {
        maxDocs = limit.MaxDocuments
    }
    
    currentCount := coll.Count()
    if currentCount >= maxDocs {
        // Пытаемся вытеснить документы
        targetBytes := int64((currentCount - maxDocs + 1) * 1024) // Примерная оценка
        if evicted, err := rlm.EvictFromCollection(coll, targetBytes); err == nil && evicted > 0 {
            // Eviction успешен, повторяем проверку
            if coll.Count() < maxDocs {
                return nil
            }
        }
        
        rlm.metrics.RejectedByDocCount.Add(1)
        return fmt.Errorf("collection has reached maximum document count %d", maxDocs)
    }
    
    return nil
}

// CheckMemoryUsage проверяет использование памяти
// ИСПРАВЛЕНО: Добавлен метод для проверки памяти
func (rlm *RuntimeLimitsManager) CheckMemoryUsage() error {
    if !rlm.enabled {
        return nil
    }
    
    var memStats runtime.MemStats
    runtime.ReadMemStats(&memStats)
    
    currentUsage := int64(memStats.Alloc)
    if currentUsage > rlm.globalMaxMemory {
        rlm.metrics.RejectedByMemory.Add(1)
        return fmt.Errorf("memory usage %d bytes exceeds limit %d bytes", currentUsage, rlm.globalMaxMemory)
    }
    
    return nil
}

// getCollectionLimit возвращает лимиты для коллекции
func (rlm *RuntimeLimitsManager) getCollectionLimit(dbName, collName string) *CollectionLimits {
    rlm.mu.RLock()
    defer rlm.mu.RUnlock()
    
    key := fmt.Sprintf("%s.%s", dbName, collName)
    if limits, ok := rlm.collectionOverrides[key]; ok {
        return limits
    }
    return nil
}

// SetCollectionLimits устанавливает лимиты для коллекции
func (rlm *RuntimeLimitsManager) SetCollectionLimits(dbName, collName string, maxDocSizeMB int, maxCollSizeMB int64, maxDocuments int64) {
    rlm.mu.Lock()
    defer rlm.mu.Unlock()
    
    key := fmt.Sprintf("%s.%s", dbName, collName)
    rlm.collectionOverrides[key] = &CollectionLimits{
        MaxDocSize:        int64(maxDocSizeMB) * 1024 * 1024,
        MaxCollectionSize: maxCollSizeMB * 1024 * 1024,
        MaxDocuments:      maxDocuments,
        EvictionPolicy:    rlm.evictionPolicy,
        LastUpdated:       time.Now().UnixMilli(),
    }
    
    if rlm.logger != nil {
        rlm.logger.Info(fmt.Sprintf("Set limits for %s: maxDoc=%dMB, maxColl=%dMB, maxDocs=%d",
            key, maxDocSizeMB, maxCollSizeMB, maxDocuments))
    }
}

// RemoveCollectionLimits удаляет переопределения для коллекции
func (rlm *RuntimeLimitsManager) RemoveCollectionLimits(dbName, collName string) {
    rlm.mu.Lock()
    defer rlm.mu.Unlock()
    
    key := fmt.Sprintf("%s.%s", dbName, collName)
    delete(rlm.collectionOverrides, key)
    
    if rlm.logger != nil {
        rlm.logger.Info(fmt.Sprintf("Removed limits override for %s", key))
    }
}

// GetMetrics возвращает метрики
// ИСПРАВЛЕНО: Добавлены метрики eviction
func (rlm *RuntimeLimitsManager) GetMetrics() map[string]interface{} {
    return map[string]interface{}{
        "rejected_by_size":        rlm.metrics.RejectedBySize.Load(),
        "rejected_by_doc_count":   rlm.metrics.RejectedByDocCount.Load(),
        "rejected_by_coll_size":   rlm.metrics.RejectedByCollSize.Load(),
        "rejected_by_memory":      rlm.metrics.RejectedByMemory.Load(),
        "evicted_documents":       rlm.metrics.EvictedDocuments.Load(),
        "evicted_bytes":           rlm.metrics.EvictedBytes.Load(),
        "last_eviction_time":      rlm.metrics.LastEvictionTime.Load(),
        "global_max_doc_size_mb":  rlm.globalMaxDocSize / (1024 * 1024),
        "global_max_coll_size_mb": rlm.globalMaxCollSize / (1024 * 1024),
        "global_max_docs_per_coll": rlm.globalMaxDocsPerColl,
        "global_max_memory_mb":    rlm.globalMaxMemory / (1024 * 1024),
        "eviction_policy":         rlm.evictionPolicy,
        "memory_threshold":        rlm.memoryThreshold,
        "enabled":                 rlm.enabled,
    }
}

// Stop останавливает менеджер ограничений
// ИСПРАВЛЕНО: Добавлен метод остановки
func (rlm *RuntimeLimitsManager) Stop() {
    close(rlm.stopChan)
    rlm.wg.Wait()
}