JCS-pub/hub/internal/pubshards/pub_shards.go

package pubshards

import (
	"fmt"

	stgtypes "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type LoadedStore struct {
	ShardStore       stgtypes.ShardStore
	Config           jcstypes.PubShards
	PasswordHash     []byte
	ClientFileHashDB *gorm.DB
}

func (s *LoadedStore) StoreShard(userID jcstypes.UserID, path jcstypes.JPath, hash jcstypes.FileHash, size int64) (stgtypes.FileInfo, error) {
	info, err := s.ShardStore.Store(path, hash, size)
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	err = s.ClientFileHashDB.Transaction(func(tx *gorm.DB) error {
		err = tx.Clauses(clause.Insert{Modifier: "or ignore"}).Create(&Shard{
			Path: info.Path,
			Hash: hash,
			Size: size,
		}).Error
		if err != nil {
			return err
		}

		err = tx.Clauses(clause.Insert{Modifier: "or ignore"}).Create(&UserRef{
			UserID: userID,
			Hash:   hash,
		}).Error
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	return info, nil
}

func (s *LoadedStore) InfoShard(hash jcstypes.FileHash) (stgtypes.FileInfo, error) {
	return s.ShardStore.Info(hash)
}

func (s *LoadedStore) ListUserAll(userID jcstypes.UserID) ([]stgtypes.FileInfo, error) {
	var files []Shard
	err := s.ClientFileHashDB.Table("UserRef").Select("Shard.*").Joins("join Shard on UserRef.Hash = Shard.Hash").Where("UserRef.UserID = ?", userID).Find(&files).Error
	if err != nil {
		return nil, err
	}

	infos := make([]stgtypes.FileInfo, len(files))
	for i, file := range files {
		infos[i] = stgtypes.FileInfo{
			Path: file.Path,
			Size: file.Size,
			Hash: file.Hash,
		}
	}

	return infos, nil
}

func (s *LoadedStore) GC(userID jcstypes.UserID, fileHashes []jcstypes.FileHash) error {
	// 从总体上看，GC是在ListUserAll之后调用的，FileHashes的内容只会小于已有的内容，所以创建Ref时可以不检查Hash是否存在
	return s.ClientFileHashDB.Transaction(func(tx *gorm.DB) error {
		if err := tx.Delete(&UserRef{}, "UserID = ?", userID).Error; err != nil {
			return fmt.Errorf("delete all user refs: %w", err)
		}

		refs := make([]UserRef, len(fileHashes))
		for i, hash := range fileHashes {
			refs[i] = UserRef{
				UserID: userID,
				Hash:   hash,
			}
		}

		return tx.Clauses(clause.Insert{Modifier: "or ignore"}).Create(&refs).Error
	})
}

func (s *LoadedStore) GetUserStats(userID jcstypes.UserID) stgtypes.ShardStoreStats {
	var stats struct {
		Cnt  int64 `gorm:"column:Cnt; type:bigint;"`
		Size int64 `gorm:"column:Size; type:bigint;"`
	}

	s.ClientFileHashDB.Table("Shard").Select("Count(Hash) as Cnt, Sum(Size) as Size").Find(&stats)

	return stgtypes.ShardStoreStats{
		Status:    stgtypes.StatusOK,
		FileCount: stats.Cnt,
		TotalSize: s.Config.ShardStore.MaxSize,
		UsedSize:  stats.Size,
	}
}

func (s *LoadedStore) CreateRefs(userID jcstypes.UserID, refs []jcstypes.FileHash) ([]jcstypes.FileHash, error) {
	var invalidHashes []jcstypes.FileHash
	err := s.ClientFileHashDB.Transaction(func(tx *gorm.DB) error {
		var avaiHashes []jcstypes.FileHash
		err := tx.Table("Shard").Select("Hash").Where("Hash IN ?", refs).Find(&avaiHashes).Error
		if err != nil {
			return fmt.Errorf("check avaiable hashes: %w", err)
		}

		var avaiRefs []UserRef
		avaiHashesMp := make(map[jcstypes.FileHash]bool)
		for _, hash := range avaiHashes {
			avaiHashesMp[hash] = true
			avaiRefs = append(avaiRefs, UserRef{
				UserID: userID,
				Hash:   hash,
			})
		}

		for _, hash := range refs {
			if _, ok := avaiHashesMp[hash]; !ok {
				invalidHashes = append(invalidHashes, hash)
			}
		}

		return tx.Clauses(clause.Insert{Modifier: "or ignore"}).Create(&avaiRefs).Error
	})
	return invalidHashes, err
}

func (s *LoadedStore) GetAllHashes() ([]jcstypes.FileHash, error) {
	var hashes []jcstypes.FileHash
	return hashes, s.ClientFileHashDB.Table("Shard").Select("Hash").Find(&hashes).Error
}

type Shard struct {
	Hash jcstypes.FileHash `gorm:"column:Hash; type:char(68); primaryKey; not null; index" json:"hash"`
	Path jcstypes.JPath    `gorm:"column:Path; type:varchar(1024); not null; serializer:string" json:"path"`
	Size int64             `gorm:"column:Size; type:bigint; not null" json:"size"`
}

func (t Shard) TableName() string {
	return "Shard"
}

type UserRef struct {
	UserID jcstypes.UserID   `gorm:"column:UserID; type:bigint; primaryKey; not null" json:"userID"`
	Hash   jcstypes.FileHash `gorm:"column:Hash; type:char(68); primaryKey; not null; index" json:"hash"`
}

func (t UserRef) TableName() string {
	return "UserRef"
}