Update internal/storage/persistence.go

internal/storage/persistence.go

@@ -16,6 +16,8 @@ package storage
 import (
     "bytes"
     "compress/gzip"
+    "crypto/sha256"
+    "encoding/hex"
     "encoding/json"
     "fmt"
     "os"
@@ -232,6 +234,7 @@ type PersistenceConfig struct {
     SyncWrites          bool          `json:"sync_writes"`
     WalPath             string        `json:"wal_path"`
     UseWalForCheckpoint bool          `json:"use_wal_for_checkpoint"`
+    AtomicWrites        bool          `json:"atomic_writes"`
 }
 
 // DefaultPersistenceConfig возвращает конфигурацию по умолчанию
@@ -244,6 +247,7 @@ func DefaultPersistenceConfig() *PersistenceConfig {
         SyncWrites:          true,
         WalPath:             "futriis.wal",
         UseWalForCheckpoint: true,
+        AtomicWrites:        true,
     }
 }
 
@@ -258,6 +262,19 @@ type DatabaseSnapshot struct {
     WalSegment  string                 `json:"wal_segment,omitempty"`
 }
 
+// CheckpointRecoveryInfo содержит информацию о восстановлении из чекпоинта
+type CheckpointRecoveryInfo struct {
+    CheckpointTime   int64   `json:"checkpoint_time"`
+    CheckpointVersion uint64  `json:"checkpoint_version"`
+    WalLSN           uint64  `json:"wal_lsn"`
+    AppliedWalCount  int     `json:"applied_wal_count"`
+    RestoredDocs     int64   `json:"restored_docs"`
+    RestoredColls    int     `json:"restored_collections"`
+    DurationMs       int64   `json:"duration_ms"`
+    Success          bool    `json:"success"`
+    Error            string  `json:"error,omitempty"`
+}
+
 // PersistenceManager управляет персистентным хранением
 type PersistenceManager struct {
     config          *PersistenceConfig
@@ -270,6 +287,8 @@ type PersistenceManager struct {
     checkpointID    atomic.Uint64
     walManager      interface{}
     walReader       WALReader
+    recoveryInfo    *CheckpointRecoveryInfo
+    isRestoring     atomic.Bool
 }
 
 // NewPersistenceManager создаёт новый менеджер персистентности
@@ -292,7 +311,6 @@ func NewPersistenceManager(config *PersistenceConfig, storage *Storage, logger L
         }
     }
 
-    // Инициализируем WAL reader для восстановления
     if config.UseWalForCheckpoint && config.WalPath != "" {
         pm.walReader = NewWALReaderImpl(filepath.Dir(config.WalPath), logger)
     }
