JCS-pub/common/pkgs/storage/local/shard_store.go

package local

import (
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/inhies/go-bytesize"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	stgtypes "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

type ShardStore struct {
	detail          *jcstypes.UserSpaceDetail
	stgRoot         string
	storeAbsRoot    string
	lock            sync.Mutex
	done            chan any
	totalShardsCnt  int64
	totalShardsSize int64
}

func NewShardStore(root string, detail *jcstypes.UserSpaceDetail) (*ShardStore, error) {
	storeAbsRoot, err := filepath.Abs(filepath.Join(root, detail.UserSpace.WorkingDir.OSString(), stgtypes.ShardStoreWorkingDir))
	if err != nil {
		return nil, fmt.Errorf("get abs root: %w", err)
	}

	return &ShardStore{
		detail:       detail,
		stgRoot:      root,
		storeAbsRoot: storeAbsRoot,
		done:         make(chan any, 1),
	}, nil
}

func (s *ShardStore) Start(ch *stgtypes.StorageEventChan) {
	s.updateStats()
	s.getLogger().Infof("component start, root: %v, max size: %v, shards count: %v, total size: %v",
		s.storeAbsRoot, s.detail.UserSpace.ShardStore.MaxSize, s.totalShardsCnt, s.totalShardsSize,
	)
}

func (s *ShardStore) Stop() {
	s.getLogger().Infof("component stop")
}

func (s *ShardStore) Store(path jcstypes.JPath, hash jcstypes.FileHash, size int64) (stgtypes.FileInfo, error) {
	fullTempPath := filepath.Join(s.stgRoot, path.OSString())

	s.lock.Lock()
	defer s.lock.Unlock()

	log := s.getLogger()

	log.Debugf("%v bypass uploaded, size: %v, hash: %v", fullTempPath, size, hash)

	blockDir := s.getFileDirFromHash(hash)
	err := os.MkdirAll(blockDir, 0755)
	if err != nil {
		log.Warnf("make block dir %v: %v", blockDir, err)
		return stgtypes.FileInfo{}, fmt.Errorf("making block dir: %w", err)
	}

	newPath := filepath.Join(blockDir, string(hash))
	_, err = os.Stat(newPath)
	if os.IsNotExist(err) {
		err = os.Rename(fullTempPath, newPath)
		if err != nil {
			log.Warnf("rename %v to %v: %v", fullTempPath, newPath, err)
			return stgtypes.FileInfo{}, fmt.Errorf("rename file: %w", err)
		}

		s.totalShardsCnt++
		s.totalShardsSize += size

	} else if err != nil {
		log.Warnf("get file %v stat: %v", newPath, err)
		return stgtypes.FileInfo{}, fmt.Errorf("get file stat: %w", err)
	}

	return stgtypes.FileInfo{
		Hash: hash,
		Size: size,
		Path: s.getJPathFromHash(hash),
	}, nil
}

func (s *ShardStore) Info(hash jcstypes.FileHash) (stgtypes.FileInfo, error) {
	s.lock.Lock()
	defer s.lock.Unlock()

	filePath := s.getFilePathFromHash(hash)
	info, err := os.Stat(filePath)
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	return stgtypes.FileInfo{
		Hash: hash,
		Size: info.Size(),
		Path: s.getJPathFromHash(hash),
	}, nil
}

func (s *ShardStore) ListAll() ([]stgtypes.FileInfo, error) {
	s.lock.Lock()
	defer s.lock.Unlock()

	var infos []stgtypes.FileInfo

	err := filepath.WalkDir(s.storeAbsRoot, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}

		if d.IsDir() {
			return nil
		}

		info, err := d.Info()
		if err != nil {
			return err
		}

		fileHash, err := jcstypes.ParseHash(filepath.Base(info.Name()))
		if err != nil {
			return nil
		}

		infos = append(infos, stgtypes.FileInfo{
			Hash: fileHash,
			Size: info.Size(),
			Path: s.getJPathFromHash(fileHash),
		})
		return nil
	})
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return nil, err
	}

	return infos, nil
}

func (s *ShardStore) GC(avaiables []jcstypes.FileHash) error {
	startTime := time.Now()

	s.lock.Lock()
	defer s.lock.Unlock()

	avais := make(map[jcstypes.FileHash]bool)
	for _, hash := range avaiables {
		avais[hash] = true
	}

	totalCnt := 0
	totalSize := int64(0)

	err := filepath.WalkDir(s.storeAbsRoot, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}

		if d.IsDir() {
			return nil
		}

		info, err := d.Info()
		if err != nil {
			return err
		}

		fileHash, err := jcstypes.ParseHash(filepath.Base(info.Name()))
		if err != nil {
			return nil
		}

		if !avais[fileHash] {
			err = os.Remove(path)
			if err != nil {
				s.getLogger().Warnf("remove file %v: %v", path, err)
			} else {
				totalCnt++
				totalSize += info.Size()
				s.totalShardsCnt--
				s.totalShardsSize -= info.Size()
			}
		}

		return nil
	})
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}

	s.getLogger().Infof(
		"gc %v(size: %v), time: %v, remains: %v(size: %v)",
		totalCnt, bytesize.ByteSize(totalSize), time.Since(startTime), s.totalShardsCnt, bytesize.ByteSize(s.totalShardsSize),
	)

	// TODO 无法保证原子性，所以删除失败只打日志
	return nil
}

func (s *ShardStore) Stats() stgtypes.ShardStoreStats {
	s.lock.Lock()
	defer s.lock.Unlock()

	return stgtypes.ShardStoreStats{
		Status:    stgtypes.StatusOK,
		FileCount: s.totalShardsCnt,
		TotalSize: s.detail.UserSpace.ShardStore.MaxSize,
		UsedSize:  s.totalShardsSize,
	}
}

func (s *ShardStore) getLogger() logger.Logger {
	return logger.WithField("ShardStore", "Local").WithField("Storage", s.detail.UserSpace.Storage.String())
}

func (s *ShardStore) getFileDirFromHash(hash jcstypes.FileHash) string {
	return filepath.Join(s.storeAbsRoot, hash.GetHashPrefix(2))
}

func (s *ShardStore) getFilePathFromHash(hash jcstypes.FileHash) string {
	return filepath.Join(s.storeAbsRoot, hash.GetHashPrefix(2), string(hash))
}

func (s *ShardStore) getJPathFromHash(hash jcstypes.FileHash) jcstypes.JPath {
	return s.detail.UserSpace.WorkingDir.PushNew(stgtypes.ShardStoreWorkingDir, hash.GetHashPrefix(2), string(hash))
}

func (s *ShardStore) updateStats() {
	log := s.getLogger()

	totalCnt := int64(0)
	totalSize := int64(0)

	err := filepath.WalkDir(s.storeAbsRoot, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}

		if d.IsDir() {
			return nil
		}

		info, err := d.Info()
		if err != nil {
			return err
		}

		_, err = jcstypes.ParseHash(filepath.Base(info.Name()))
		if err != nil {
			return nil
		}

		totalCnt++
		totalSize += info.Size()
		return nil
	})
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		log.Warnf("walk dir %v: %v", s.storeAbsRoot, err)
		return
	}

	s.lock.Lock()
	defer s.lock.Unlock()

	s.totalShardsCnt = totalCnt
	s.totalShardsSize = totalSize
}