gocache/kvstore.go

package gocache

import (
	"bytes"
	"encoding/json"
	"os"
	"sync"
)

// KVStore represents a simple in-memory key/value store.
type KVStore struct {
	mu          sync.RWMutex
	store       map[string]string
	file        string
	transaction *Transaction
	logFile     string // 添加日志文件路径
}

// Transaction represents an ongoing transaction
type Transaction struct {
	store map[string]string
}

// NewKVStore creates a new instance of KVStore.
func NewKVStore(file string) *KVStore {
	store := &KVStore{
		store:   make(map[string]string),
		file:    file,
		logFile: file + ".log",
	}
	// 启动时自动恢复日志
	if err := store.RecoverFromLog(); err != nil {
		panic(err)
	}
	return store
}

// RecoverFromLog recovers data from log file
func (k *KVStore) RecoverFromLog() error {
	k.mu.Lock()
	defer k.mu.Unlock()

	// 检查日志文件是否存在
	if _, err := os.Stat(k.logFile); os.IsNotExist(err) {
		return nil
	}

	// 读取日志文件
	data, err := os.ReadFile(k.logFile)
	if err != nil {
		return err
	}

	// 逐行处理日志
	lines := bytes.Split(data, []byte{'\n'})
	for _, line := range lines {
		if len(line) ==  {
			continue
		}
		var logEntry map[string]string
		if err := json.Unmarshal(line, &logEntry); err != nil {
			return err
		}
		switch logEntry["op"] {
		case "put":
			k.store[logEntry["key"]] = logEntry["value"]
		case "delete":
			delete(k.store, logEntry["key"])
		}
	}

	// 清空日志文件
	return os.Truncate(k.logFile, 0)
}

// SaveToFile saves the current store to a file.
func (k *KVStore) SaveToFile() error {
	k.mu.RLock()
	defer k.mu.RUnlock()
	data, err := json.Marshal(k.store)
	if err != nil {
		return err
	}
	return os.WriteFile(k.file, data, 0644)
}

// LogOperation logs a key/value operation to the log file
func (k *KVStore) LogOperation(op string, key, value string) error {
	logEntry := map[string]string{
		"op":    op,
		"key":   key,
		"value": value,
	}
	data, err := json.Marshal(logEntry)
	if err != nil {
		return err
	}
	data = append(data, '\n')
	f, err := os.OpenFile(k.logFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.Write(data)
	return err
}

// LoadFromFile loads the store from a file with partial loading support
func (k *KVStore) LoadFromFile(keys ...string) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	// 先加载主数据文件
	if _, err := os.Stat(k.file); err == nil {
		data, err := os.ReadFile(k.file)
		if err != nil {
			return err
		}

		var tempStore map[string]string
		if err := json.Unmarshal(data, &tempStore); err != nil {
			return err
		}

		// Partial loading
		if len(keys) > 0 {
			for _, key := range keys {
				if value, exists := tempStore[key]; exists {
					k.store[key] = value
				}
			}
		}
	}

	// 然后应用日志文件中的操作
	if _, err := os.Stat(k.logFile); err == nil {
		data, err := os.ReadFile(k.logFile)
		if err != nil {
			return err
		}

		lines := bytes.Split(data, []byte{'\n'})
		for _, line := range lines {
			if len(line) == 0 {
				continue
			}
			var logEntry map[string]string
			if err := json.Unmarshal(line, &logEntry); err != nil {
				return err
			}
			switch logEntry["op"] {
			case "put":
				k.store[logEntry["key"]] = logEntry["value"]
			case "delete":
				delete(k.store, logEntry["key"])
			}
		}
	}

	return nil
}

// Put adds a key/value pair to the store.
func (k *KVStore) Put(key, value string) {
	k.mu.Lock()
	defer k.mu.Unlock()
	k.store[key] = value
	go func() {
		k.mu.Lock()
		defer k.mu.Unlock()
		k.LogOperation("put", key, value) // 异步记录操作日志
	}()
}

// Get retrieves the value associated with the given key.
func (k *KVStore) Get(key string) (string, bool) {
	k.mu.RLock()
	defer k.mu.RUnlock()
	if k.transaction != nil {
		value, exists := k.transaction.store[key]
		if exists {
			return value, true
		}
	}
	value, exists := k.store[key]
	return value, exists
}

// Delete removes the key/value pair associated with the given key.
func (k *KVStore) Delete(key string) {
	k.mu.Lock()
	defer k.mu.Unlock()
	delete(k.store, key)
	go func() {
		k.mu.Lock()
		defer k.mu.Unlock()
		k.LogOperation("delete", key, "") // 异步记录操作日志
	}()
}

// BeginTransaction starts a new transaction
func (k *KVStore) BeginTransaction() {
	k.mu.Lock()
	defer k.mu.Unlock()
	k.transaction = &Transaction{
		store: make(map[string]string),
	}
}

// Commit commits the current transaction
func (k *KVStore) Commit() error {
	k.mu.Lock()
	defer k.mu.Unlock()
	if k.transaction == nil {
		return nil
	}

	for key, value := range k.transaction.store {
		k.store[key] = value
	}
	k.transaction = nil
	return nil
}

// Rollback rolls back the current transaction
func (k *KVStore) Rollback() {
	k.mu.Lock()
	defer k.mu.Unlock()
	k.transaction = nil
}

// PutInTransaction puts a key/value pair in the current transaction
func (k *KVStore) PutInTransaction(key, value string) {
	k.mu.Lock()
	defer k.mu.Unlock()
	if k.transaction != nil {
		k.transaction.store[key] = value
	} else {
		k.store[key] = value
	}
}