90CNServer/server/internal/platform/db/hash_backend.go

package db

import (
	"context"
	"errors"
	"sort"
	"strings"
	"sync"
	"time"
)

// ErrHashBackendRequired 表示 hash 后端缺失或没有完成初始化。
var ErrHashBackendRequired = errors.New("hash backend is required")

// HashBackend 是按记录主键和字段名读写二进制字段的后端接口。
type HashBackend interface {
	LoadHash(context.Context, string, []string) (map[string][]byte, bool, error)
	SaveHash(context.Context, string, map[string][]byte, time.Duration) error
}

// HashSaveBatch 表示一次 hash 字段批量保存项。
type HashSaveBatch struct {
	Key    string
	Fields map[string][]byte
	TTL    time.Duration
}

// HashBatchBackend 支持批量保存 hash 字段。
type HashBatchBackend interface {
	SaveHashBatch(context.Context, []HashSaveBatch) error
}

// HashLoadRequest 描述一次 hash 批量读取请求。
type HashLoadRequest struct {
	Key    string
	Fields []string
}

// HashLoadResult 描述一次 hash 批量读取结果。
type HashLoadResult struct {
	Fields map[string][]byte
	Found  bool
}

// HashBatchLoaderBackend 支持批量读取 hash 字段。
type HashBatchLoaderBackend interface {
	LoadHashBatch(context.Context, []HashLoadRequest) (map[string]HashLoadResult, error)
}

// MemoryHashBackend 是线程安全的内存 hash 后端。
type MemoryHashBackend struct {
	mu      sync.RWMutex
	records map[string]map[string][]byte
	ttl     map[string]time.Duration
}

// NewMemoryHashBackend 创建内存 hash 后端。
func NewMemoryHashBackend() *MemoryHashBackend {
	return &MemoryHashBackend{
		records: make(map[string]map[string][]byte),
		ttl:     make(map[string]time.Duration),
	}
}

// LoadHash 从内存读取指定 hash 字段。
func (b *MemoryHashBackend) LoadHash(ctx context.Context, key string, fields []string) (map[string][]byte, bool, error) {
	if b == nil {
		return nil, false, ErrHashBackendRequired
	}
	ctx = contextOrBackground(ctx)
	if err := ctx.Err(); err != nil {
		return nil, false, err
	}
	key = strings.TrimSpace(key)
	if key == "" {
		return nil, false, ErrRecordKeyRequired
	}
	fields = normHashFields(fields)

	b.mu.RLock()
	record, ok := b.records[key]
	if !ok {
		b.mu.RUnlock()
		return nil, false, nil
	}
	out := make(map[string][]byte, len(record))
	if len(fields) == 0 {
		for name, data := range record {
			out[name] = append([]byte(nil), data...)
		}
	} else {
		for _, name := range fields {
			if data, exists := record[name]; exists {
				out[name] = append([]byte(nil), data...)
			}
		}
	}
	b.mu.RUnlock()
	return out, true, nil
}

// LoadHashBatch 批量读取 hash 字段。
func (b *MemoryHashBackend) LoadHashBatch(ctx context.Context, requests []HashLoadRequest) (map[string]HashLoadResult, error) {
	if b == nil {
		return nil, ErrHashBackendRequired
	}
	ctx = contextOrBackground(ctx)
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	out := make(map[string]HashLoadResult, len(requests))
	for _, request := range requests {
		fields, found, err := b.LoadHash(ctx, request.Key, request.Fields)
		if err != nil {
			return nil, err
		}
		out[strings.TrimSpace(request.Key)] = HashLoadResult{Fields: fields, Found: found}
	}
	return out, nil
}

// SaveHash 保存指定 hash 字段。
func (b *MemoryHashBackend) SaveHash(ctx context.Context, key string, fields map[string][]byte, ttl time.Duration) error {
	if b == nil {
		return ErrHashBackendRequired
	}
	ctx = contextOrBackground(ctx)
	if err := ctx.Err(); err != nil {
		return err
	}
	key = strings.TrimSpace(key)
	if key == "" {
		return ErrRecordKeyRequired
	}
	fields, err := cloneHashFields(fields)
	if err != nil {
		return err
	}
	if len(fields) == 0 {
		return nil
	}

	b.mu.Lock()
	b.ensureReadyLocked()
	record := b.records[key]
	if record == nil {
		record = make(map[string][]byte, len(fields))
		b.records[key] = record
	}
	for name, data := range fields {
		record[name] = append([]byte(nil), data...)
	}
	if ttl > 0 {
		b.ttl[key] = ttl
	}
	b.mu.Unlock()
	return nil
}

