cloudreve/pkg/cluster/controller.go

package cluster

import (
	"bytes"
	"encoding/gob"
	"fmt"
	model "github.com/cloudreve/Cloudreve/v3/models"
	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
	"github.com/cloudreve/Cloudreve/v3/pkg/request"
	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
	"github.com/jinzhu/gorm"
	"net/url"
	"sync"
)

var DefaultController Controller

// Controller controls communications between master and slave
type Controller interface {
	// Handle heartbeat sent from master
	HandleHeartBeat(*serializer.NodePingReq) (serializer.NodePingResp, error)

	// Get Aria2 Instance by master node ID
	GetAria2Instance(string) (common.Aria2, error)

	// Send event change message to master node
	SendNotification(string, string, mq.Message) error

	// Submit async task into task pool
	SubmitTask(string, interface{}, string, func(interface{})) error

	// Get master node info
	GetMasterInfo(string) (*MasterInfo, error)

	// Get master OneDrive policy credential
	GetOneDriveToken(string, uint) (string, error)
}

type slaveController struct {
	masters map[string]MasterInfo
	lock    sync.RWMutex
}

// info of master node
type MasterInfo struct {
	ID  string
	TTL int
	URL *url.URL
	// used to invoke aria2 rpc calls
	Instance Node
	Client   request.Client

	jobTracker map[string]bool
}

func InitController() {
	DefaultController = &slaveController{
		masters: make(map[string]MasterInfo),
	}
	gob.Register(rpc.StatusInfo{})
}

func (c *slaveController) HandleHeartBeat(req *serializer.NodePingReq) (serializer.NodePingResp, error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	req.Node.AfterFind()

	// close old node if exist
	origin, ok := c.masters[req.SiteID]

	if (ok && req.IsUpdate) || !ok {
		if ok {
			origin.Instance.Kill()
		}

		masterUrl, err := url.Parse(req.SiteURL)
		if err != nil {
			return serializer.NodePingResp{}, err
		}

		c.masters[req.SiteID] = MasterInfo{
			ID:  req.SiteID,
			URL: masterUrl,
			TTL: req.CredentialTTL,
			Client: request.NewClient(
				request.WithEndpoint(masterUrl.String()),
				request.WithSlaveMeta(fmt.Sprintf("%d", req.Node.ID)),
				request.WithCredential(auth.HMACAuth{
					SecretKey: []byte(req.Node.MasterKey),
				}, int64(req.CredentialTTL)),
			),
			jobTracker: make(map[string]bool),
			Instance: NewNodeFromDBModel(&model.Node{
				Model:                  gorm.Model{ID: req.Node.ID},
				MasterKey:              req.Node.MasterKey,
				Type:                   model.MasterNodeType,
				Aria2Enabled:           req.Node.Aria2Enabled,
				Aria2OptionsSerialized: req.Node.Aria2OptionsSerialized,
			}),
		}
	}

	return serializer.NodePingResp{}, nil
}

func (c *slaveController) GetAria2Instance(id string) (common.Aria2, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	if node, ok := c.masters[id]; ok {
		return node.Instance.GetAria2Instance(), nil
	}

	return nil, ErrMasterNotFound
}

func (c *slaveController) SendNotification(id, subject string, msg mq.Message) error {
	c.lock.RLock()

	if node, ok := c.masters[id]; ok {
		c.lock.RUnlock()

		body := bytes.Buffer{}
		enc := gob.NewEncoder(&body)
		if err := enc.Encode(&msg); err != nil {
			return err
		}

		res, err := node.Client.Request(
			"PUT",
			fmt.Sprintf("/api/v3/slave/notification/%s", subject),
			&body,
		).CheckHTTPResponse(200).DecodeResponse()
		if err != nil {
			return err
		}

		if res.Code != 0 {
			return serializer.NewErrorFromResponse(res)
		}

		return nil
	}

	c.lock.RUnlock()
	return ErrMasterNotFound
}

// SubmitTask 提交异步任务
func (c *slaveController) SubmitTask(id string, job interface{}, hash string, submitter func(interface{})) error {
	c.lock.RLock()
	defer c.lock.RUnlock()

	if node, ok := c.masters[id]; ok {
		if _, ok := node.jobTracker[hash]; ok {
			// 任务已存在，直接返回
			return nil
		}

		node.jobTracker[hash] = true
		submitter(job)
		return nil
	}

	return ErrMasterNotFound
}

// GetMasterInfo 获取主机节点信息
func (c *slaveController) GetMasterInfo(id string) (*MasterInfo, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	if node, ok := c.masters[id]; ok {
		return &node, nil
	}

	return nil, ErrMasterNotFound
}

// GetOneDriveToken 获取主机OneDrive凭证
func (c *slaveController) GetOneDriveToken(id string, policyID uint) (string, error) {
	c.lock.RLock()

	if node, ok := c.masters[id]; ok {
		c.lock.RUnlock()

		res, err := node.Client.Request(
			"GET",
			fmt.Sprintf("/api/v3/slave/credential/onedrive/%d", policyID),
			nil,
		).CheckHTTPResponse(200).DecodeResponse()
		if err != nil {
			return "", err
		}

		if res.Code != 0 {
			return "", serializer.NewErrorFromResponse(res)
		}

		return res.Data.(string), nil
	}

	c.lock.RUnlock()
	return "", ErrMasterNotFound
}