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

package pubshards

import (
	"context"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/glebarez/sqlite"
	"gitlink.org.cn/cloudream/common/pkgs/async"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	stgglb "gitlink.org.cn/cloudream/jcs-pub/common/globals"
	corrpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/coordinator"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/factory"
	stgtypes "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
	"golang.org/x/crypto/bcrypt"
	"gorm.io/gorm"
)

type Pool struct {
	cfg          Config
	localHubID   jcstypes.HubID
	stores       map[jcstypes.PubShardsID]*LoadedStore
	stgEventChan *stgtypes.StorageEventChan
	done         chan any
	lock         sync.Mutex
}

func New(cfg Config, localHubID jcstypes.HubID) *Pool {
	return &Pool{
		cfg:          cfg,
		localHubID:   localHubID,
		stores:       make(map[jcstypes.PubShardsID]*LoadedStore),
		stgEventChan: async.NewUnboundChannel[stgtypes.StorageEvent](),
		done:         make(chan any, 1),
	}
}

func (p *Pool) GetOrLoad(pubStoreID jcstypes.PubShardsID, password string) (*LoadedStore, error) {
	log := logger.WithField("Mod", "PubShards")

	p.lock.Lock()
	defer p.lock.Unlock()

	loaded := p.stores[pubStoreID]
	if loaded == nil {
		corCli := stgglb.CoordinatorRPCPool.Get()
		defer corCli.Release()

		resp, cerr := corCli.UserGetPubShards(context.Background(), &corrpc.UserGetPubShards{
			PubShardsID: pubStoreID,
			Password:    password,
		})
		if cerr != nil {
			return nil, cerr.ToError()
		}

		if resp.PubShards.MasterHub != p.localHubID {
			return nil, fmt.Errorf("this hub is not the master hub of the public shard store")
		}

		pwdHash, err := hex.DecodeString(resp.PubShards.Password)
		if err != nil {
			return nil, fmt.Errorf("decode password: %w", err)
		}

		detail := jcstypes.UserSpaceDetail{
			UserSpace: jcstypes.UserSpace{
				Name:       resp.PubShards.Name,
				Storage:    resp.PubShards.Storage,
				Credential: resp.PubShards.Credential,
				ShardStore: &resp.PubShards.ShardStore,
				Features:   resp.PubShards.Features,
				WorkingDir: resp.PubShards.WorkingDir,
			},
			RecommendHub: &resp.MasterHub,
		}

		blder := factory.GetBuilder(&detail)
		ss, err := blder.CreateShardStore(false)
		if err != nil {
			return nil, err
		}

		err = os.MkdirAll(p.cfg.DBDir, 0755)
		if err != nil {
			return nil, err
		}

		dbFilePath := filepath.Join(p.cfg.DBDir, fmt.Sprintf("%s.db", pubStoreID))
		db, err := gorm.Open(sqlite.Open(dbFilePath), &gorm.Config{})
		if err != nil {
			return nil, err
		}
		err = db.AutoMigrate(Shard{})
		if err != nil {
			return nil, fmt.Errorf("migrate Shard: %w", err)
		}
		err = db.AutoMigrate(UserRef{})
		if err != nil {
			return nil, fmt.Errorf("migrate UserRef: %w", err)
		}

		ss.Start(p.stgEventChan)

		loaded = &LoadedStore{
			ShardStore:       ss,
			Config:           resp.PubShards,
			PasswordHash:     pwdHash,
			ClientFileHashDB: db,
		}
		p.stores[pubStoreID] = loaded

		log.Infof("%v loaded", loaded.Config.String())

	} else {
		// 如果已经被加载，那么就要验证一下密码是否正确
		if bcrypt.CompareHashAndPassword(loaded.PasswordHash, []byte(password)) != nil {
			return nil, fmt.Errorf("wrong password")
		}
	}

	return loaded, nil
}

func (p *Pool) Start() {
	go func() {
		log := logger.WithField("Mod", "PubShards")

		ticker := time.NewTicker(time.Minute)
		defer ticker.Stop()

		gced := make(map[jcstypes.PubShardsID]bool)

	loop:
		for {
			select {
			case <-ticker.C:
			case <-p.done:
				break loop
			}

			// 凌晨5点开始GC
			if time.Now().Hour() != 5 {
				gced = make(map[jcstypes.PubShardsID]bool)
				continue
			}

			p.lock.Lock()
			for pubStoreID, loaded := range p.stores {
				if gced[pubStoreID] {
					continue
				}

				allHashes, err := loaded.GetAllHashes()
				if err != nil {
					log.Warnf("get all hashes of %v: %v", loaded.Config.String(), err)
					continue
				}

				err = loaded.ShardStore.GC(allHashes)
				if err != nil {
					log.Warnf("gc %v: %v", loaded.Config.String(), err)
					continue
				}

				gced[pubStoreID] = true
				log.Infof("%v gc done", loaded.Config.String())
			}
			p.lock.Unlock()
		}

		p.lock.Lock()
		for _, loaded := range p.stores {
			loaded.ShardStore.Stop()
			d, err := loaded.ClientFileHashDB.DB()
			if err != nil {
				log.Warnf("get sql db of %v: %v", loaded.Config.String(), err)
				continue
			}
			d.Close()
		}
		p.stores = make(map[jcstypes.PubShardsID]*LoadedStore)
		p.lock.Unlock()
	}()
}

func (p *Pool) Stop() {
	select {
	case p.done <- true:
	default:
	}
}