github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/aria2/aria2.go (about)

     1  package aria2
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/url"
     7  	"sync"
     8  	"time"
     9  
    10  	model "github.com/cloudreve/Cloudreve/v3/models"
    11  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
    12  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/monitor"
    13  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
    14  	"github.com/cloudreve/Cloudreve/v3/pkg/balancer"
    15  	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
    16  	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
    17  )
    18  
    19  // Instance 默认使用的Aria2处理实例
    20  var Instance common.Aria2 = &common.DummyAria2{}
    21  
    22  // LB 获取 Aria2 节点的负载均衡器
    23  var LB balancer.Balancer
    24  
    25  // Lock Instance的读写锁
    26  var Lock sync.RWMutex
    27  
    28  // GetLoadBalancer 返回供Aria2使用的负载均衡器
    29  func GetLoadBalancer() balancer.Balancer {
    30  	Lock.RLock()
    31  	defer Lock.RUnlock()
    32  	return LB
    33  }
    34  
    35  // Init 初始化
    36  func Init(isReload bool, pool cluster.Pool, mqClient mq.MQ) {
    37  	Lock.Lock()
    38  	LB = balancer.NewBalancer("RoundRobin")
    39  	Lock.Unlock()
    40  
    41  	if !isReload {
    42  		// 从数据库中读取未完成任务，创建监控
    43  		unfinished := model.GetDownloadsByStatus(common.Ready, common.Paused, common.Downloading, common.Seeding)
    44  
    45  		for i := 0; i < len(unfinished); i++ {
    46  			// 创建任务监控
    47  			monitor.NewMonitor(&unfinished[i], pool, mqClient)
    48  		}
    49  	}
    50  }
    51  
    52  // TestRPCConnection 发送测试用的 RPC 请求，测试服务连通性
    53  func TestRPCConnection(server, secret string, timeout int) (rpc.VersionInfo, error) {
    54  	// 解析RPC服务地址
    55  	rpcServer, err := url.Parse(server)
    56  	if err != nil {
    57  		return rpc.VersionInfo{}, fmt.Errorf("cannot parse RPC server: %w", err)
    58  	}
    59  
    60  	rpcServer.Path = "/jsonrpc"
    61  	caller, err := rpc.New(context.Background(), rpcServer.String(), secret, time.Duration(timeout)*time.Second, nil)
    62  	if err != nil {
    63  		return rpc.VersionInfo{}, fmt.Errorf("cannot initialize rpc connection: %w", err)
    64  	}
    65  
    66  	return caller.GetVersion()
    67  }