@@ -345,6 +363,12 @@ func (pm *PersistenceManager) checkpointLoop() {
         case <-pm.stopChan:
             return
         case <-ticker.C:
+            if pm.isRestoring.Load() {
+                if pm.logger != nil {
+                    pm.logger.Info("Skipping checkpoint during restore")
+                }
+                continue
+            }
             if err := pm.SaveAll(); err != nil {
                 if pm.logger != nil {
                     pm.logger.Error(fmt.Sprintf("Failed to save checkpoint: %v", err))
@@ -356,8 +380,12 @@ func (pm *PersistenceManager) checkpointLoop() {
     }
 }
 
-// SaveDatabase сохраняет базу данных на диск с синхронизацией WAL
+// SaveDatabase сохраняет базу данных на диск с атомарной записью
 func (pm *PersistenceManager) SaveDatabase(dbName string) error {
+    if pm.isRestoring.Load() {
+        return fmt.Errorf("cannot save checkpoint during restore")
+    }
+
     pm.mu.Lock()
     defer pm.mu.Unlock()
 
@@ -403,6 +431,7 @@ func (pm *PersistenceManager) SaveDatabase(dbName string) error {
                 "fields":     doc.GetFields(),
                 "created_at": doc.CreatedAt,
                 "updated_at": doc.UpdatedAt,
+                "deleted_at": doc.DeletedAt,
                 "version":    doc.Version,
             })
         }
@@ -414,6 +443,16 @@ func (pm *PersistenceManager) SaveDatabase(dbName string) error {
         return err
     }
 
+    // Вычисляем контрольную сумму
+    hash := sha256.Sum256(data)
+    snapshot.Checksum = hex.EncodeToString(hash[:])
+
+    // Пере-сериализуем с контрольной суммой
+    data, err = json.Marshal(snapshot)
+    if err != nil {
+        return err
+    }
+
     if pm.config.CompressEnabled {
         data, err = pm.compress(data)
         if err != nil {
@@ -421,29 +460,61 @@ func (pm *PersistenceManager) SaveDatabase(dbName string) error {
         }
     }
 
+    // Атомарная запись с использованием временного файла
     filename := pm.getSnapshotFilename(dbName, snapshot.Version)
-    if err := os.WriteFile(filename, data, 0644); err != nil {
+    if err := pm.atomicWriteFile(filename, data); err != nil {
         return err
     }
 
-    if pm.config.SyncWrites {
+    pm.lastCheckpoint = snapshot.CreatedAt
-        if f, err := os.OpenFile(filename, os.O_RDWR, 0644); err == nil {
-            RealFsync(f)
-            f.Close()
-        }
-        FsyncDir(pm.config.DataDir)
-    }
 
     if pm.logger != nil {
-        pm.logger.Info(fmt.Sprintf("Saved database %s checkpoint %d (%d bytes, WAL LSN: %d)",
+        pm.logger.Info(fmt.Sprintf("Saved database %s checkpoint %d (%d bytes, checksum: %s, WAL LSN: %d)",
-            dbName, snapshot.Version, len(data), walLSN))
+            dbName, snapshot.Version, len(data), snapshot.Checksum[:8], walLSN))
     }
 
     return nil
 }
 
+// atomicWriteFile выполняет атомарную запись файла
+func (pm *PersistenceManager) atomicWriteFile(filename string, data []byte) error {
+    if pm.config.AtomicWrites {
+        // Создаём временный файл в той же директории
+        tempFile := filename + ".tmp"
+        if err := os.WriteFile(tempFile, data, 0644); err != nil {
+            return fmt.Errorf("failed to write temp file: %v", err)
+        }
+
+        // Синхронизируем временный файл
+        if pm.config.SyncWrites {
+            if f, err := os.OpenFile(tempFile, os.O_RDWR, 0644); err == nil {
+                RealFsync(f)
+                f.Close()
+            }
+        }
+
+        // Атомарное переименование
+        if err := os.Rename(tempFile, filename); err != nil {
+            os.Remove(tempFile)
+            return fmt.Errorf("failed to rename temp file: %v", err)
+        }
+
+        // Синхронизируем директорию
+        FsyncDir(pm.config.DataDir)
+
+        return nil
+    }
+
+    // Обычная запись (не атомарная)
+    return os.WriteFile(filename, data, 0644)
+}
+
 // SaveAll сохраняет все базы данных
 func (pm *PersistenceManager) SaveAll() error {
+    if pm.isRestoring.Load() {
+        return fmt.Errorf("cannot save checkpoints during restore")
+    }
+
     databases := pm.storage.ListDatabases()
 
     var lastErr error
@@ -459,51 +530,80 @@ func (pm *PersistenceManager) SaveAll() error {
     return lastErr
 }
 
-// LoadDatabase загружает базу данных с диска
+// LoadDatabase загружает базу данных с диска с применением WAL
 func (pm *PersistenceManager) LoadDatabase(dbName string) error {
+    pm.isRestoring.Store(true)
+    defer pm.isRestoring.Store(false)
+
     pm.mu.Lock()
     defer pm.mu.Unlock()
 
+    startTime := time.Now()
+
+    recoveryInfo := &CheckpointRecoveryInfo{
+        CheckpointTime: time.Now().UnixMilli(),
+        Success:        false,
+    }
+
     snapshot, err := pm.findLatestSnapshot(dbName)
     if err != nil {
+        recoveryInfo.Error = err.Error()
+        pm.recoveryInfo = recoveryInfo
         return err
     }
 
     if snapshot == nil {
+        recoveryInfo.Error = "no snapshot found"
+        pm.recoveryInfo = recoveryInfo
         return fmt.Errorf("no snapshot found for database: %s", dbName)
     }
 
-    if pm.config.UseWalForCheckpoint && snapshot.WalLSN > 0 && pm.walReader != nil {
+    recoveryInfo.CheckpointVersion = snapshot.Version
-        currentLSN, err := pm.walReader.GetCurrentIndex()
+    recoveryInfo.WalLSN = snapshot.WalLSN
-        if err == nil && currentLSN > snapshot.WalLSN {
+
-            if pm.logger != nil {
+    // Проверяем целостность чекпоинта
-                pm.logger.Info(fmt.Sprintf("Database %s: WAL has %d entries after checkpoint, applying...",
+    if err := pm.verifySnapshot(snapshot); err != nil {
-                    dbName, currentLSN-snapshot.WalLSN))
+        recoveryInfo.Error = err.Error()
-            }
+        pm.recoveryInfo = recoveryInfo
-        }
+        return fmt.Errorf("snapshot integrity check failed: %v", err)
     }
 
+    // Создаём базу данных если не существует
     if !pm.storage.ExistsDatabase(dbName) {
         if err := pm.storage.CreateDatabase(dbName); err != nil {
+            recoveryInfo.Error = err.Error()
+            pm.recoveryInfo = recoveryInfo
             return err
         }
     }
 
     db, err := pm.storage.GetDatabase(dbName)
     if err != nil {
+        recoveryInfo.Error = err.Error()
+        pm.recoveryInfo = recoveryInfo
         return err
     }
 
+    // Восстанавливаем коллекции
+    restoredColls := 0
+    restoredDocs := int64(0)
+
     for collName, collDataRaw := range snapshot.Collections {
         collData, ok := collDataRaw.([]interface{})
         if !ok {
             continue
         }
 
-        if _, err := db.GetCollection(collName); err != nil {
+        // Удаляем существующую коллекцию если есть
-            if err := db.CreateCollection(collName); err != nil {
+        if _, err := db.GetCollection(collName); err == nil {
-                continue
+            db.DropCollection(collName)
         }
+
+        if err := db.CreateCollection(collName); err != nil {
+            if pm.logger != nil {
+                pm.logger.Error(fmt.Sprintf("Failed to create collection %s: %v", collName, err))
+            }
+            continue
         }
 
         coll, err := db.GetCollection(collName)
@@ -511,6 +611,8 @@ func (pm *PersistenceManager) LoadDatabase(dbName string) error {
             continue
         }
 
+        restoredColls++
+
         for _, docRaw := range collData {
             docMap, ok := docRaw.(map[string]interface{})
             if !ok {
@@ -532,27 +634,172 @@ func (pm *PersistenceManager) LoadDatabase(dbName string) error {
 
             if createdAt, ok := docMap["created_at"].(float64); ok {
                 doc.CreatedAt = int64(createdAt)
+            } else if createdAt, ok := docMap["created_at"].(int64); ok {
+                doc.CreatedAt = createdAt
             }
 
             if updatedAt, ok := docMap["updated_at"].(float64); ok {
                 doc.UpdatedAt = int64(updatedAt)
+            } else if updatedAt, ok := docMap["updated_at"].(int64); ok {
+                doc.UpdatedAt = updatedAt
+            }
+
+            if deletedAt, ok := docMap["deleted_at"].(float64); ok {
+                doc.DeletedAt = int64(deletedAt)
+            } else if deletedAt, ok := docMap["deleted_at"].(int64); ok {
+                doc.DeletedAt = deletedAt
             }
 
             if version, ok := docMap["version"].(float64); ok {
                 doc.Version = uint64(version)
+            } else if version, ok := docMap["version"].(uint64); ok {
+                doc.Version = version
             }
 
             if err := coll.Insert(doc); err != nil {
                 if pm.logger != nil {
-                    pm.logger.Error(fmt.Sprintf("Failed to restore document %s: %v", docID, err))
+                    pm.logger.Warn(fmt.Sprintf("Failed to restore document %s: %v", docID, err))
+                }
+                continue
+            }
+            restoredDocs++
+        }
+    }
+
+    recoveryInfo.RestoredDocs = restoredDocs
+    recoveryInfo.RestoredColls = restoredColls
+
+    // Применяем WAL записи после чекпоинта
+    appliedWalCount := 0
+    if pm.config.UseWalForCheckpoint && snapshot.WalLSN > 0 && pm.walReader != nil {
+        currentLSN, err := pm.walReader.GetCurrentIndex()
+        if err == nil && currentLSN > snapshot.WalLSN {
+            if pm.logger != nil {
+                pm.logger.Info(fmt.Sprintf("Database %s: Applying %d WAL entries after checkpoint (LSN %d -> %d)",
+                    dbName, currentLSN-snapshot.WalLSN, snapshot.WalLSN, currentLSN))
+            }
+
+            entries, err := pm.walReader.ReadSince(snapshot.WalLSN)
+            if err != nil {
+                if pm.logger != nil {
+                    pm.logger.Error(fmt.Sprintf("Failed to read WAL entries: %v", err))
+                }
+            } else {
+                for _, entry := range entries {
+                    if err := pm.applyWALEntry(db, entry); err != nil {
+                        if pm.logger != nil {
+                            pm.logger.Error(fmt.Sprintf("Failed to apply WAL entry %d: %v", entry.Index, err))
+                        }
+                        continue
+                    }
+                    appliedWalCount++
                 }
             }
         }
     }
 
+    recoveryInfo.AppliedWalCount = appliedWalCount
+    recoveryInfo.DurationMs = time.Since(startTime).Milliseconds()
+    recoveryInfo.Success = true
+
+    pm.recoveryInfo = recoveryInfo
+
     if pm.logger != nil {
-        pm.logger.Info(fmt.Sprintf("Loaded database %s from snapshot (version %d, WAL LSN: %d)",
+        pm.logger.Info(fmt.Sprintf("Loaded database %s from snapshot (version %d, WAL LSN: %d, applied %d WAL entries, %d docs, %d colls, duration: %dms)",
-            dbName, snapshot.Version, snapshot.WalLSN))
+            dbName, snapshot.Version, snapshot.WalLSN, appliedWalCount, restoredDocs, restoredColls, recoveryInfo.DurationMs))
+    }
+
+    return nil
+}
+
+// applyWALEntry применяет одну запись WAL к базе данных
+func (pm *PersistenceManager) applyWALEntry(db *Database, entry WALEntry) error {
+    var txRecord struct {
+        ID         uint64 `json:"id"`
+        State      int32  `json:"state"`
+        Timestamp  int64  `json:"timestamp"`
+        Operations []struct {
+            Type       string                 `json:"type"`
+            Database   string                 `json:"database"`
+            Collection string                 `json:"collection"`
+            DocumentID string                 `json:"document_id"`
+            Data       map[string]interface{} `json:"data"`
+            Version    uint64                 `json:"version"`
+        } `json:"operations"`
+    }
+
+    if err := json.Unmarshal(entry.Data, &txRecord); err != nil {
+        return err
+    }
+
+    // Применяем только закоммиченные транзакции
+    if txRecord.State != 1 { // TransactionCommitted
+        return nil
+    }
+
+    for _, op := range txRecord.Operations {
+        // Пропускаем операции для других баз данных
+        if op.Database != db.name {
+            continue
+        }
+
+        coll, err := db.GetCollection(op.Collection)
+        if err != nil {
+            continue
+        }
+
+        switch op.Type {
+        case "insert":
+            doc := NewDocumentWithID(op.DocumentID)
+            for k, v := range op.Data {
+                doc.SetField(k, v)
+            }
+            doc.Version = op.Version
+            if err := coll.Insert(doc); err != nil {
+                return err
+            }
+
+        case "update":
+            if err := coll.Update(op.DocumentID, op.Data); err != nil {
+                return err
+            }
+
+        case "delete":
+            if err := coll.Delete(op.DocumentID); err != nil {
+                return err
+            }
+        }
+    }
+
+    return nil
+}
+
+// verifySnapshot проверяет целостность снапшота
+func (pm *PersistenceManager) verifySnapshot(snapshot *DatabaseSnapshot) error {
+    if snapshot.Checksum == "" {
+        return fmt.Errorf("snapshot has no checksum")
+    }
+
+    // Пересоздаём данные без контрольной суммы для проверки
+    tempSnapshot := &DatabaseSnapshot{
+        Name:        snapshot.Name,
+        Collections: snapshot.Collections,
+        CreatedAt:   snapshot.CreatedAt,
+        Version:     snapshot.Version,
+        WalLSN:      snapshot.WalLSN,
+        WalSegment:  snapshot.WalSegment,
+    }
+
+    data, err := json.Marshal(tempSnapshot)
+    if err != nil {
+        return fmt.Errorf("failed to marshal for checksum: %v", err)
+    }
+
+    hash := sha256.Sum256(data)
+    checksum := hex.EncodeToString(hash[:])
+
+    if checksum != snapshot.Checksum {
+        return fmt.Errorf("checksum mismatch: expected %s, got %s", snapshot.Checksum[:8], checksum[:8])
     }
 
     return nil
@@ -703,11 +950,31 @@ func (pm *PersistenceManager) GetLastCheckpointInfo() map[string]interface{} {
     pm.mu.RLock()
     defer pm.mu.RUnlock()
 
-    return map[string]interface{}{
+    info := map[string]interface{}{
         "last_checkpoint_time":      pm.lastCheckpoint,
         "last_checkpoint_time_str":  time.UnixMilli(pm.lastCheckpoint).Format("2006-01-02 15:04:05.000"),
         "checkpoint_id":             pm.checkpointID.Load(),
         "wal_path":                  pm.config.WalPath,
         "use_wal":                   pm.config.UseWalForCheckpoint,
+        "atomic_writes":             pm.config.AtomicWrites,
+        "is_restoring":              pm.isRestoring.Load(),
     }
+
+    if pm.recoveryInfo != nil {
+        info["last_recovery"] = pm.recoveryInfo
+    }
+
+    return info
+}
+
+// GetRecoveryInfo возвращает информацию о последнем восстановлении
+func (pm *PersistenceManager) GetRecoveryInfo() *CheckpointRecoveryInfo {
+    pm.mu.RLock()
+    defer pm.mu.RUnlock()
+    return pm.recoveryInfo
+}
+
+// IsRestoring возвращает статус восстановления
+func (pm *PersistenceManager) IsRestoring() bool {
+    return pm.isRestoring.Load()
 }