SaveAny-Bot/storage/load.go

package storage

import (
	"context"
	"fmt"
	"maps"
	"sync"

	"github.com/charmbracelet/log"
	"github.com/krau/SaveAny-Bot/config"
	storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
	"golang.org/x/sync/singleflight"
)

var (
	storageMu sync.RWMutex
	// Storages maps storage names to initialized storage instances.
	Storages = make(map[string]Storage)

	userStoragesMu sync.RWMutex
	// UserStorages maps user IDs to their available storage instances.
	UserStorages = make(map[int64][]Storage)

	initFlight singleflight.Group
)

// GetStorage returns the initialized storage instance for name, without
// creating one on demand.
func GetStorage(name string) (Storage, bool) {
	storageMu.RLock()
	defer storageMu.RUnlock()
	s, ok := Storages[name]
	return s, ok
}

// AllStorages returns a snapshot copy of all initialized storages.
func AllStorages() map[string]Storage {
	storageMu.RLock()
	defer storageMu.RUnlock()
	out := make(map[string]Storage, len(Storages))
	maps.Copy(out, Storages)
	return out
}

// GetStorageByName returns storage by name from cache or creates new one
// It should NOT be used to get storage for user, use GetStorageByUserIDAndName instead
func GetStorageByName(ctx context.Context, name string) (Storage, error) {
	if name == "" {
		return nil, ErrStorageNameEmpty
	}

	storageMu.RLock()
	storage, ok := Storages[name]
	storageMu.RUnlock()
	if ok {
		return storage, nil
	}
	cfg := config.C().GetStorageByName(name)
	if cfg == nil {
		return nil, fmt.Errorf("storage %s not found", name)
	}

	// Merge concurrent first-time initializations.
	v, err, _ := initFlight.Do("storage:"+name, func() (any, error) {
		storageMu.RLock()
		if existing, ok := Storages[name]; ok {
			storageMu.RUnlock()
			return existing, nil
		}
		storageMu.RUnlock()
		storage, err := NewStorage(ctx, cfg)
		if err != nil {
			return nil, err
		}
		storageMu.Lock()
		defer storageMu.Unlock()
		if existing, ok := Storages[name]; ok {
			return existing, nil
		}
		Storages[name] = storage
		return storage, nil
	})
	if err != nil {
		return nil, err
	}
	return v.(Storage), nil
}

// 检查 user 是否可用指定的 storage, 若不可用则返回未找到错误
func GetStorageByUserIDAndName(ctx context.Context, chatID int64, name string) (Storage, error) {
	if name == "" {
		return nil, ErrStorageNameEmpty
	}

	if !config.C().HasStorage(chatID, name) {
		return nil, fmt.Errorf("no storage %s for user %d", name, chatID)
	}

	return GetStorageByName(ctx, name)
}

func GetUserStorages(ctx context.Context, chatID int64) []Storage {
	if chatID <= 0 {
		return nil
	}
	userStoragesMu.RLock()
	cached, ok := UserStorages[chatID]
	userStoragesMu.RUnlock()
	if ok {
		return cached
	}
	var storages []Storage
	for _, name := range config.C().GetStorageNamesByUserID(chatID) {
		storage, err := GetStorageByName(ctx, name)
		if err != nil {
			continue
		}
		storages = append(storages, storage)
	}
	return storages
}

func LoadStorages(ctx context.Context) {
	logger := log.FromContext(ctx)
	logger.Debug("loading storages...")
	for _, storage := range config.C().Storages {
		_, err := GetStorageByName(ctx, storage.GetName())
		if err != nil {
			logger.Errorf("failed to load storage %s: %v", storage.GetName(), err)
		}
	}
	storageMu.RLock()
	loaded := len(Storages)
	storageMu.RUnlock()
	logger.Infof("successfully loaded %d storages", loaded)
	for user := range config.C().GetUsersID() {
		uid := int64(user)
		storages := GetUserStorages(ctx, uid)
		userStoragesMu.Lock()
		UserStorages[uid] = storages
		userStoragesMu.Unlock()
	}
}

// GetTelegramStorageByUserID returns the first enabled Telegram storage for the user
func GetTelegramStorageByUserID(ctx context.Context, chatID int64) (Storage, error) {
	storages := GetUserStorages(ctx, chatID)
	for _, stor := range storages {
		if stor.Type() == storenum.Telegram {
			return stor, nil
		}
	}
	return nil, fmt.Errorf("no telegram storage found for user %d", chatID)
}