// SaveHashBatch 批量保存 hash 字段。
func (b *MemoryHashBackend) SaveHashBatch(ctx context.Context, batches []HashSaveBatch) error {
	if b == nil {
		return ErrHashBackendRequired
	}
	ctx = contextOrBackground(ctx)
	merged, err := mergeHashSaveBatches(batches)
	if err != nil {
		return err
	}
	for _, batch := range merged {
		if err := ctx.Err(); err != nil {
			return err
		}
		if err := b.SaveHash(ctx, batch.Key, batch.Fields, batch.TTL); err != nil {
			return err
		}
	}
	return nil
}

func mergeHashSaveBatches(batches []HashSaveBatch) ([]HashSaveBatch, error) {
	if len(batches) == 0 {
		return nil, nil
	}
	merged := make(map[string]HashSaveBatch, len(batches))
	for _, batch := range batches {
		if len(batch.Fields) == 0 {
			continue
		}
		key := strings.TrimSpace(batch.Key)
		if key == "" {
			return nil, ErrRecordKeyRequired
		}
		fields, err := cloneHashFields(batch.Fields)
		if err != nil {
			return nil, err
		}
		if len(fields) == 0 {
			continue
		}
		current := merged[key]
		if current.Fields == nil {
			current = HashSaveBatch{
				Key:    key,
				Fields: make(map[string][]byte, len(fields)),
			}
		}
		for name, data := range fields {
			current.Fields[name] = append([]byte(nil), data...)
		}
		current.TTL = batch.TTL
		merged[key] = current
	}
	keys := make([]string, 0, len(merged))
	for key := range merged {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	out := make([]HashSaveBatch, 0, len(keys))
	for _, key := range keys {
		batch := merged[key]
		batch.Fields = cloneBytesMap(batch.Fields)
		out = append(out, batch)
	}
	return out, nil
}

// Check 检查上下文是否已取消。
func (b *MemoryHashBackend) Check(ctx context.Context) error {
	if b == nil {
		return ErrHashBackendRequired
	}
	ctx = contextOrBackground(ctx)
	return ctx.Err()
}

// Snapshot 返回全部 hash 记录副本。
func (b *MemoryHashBackend) Snapshot() map[string]map[string][]byte {
	if b == nil {
		return map[string]map[string][]byte{}
	}
	b.mu.RLock()
	out := make(map[string]map[string][]byte, len(b.records))
	for key, fields := range b.records {
		out[key] = cloneBytesMap(fields)
	}
	b.mu.RUnlock()
	return out
}

func (b *MemoryHashBackend) ensureReadyLocked() {
	if b.records == nil {
		b.records = make(map[string]map[string][]byte)
	}
	if b.ttl == nil {
		b.ttl = make(map[string]time.Duration)
	}
}

func cloneHashFields(fields map[string][]byte) (map[string][]byte, error) {
	if len(fields) == 0 {
		return nil, nil
	}
	out := make(map[string][]byte, len(fields))
	names := make([]string, 0, len(fields))
	for rawName := range fields {
		names = append(names, rawName)
	}
	sort.Strings(names)
	for _, rawName := range names {
		name := strings.TrimSpace(rawName)
		if name == "" {
			return nil, ErrRecordKeyRequired
		}
		if _, exists := out[name]; exists {
			continue
		}
		data := fields[rawName]
		if data == nil {
			return nil, errors.New("hash field value is nil")
		}
		out[name] = append([]byte(nil), data...)
	}
	return out, nil
}

func cloneBytesMap(in map[string][]byte) map[string][]byte {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string][]byte, len(in))
	for key, data := range in {
		out[key] = append([]byte(nil), data...)
	}
	return out
}

func normHashFields(fields []string) []string {
	if len(fields) == 0 {
		return nil
	}
	seen := make(map[string]struct{}, len(fields))
	for _, field := range fields {
		field = strings.TrimSpace(field)
		if field != "" {
			seen[field] = struct{}{}
		}
	}
	out := make([]string, 0, len(seen))
	for field := range seen {
		out = append(out, field)
	}
	sort.Strings(out)
	return out
}