diff --git a/internal/iplibrary/manager_ip_list.go b/internal/iplibrary/manager_ip_list.go index 85c122d..8458d6c 100644 --- a/internal/iplibrary/manager_ip_list.go +++ b/internal/iplibrary/manager_ip_list.go @@ -144,7 +144,7 @@ func (this *IPListManager) fetch() (hasNext bool, err error) { if err != nil { return false, err } - itemsResp, err := rpcClient.IPItemRPC().ListIPItemsAfterVersion(rpcClient.Context(), &pb.ListIPItemsAfterVersionRequest{ + itemsResp, err := rpcClient.IPItemRPC.ListIPItemsAfterVersion(rpcClient.Context(), &pb.ListIPItemsAfterVersionRequest{ Version: this.version, Size: this.pageSize, }) diff --git a/internal/metrics/task.go b/internal/metrics/task.go index aa4a45b..8d33450 100644 --- a/internal/metrics/task.go +++ b/internal/metrics/task.go @@ -424,7 +424,7 @@ func (this *Task) Upload(pauseDuration time.Duration) error { return nil, err } - _, err = rpcClient.MetricStatRPC().UploadMetricStats(rpcClient.Context(), &pb.UploadMetricStatsRequest{ + _, err = rpcClient.MetricStatRPC.UploadMetricStats(rpcClient.Context(), &pb.UploadMetricStatsRequest{ MetricStats: pbStats, Time: currentTime, ServerId: serverId, diff --git a/internal/monitor/value_queue.go b/internal/monitor/value_queue.go index 52b59e6..b46936b 100644 --- a/internal/monitor/value_queue.go +++ b/internal/monitor/value_queue.go @@ -69,7 +69,7 @@ func (this *ValueQueue) Loop() error { } for value := range this.valuesChan { - _, err = rpcClient.NodeValueRPC().CreateNodeValue(rpcClient.Context(), &pb.CreateNodeValueRequest{ + _, err = rpcClient.NodeValueRPC.CreateNodeValue(rpcClient.Context(), &pb.CreateNodeValueRequest{ Item: value.Item, ValueJSON: value.ValueJSON, CreatedAt: value.CreatedAt, diff --git a/internal/monitor/value_queue_test.go b/internal/monitor/value_queue_test.go index c35c7d0..6e3b890 100644 --- a/internal/monitor/value_queue_test.go +++ b/internal/monitor/value_queue_test.go @@ -16,9 +16,9 @@ func TestValueQueue_RPC(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = rpcClient.NodeValueRPC().CreateNodeValue(rpcClient.Context(), &pb.CreateNodeValueRequest{}) + _, err = rpcClient.NodeValueRPC.CreateNodeValue(rpcClient.Context(), &pb.CreateNodeValueRequest{}) if err != nil { - statusErr, ok:= status.FromError(err) + statusErr, ok := status.FromError(err) if ok { logs.Println(statusErr.Code()) } diff --git a/internal/nodes/api_stream.go b/internal/nodes/api_stream.go index 9fdf33a..ac64c9d 100644 --- a/internal/nodes/api_stream.go +++ b/internal/nodes/api_stream.go @@ -73,7 +73,7 @@ func (this *APIStream) loop() error { cancelFunc() }() - nodeStream, err := rpcClient.NodeRPC().NodeStream(ctx) + nodeStream, err := rpcClient.NodeRPC.NodeStream(ctx) if err != nil { if this.isQuiting { return nil diff --git a/internal/nodes/http_access_log_queue.go b/internal/nodes/http_access_log_queue.go index a091035..6fb0029 100644 --- a/internal/nodes/http_access_log_queue.go +++ b/internal/nodes/http_access_log_queue.go @@ -96,7 +96,7 @@ Loop: this.rpcClient = client } - _, err := this.rpcClient.HTTPAccessLogRPC().CreateHTTPAccessLogs(this.rpcClient.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: accessLogs}) + _, err := this.rpcClient.HTTPAccessLogRPC.CreateHTTPAccessLogs(this.rpcClient.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: accessLogs}) if err != nil { // 是否包含了invalid UTF-8 if strings.Contains(err.Error(), "string field contains invalid UTF-8") { @@ -105,7 +105,7 @@ Loop: } // 重新提交 - _, err = this.rpcClient.HTTPAccessLogRPC().CreateHTTPAccessLogs(this.rpcClient.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: accessLogs}) + _, err = this.rpcClient.HTTPAccessLogRPC.CreateHTTPAccessLogs(this.rpcClient.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: accessLogs}) return err } diff --git a/internal/nodes/http_access_log_queue_test.go b/internal/nodes/http_access_log_queue_test.go index b56f8b4..d2d81fa 100644 --- a/internal/nodes/http_access_log_queue_test.go +++ b/internal/nodes/http_access_log_queue_test.go @@ -55,7 +55,7 @@ func TestHTTPAccessLogQueue_Push(t *testing.T) { // logs.PrintAsJSON(accessLog) //t.Log(strings.ToValidUTF8(string(utf8Bytes), "")) - _, err = client.HTTPAccessLogRPC().CreateHTTPAccessLogs(client.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: []*pb.HTTPAccessLog{ + _, err = client.HTTPAccessLogRPC.CreateHTTPAccessLogs(client.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: []*pb.HTTPAccessLog{ accessLog, }}) if err != nil { @@ -99,7 +99,7 @@ func TestHTTPAccessLogQueue_Push2(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = client.HTTPAccessLogRPC().CreateHTTPAccessLogs(client.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: []*pb.HTTPAccessLog{ + _, err = client.HTTPAccessLogRPC.CreateHTTPAccessLogs(client.Context(), &pb.CreateHTTPAccessLogsRequest{HttpAccessLogs: []*pb.HTTPAccessLog{ accessLog, }}) if err != nil { diff --git a/internal/nodes/http_cache_task_manager.go b/internal/nodes/http_cache_task_manager.go index 9260781..b3c4f68 100644 --- a/internal/nodes/http_cache_task_manager.go +++ b/internal/nodes/http_cache_task_manager.go @@ -81,7 +81,7 @@ func (this *HTTPCacheTaskManager) Start() { if rpcClient != nil { for taskReq := range this.taskQueue { - _, err := rpcClient.ServerRPC().PurgeServerCache(rpcClient.Context(), taskReq) + _, err := rpcClient.ServerRPC.PurgeServerCache(rpcClient.Context(), taskReq) if err != nil { remotelogs.Error("HTTP_CACHE_TASK_MANAGER", "create purge task failed: "+err.Error()) } @@ -104,7 +104,7 @@ func (this *HTTPCacheTaskManager) Loop() error { return err } - resp, err := rpcClient.HTTPCacheTaskKeyRPC().FindDoingHTTPCacheTaskKeys(rpcClient.Context(), &pb.FindDoingHTTPCacheTaskKeysRequest{}) + resp, err := rpcClient.HTTPCacheTaskKeyRPC.FindDoingHTTPCacheTaskKeys(rpcClient.Context(), &pb.FindDoingHTTPCacheTaskKeysRequest{}) if err != nil { // 忽略连接错误 if rpc.IsConnError(err) { @@ -135,7 +135,7 @@ func (this *HTTPCacheTaskManager) Loop() error { pbResults = append(pbResults, pbResult) } - _, err = rpcClient.HTTPCacheTaskKeyRPC().UpdateHTTPCacheTaskKeysStatus(rpcClient.Context(), &pb.UpdateHTTPCacheTaskKeysStatusRequest{KeyResults: pbResults}) + _, err = rpcClient.HTTPCacheTaskKeyRPC.UpdateHTTPCacheTaskKeysStatus(rpcClient.Context(), &pb.UpdateHTTPCacheTaskKeysStatusRequest{KeyResults: pbResults}) if err != nil { return err } diff --git a/internal/nodes/http_request_acme.go b/internal/nodes/http_request_acme.go index 75f8c06..2b831ed 100644 --- a/internal/nodes/http_request_acme.go +++ b/internal/nodes/http_request_acme.go @@ -21,7 +21,7 @@ func (this *HTTPRequest) doACME() (shouldStop bool) { return false } - keyResp, err := rpcClient.ACMEAuthenticationRPC().FindACMEAuthenticationKeyWithToken(rpcClient.Context(), &pb.FindACMEAuthenticationKeyWithTokenRequest{Token: token}) + keyResp, err := rpcClient.ACMEAuthenticationRPC.FindACMEAuthenticationKeyWithToken(rpcClient.Context(), &pb.FindACMEAuthenticationKeyWithTokenRequest{Token: token}) if err != nil { remotelogs.Error("RPC", "[ACME]read key for token failed: "+err.Error()) return false diff --git a/internal/nodes/ip_library_updater.go b/internal/nodes/ip_library_updater.go index df0c303..8a0457b 100644 --- a/internal/nodes/ip_library_updater.go +++ b/internal/nodes/ip_library_updater.go @@ -46,7 +46,7 @@ func (this *IPLibraryUpdater) FindLatestFile() (code string, fileId int64, err e if err != nil { return "", 0, err } - resp, err := rpcClient.IPLibraryArtifactRPC().FindPublicIPLibraryArtifact(rpcClient.Context(), &pb.FindPublicIPLibraryArtifactRequest{}) + resp, err := rpcClient.IPLibraryArtifactRPC.FindPublicIPLibraryArtifact(rpcClient.Context(), &pb.FindPublicIPLibraryArtifactRequest{}) if err != nil { return "", 0, err } @@ -68,12 +68,12 @@ func (this *IPLibraryUpdater) DownloadFile(fileId int64, writer io.Writer) error return err } - chunkIdsResp, err := rpcClient.FileChunkRPC().FindAllFileChunkIds(rpcClient.Context(), &pb.FindAllFileChunkIdsRequest{FileId: fileId}) + chunkIdsResp, err := rpcClient.FileChunkRPC.FindAllFileChunkIds(rpcClient.Context(), &pb.FindAllFileChunkIdsRequest{FileId: fileId}) if err != nil { return err } for _, chunkId := range chunkIdsResp.FileChunkIds { - chunkResp, err := rpcClient.FileChunkRPC().DownloadFileChunk(rpcClient.Context(), &pb.DownloadFileChunkRequest{FileChunkId: chunkId}) + chunkResp, err := rpcClient.FileChunkRPC.DownloadFileChunk(rpcClient.Context(), &pb.DownloadFileChunkRequest{FileChunkId: chunkId}) if err != nil { return err } diff --git a/internal/nodes/node.go b/internal/nodes/node.go index ec19c61..fcf81d1 100644 --- a/internal/nodes/node.go +++ b/internal/nodes/node.go @@ -79,7 +79,7 @@ func (this *Node) Test() error { if err != nil { return errors.New("test rpc failed: " + err.Error()) } - _, err = rpcClient.APINodeRPC().FindCurrentAPINodeVersion(rpcClient.Context(), &pb.FindCurrentAPINodeVersionRequest{}) + _, err = rpcClient.APINodeRPC.FindCurrentAPINodeVersion(rpcClient.Context(), &pb.FindCurrentAPINodeVersionRequest{}) if err != nil { return errors.New("test rpc failed: " + err.Error()) } @@ -307,7 +307,7 @@ func (this *Node) loop() error { } var nodeCtx = rpcClient.Context() - tasksResp, err := rpcClient.NodeTaskRPC().FindNodeTasks(nodeCtx, &pb.FindNodeTasksRequest{}) + tasksResp, err := rpcClient.NodeTaskRPC.FindNodeTasks(nodeCtx, &pb.FindNodeTasksRequest{}) if err != nil { if rpc.IsConnError(err) && !Tea.IsTesting() { return nil @@ -325,7 +325,7 @@ func (this *Node) loop() error { } // 修改为已同步 - _, err = rpcClient.NodeTaskRPC().ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ + _, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ NodeTaskId: task.Id, IsOk: true, Error: "", @@ -344,13 +344,13 @@ func (this *Node) loop() error { err = this.syncConfig(task.Version) } if err != nil { - _, err = rpcClient.NodeTaskRPC().ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ + _, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ NodeTaskId: task.Id, IsOk: false, Error: err.Error(), }) } else { - _, err = rpcClient.NodeTaskRPC().ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ + _, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ NodeTaskId: task.Id, IsOk: true, Error: "", @@ -372,7 +372,7 @@ func (this *Node) loop() error { } // 修改为已同步 - _, err = rpcClient.NodeTaskRPC().ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ + _, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ NodeTaskId: task.Id, IsOk: true, Error: "", @@ -381,7 +381,7 @@ func (this *Node) loop() error { return err } case "nodeLevelChanged": - levelInfoResp, err := rpcClient.NodeRPC().FindNodeLevelInfo(nodeCtx, &pb.FindNodeLevelInfoRequest{}) + levelInfoResp, err := rpcClient.NodeRPC.FindNodeLevelInfo(nodeCtx, &pb.FindNodeLevelInfoRequest{}) if err != nil { return err } @@ -398,7 +398,7 @@ func (this *Node) loop() error { sharedNodeConfig.ParentNodes = parentNodes // 修改为已同步 - _, err = rpcClient.NodeTaskRPC().ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ + _, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ NodeTaskId: task.Id, IsOk: true, Error: "", @@ -407,7 +407,7 @@ func (this *Node) loop() error { return err } case "ddosProtectionChanged": - resp, err := rpcClient.NodeRPC().FindNodeDDoSProtection(nodeCtx, &pb.FindNodeDDoSProtectionRequest{}) + resp, err := rpcClient.NodeRPC.FindNodeDDoSProtection(nodeCtx, &pb.FindNodeDDoSProtectionRequest{}) if err != nil { return err } @@ -434,7 +434,7 @@ func (this *Node) loop() error { } // 修改为已同步 - _, err = rpcClient.NodeTaskRPC().ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ + _, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ NodeTaskId: task.Id, IsOk: true, Error: "", @@ -479,7 +479,7 @@ func (this *Node) syncConfig(taskVersion int64) error { nodeCtx := rpcClient.Context() // TODO 这里考虑只同步版本号有变更的 - configResp, err := rpcClient.NodeRPC().FindCurrentNodeConfig(nodeCtx, &pb.FindCurrentNodeConfigRequest{ + configResp, err := rpcClient.NodeRPC.FindCurrentNodeConfig(nodeCtx, &pb.FindCurrentNodeConfigRequest{ Version: -1, // 更新所有版本 Compress: true, NodeTaskVersion: taskVersion, @@ -570,7 +570,7 @@ func (this *Node) syncServerConfig(serverId int64) error { if err != nil { return err } - resp, err := rpcClient.ServerRPC().ComposeServerConfig(rpcClient.Context(), &pb.ComposeServerConfigRequest{ServerId: serverId}) + resp, err := rpcClient.ServerRPC.ComposeServerConfig(rpcClient.Context(), &pb.ComposeServerConfigRequest{ServerId: serverId}) if err != nil { return err } @@ -652,7 +652,7 @@ func (this *Node) checkClusterConfig() error { } logs.Println("[NODE]registering node to cluster ...") - resp, err := rpcClient.NodeRPC().RegisterClusterNode(rpcClient.ClusterContext(config.ClusterId, config.Secret), &pb.RegisterClusterNodeRequest{Name: HOSTNAME}) + resp, err := rpcClient.NodeRPC.RegisterClusterNode(rpcClient.ClusterContext(config.ClusterId, config.Secret), &pb.RegisterClusterNodeRequest{Name: HOSTNAME}) if err != nil { return err } diff --git a/internal/nodes/node_status_executor.go b/internal/nodes/node_status_executor.go index 856f6a7..1a69037 100644 --- a/internal/nodes/node_status_executor.go +++ b/internal/nodes/node_status_executor.go @@ -125,7 +125,7 @@ func (this *NodeStatusExecutor) update() { remotelogs.Error("NODE_STATUS", "failed to open rpc: "+err.Error()) return } - _, err = rpcClient.NodeRPC().UpdateNodeStatus(rpcClient.Context(), &pb.UpdateNodeStatusRequest{ + _, err = rpcClient.NodeRPC.UpdateNodeStatus(rpcClient.Context(), &pb.UpdateNodeStatusRequest{ StatusJSON: jsonData, }) if err != nil { diff --git a/internal/nodes/task_ocsp_update.go b/internal/nodes/task_ocsp_update.go index 9e7bd16..068738f 100644 --- a/internal/nodes/task_ocsp_update.go +++ b/internal/nodes/task_ocsp_update.go @@ -60,7 +60,7 @@ func (this *OCSPUpdateTask) Loop() error { return err } - resp, err := rpcClient.SSLCertRPC().ListUpdatedSSLCertOCSP(rpcClient.Context(), &pb.ListUpdatedSSLCertOCSPRequest{ + resp, err := rpcClient.SSLCertRPC.ListUpdatedSSLCertOCSP(rpcClient.Context(), &pb.ListUpdatedSSLCertOCSPRequest{ Version: this.version, Size: 100, }) diff --git a/internal/nodes/task_sync_api_nodes.go b/internal/nodes/task_sync_api_nodes.go index dbc8a92..8d590cb 100644 --- a/internal/nodes/task_sync_api_nodes.go +++ b/internal/nodes/task_sync_api_nodes.go @@ -81,7 +81,7 @@ func (this *SyncAPINodesTask) Loop() error { if err != nil { return err } - resp, err := rpcClient.APINodeRPC().FindAllEnabledAPINodes(rpcClient.Context(), &pb.FindAllEnabledAPINodesRequest{}) + resp, err := rpcClient.APINodeRPC.FindAllEnabledAPINodes(rpcClient.Context(), &pb.FindAllEnabledAPINodesRequest{}) if err != nil { return err } diff --git a/internal/nodes/upgrade_manager.go b/internal/nodes/upgrade_manager.go index 24e52e1..64dd2d8 100644 --- a/internal/nodes/upgrade_manager.go +++ b/internal/nodes/upgrade_manager.go @@ -126,7 +126,7 @@ func (this *UpgradeManager) install() error { var sum = "" var filename = "" for { - resp, err := client.NodeRPC().DownloadNodeInstallationFile(client.Context(), &pb.DownloadNodeInstallationFileRequest{ + resp, err := client.NodeRPC.DownloadNodeInstallationFile(client.Context(), &pb.DownloadNodeInstallationFileRequest{ Os: runtime.GOOS, Arch: runtime.GOARCH, ChunkOffset: offset, diff --git a/internal/remotelogs/utils.go b/internal/remotelogs/utils.go index 9774c3d..d2ef3d9 100644 --- a/internal/remotelogs/utils.go +++ b/internal/remotelogs/utils.go @@ -173,7 +173,6 @@ func ServerSuccess(serverId int64, tag string, description string, logType nodec } } - // ServerLog 打印服务相关日志信息 func ServerLog(serverId int64, tag string, description string, logType nodeconfigs.NodeLogType, params maps.Map) { logs.Println("[" + tag + "]" + description) @@ -253,6 +252,6 @@ Loop: return nil } - _, err = rpcClient.NodeLogRPC().CreateNodeLogs(rpcClient.Context(), &pb.CreateNodeLogsRequest{NodeLogs: logList}) + _, err = rpcClient.NodeLogRPC.CreateNodeLogs(rpcClient.Context(), &pb.CreateNodeLogsRequest{NodeLogs: logList}) return err } diff --git a/internal/rpc/rpc_client.go b/internal/rpc/rpc_client.go index 6dfe139..ce1ea3a 100644 --- a/internal/rpc/rpc_client.go +++ b/internal/rpc/rpc_client.go @@ -28,6 +28,27 @@ type RPCClient struct { conns []*grpc.ClientConn locker sync.RWMutex + + NodeRPC pb.NodeServiceClient + NodeLogRPC pb.NodeLogServiceClient + NodeTaskRPC pb.NodeTaskServiceClient + NodeValueRPC pb.NodeValueServiceClient + HTTPAccessLogRPC pb.HTTPAccessLogServiceClient + HTTPCacheTaskKeyRPC pb.HTTPCacheTaskKeyServiceClient + APINodeRPC pb.APINodeServiceClient + IPLibraryArtifactRPC pb.IPLibraryArtifactServiceClient + IPListRPC pb.IPListServiceClient + IPItemRPC pb.IPItemServiceClient + FileRPC pb.FileServiceClient + FileChunkRPC pb.FileChunkServiceClient + ACMEAuthenticationRPC pb.ACMEAuthenticationServiceClient + ServerRPC pb.ServerServiceClient + ServerDailyStatRPC pb.ServerDailyStatServiceClient + ServerBandwidthStatRPC pb.ServerBandwidthStatServiceClient + MetricStatRPC pb.MetricStatServiceClient + FirewallRPC pb.FirewallServiceClient + SSLCertRPC pb.SSLCertServiceClient + ScriptRPC pb.ScriptServiceClient } func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { @@ -35,10 +56,32 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { return nil, errors.New("api config should not be nil") } - client := &RPCClient{ + var client = &RPCClient{ apiConfig: apiConfig, } + // 初始化RPC实例 + client.NodeRPC = pb.NewNodeServiceClient(client) + client.NodeLogRPC = pb.NewNodeLogServiceClient(client) + client.NodeTaskRPC = pb.NewNodeTaskServiceClient(client) + client.NodeValueRPC = pb.NewNodeValueServiceClient(client) + client.HTTPAccessLogRPC = pb.NewHTTPAccessLogServiceClient(client) + client.HTTPCacheTaskKeyRPC = pb.NewHTTPCacheTaskKeyServiceClient(client) + client.APINodeRPC = pb.NewAPINodeServiceClient(client) + client.IPLibraryArtifactRPC = pb.NewIPLibraryArtifactServiceClient(client) + client.IPListRPC = pb.NewIPListServiceClient(client) + client.IPItemRPC = pb.NewIPItemServiceClient(client) + client.FileRPC = pb.NewFileServiceClient(client) + client.FileChunkRPC = pb.NewFileChunkServiceClient(client) + client.ACMEAuthenticationRPC = pb.NewACMEAuthenticationServiceClient(client) + client.ServerRPC = pb.NewServerServiceClient(client) + client.ServerDailyStatRPC = pb.NewServerDailyStatServiceClient(client) + client.ServerBandwidthStatRPC = pb.NewServerBandwidthStatServiceClient(client) + client.MetricStatRPC = pb.NewMetricStatServiceClient(client) + client.FirewallRPC = pb.NewFirewallServiceClient(client) + client.SSLCertRPC = pb.NewSSLCertServiceClient(client) + client.ScriptRPC = pb.NewScriptServiceClient(client) + err := client.init() if err != nil { return nil, err @@ -47,106 +90,10 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { return client, nil } -func (this *RPCClient) NodeRPC() pb.NodeServiceClient { - return pb.NewNodeServiceClient(this.pickConn()) -} - -func (this *RPCClient) NodeLogRPC() pb.NodeLogServiceClient { - return pb.NewNodeLogServiceClient(this.pickConn()) -} - -func (this *RPCClient) NodeTaskRPC() pb.NodeTaskServiceClient { - return pb.NewNodeTaskServiceClient(this.pickConn()) -} - -func (this *RPCClient) NodeValueRPC() pb.NodeValueServiceClient { - return pb.NewNodeValueServiceClient(this.pickConn()) -} - -func (this *RPCClient) HTTPAccessLogRPC() pb.HTTPAccessLogServiceClient { - return pb.NewHTTPAccessLogServiceClient(this.pickConn()) -} - -func (this *RPCClient) HTTPCacheTaskKeyRPC() pb.HTTPCacheTaskKeyServiceClient { - return pb.NewHTTPCacheTaskKeyServiceClient(this.pickConn()) -} - -func (this *RPCClient) APINodeRPC() pb.APINodeServiceClient { - return pb.NewAPINodeServiceClient(this.pickConn()) -} - -func (this *RPCClient) IPLibraryRPC() pb.IPLibraryServiceClient { - return pb.NewIPLibraryServiceClient(this.pickConn()) -} - -func (this *RPCClient) IPLibraryArtifactRPC() pb.IPLibraryArtifactServiceClient { - return pb.NewIPLibraryArtifactServiceClient(this.pickConn()) -} - -func (this *RPCClient) RegionCountryRPC() pb.RegionCountryServiceClient { - return pb.NewRegionCountryServiceClient(this.pickConn()) -} - -func (this *RPCClient) RegionProvinceRPC() pb.RegionProvinceServiceClient { - return pb.NewRegionProvinceServiceClient(this.pickConn()) -} - -func (this *RPCClient) RegionCityRPC() pb.RegionCityServiceClient { - return pb.NewRegionCityServiceClient(this.pickConn()) -} - -func (this *RPCClient) RegionProviderRPC() pb.RegionProviderServiceClient { - return pb.NewRegionProviderServiceClient(this.pickConn()) -} - -func (this *RPCClient) IPListRPC() pb.IPListServiceClient { - return pb.NewIPListServiceClient(this.pickConn()) -} - -func (this *RPCClient) IPItemRPC() pb.IPItemServiceClient { - return pb.NewIPItemServiceClient(this.pickConn()) -} - -func (this *RPCClient) FileRPC() pb.FileServiceClient { - return pb.NewFileServiceClient(this.pickConn()) -} - -func (this *RPCClient) FileChunkRPC() pb.FileChunkServiceClient { - return pb.NewFileChunkServiceClient(this.pickConn()) -} - -func (this *RPCClient) ACMEAuthenticationRPC() pb.ACMEAuthenticationServiceClient { - return pb.NewACMEAuthenticationServiceClient(this.pickConn()) -} - -func (this *RPCClient) ServerRPC() pb.ServerServiceClient { - return pb.NewServerServiceClient(this.pickConn()) -} - -func (this *RPCClient) ServerDailyStatRPC() pb.ServerDailyStatServiceClient { - return pb.NewServerDailyStatServiceClient(this.pickConn()) -} - -func (this *RPCClient) ServerBandwidthStatRPC() pb.ServerBandwidthStatServiceClient { - return pb.NewServerBandwidthStatServiceClient(this.pickConn()) -} - -func (this *RPCClient) MetricStatRPC() pb.MetricStatServiceClient { - return pb.NewMetricStatServiceClient(this.pickConn()) -} - -func (this *RPCClient) FirewallRPC() pb.FirewallServiceClient { - return pb.NewFirewallServiceClient(this.pickConn()) -} - -func (this *RPCClient) SSLCertRPC() pb.SSLCertServiceClient { - return pb.NewSSLCertServiceClient(this.pickConn()) -} - // Context 节点上下文信息 func (this *RPCClient) Context() context.Context { var ctx = context.Background() - m := maps.Map{ + var m = maps.Map{ "timestamp": time.Now().Unix(), "type": "node", "userId": 0, @@ -213,7 +160,7 @@ func (this *RPCClient) UpdateConfig(config *configs.APIConfig) error { // 初始化 func (this *RPCClient) init() error { // 重新连接 - conns := []*grpc.ClientConn{} + var conns = []*grpc.ClientConn{} for _, endpoint := range this.apiConfig.RPC.Endpoints { u, err := url.Parse(endpoint) if err != nil { @@ -252,20 +199,24 @@ func (this *RPCClient) pickConn() *grpc.ClientConn { // 检查连接状态 if len(this.conns) > 0 { - availableConns := []*grpc.ClientConn{} - for _, state := range []connectivity.State{connectivity.Ready, connectivity.Idle, connectivity.Connecting} { + var availableConns = []*grpc.ClientConn{} + for _, stateArray := range [][2]connectivity.State{ + {connectivity.Ready, connectivity.Idle}, // 优先Ready和Idle + {connectivity.Connecting, connectivity.Connecting}, + } { for _, conn := range this.conns { - if conn.GetState() == state { + var state = conn.GetState() + if state == stateArray[0] || state == stateArray[1] { availableConns = append(availableConns, conn) } } if len(availableConns) > 0 { - break + return this.randConn(availableConns) } } if len(availableConns) > 0 { - return availableConns[rands.Int(0, len(availableConns)-1)] + return this.randConn(availableConns) } // 关闭 @@ -285,5 +236,31 @@ func (this *RPCClient) pickConn() *grpc.ClientConn { return nil } - return this.conns[rands.Int(0, len(this.conns)-1)] + return this.randConn(this.conns) +} + +func (this *RPCClient) Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error { + var conn = this.pickConn() + if conn == nil { + return errors.New("can not get available grpc connection") + } + return conn.Invoke(ctx, method, args, reply, opts...) +} +func (this *RPCClient) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { + var conn = this.pickConn() + if conn == nil { + return nil, errors.New("can not get available grpc connection") + } + return conn.NewStream(ctx, desc, method, opts...) +} + +func (this *RPCClient) randConn(conns []*grpc.ClientConn) *grpc.ClientConn { + var l = len(conns) + if l == 0 { + return nil + } + if l == 1 { + return conns[0] + } + return conns[rands.Int(0, l-1)] } diff --git a/internal/stats/bandwidth_stat_manager.go b/internal/stats/bandwidth_stat_manager.go index 00912f8..088ba5f 100644 --- a/internal/stats/bandwidth_stat_manager.go +++ b/internal/stats/bandwidth_stat_manager.go @@ -96,7 +96,7 @@ func (this *BandwidthStatManager) Loop() error { if err != nil { return err } - _, err = rpcClient.ServerBandwidthStatRPC().UploadServerBandwidthStats(rpcClient.Context(), &pb.UploadServerBandwidthStatsRequest{ServerBandwidthStats: pbStats}) + _, err = rpcClient.ServerBandwidthStatRPC.UploadServerBandwidthStats(rpcClient.Context(), &pb.UploadServerBandwidthStatsRequest{ServerBandwidthStats: pbStats}) if err != nil { return err } diff --git a/internal/stats/http_request_stat_manager.go b/internal/stats/http_request_stat_manager.go index 833b98e..2067164 100644 --- a/internal/stats/http_request_stat_manager.go +++ b/internal/stats/http_request_stat_manager.go @@ -324,7 +324,7 @@ func (this *HTTPRequestStatManager) Upload() error { this.dailyFirewallRuleGroupMap = map[string]int64{} // 上传数据 - _, err = rpcClient.ServerRPC().UploadServerHTTPRequestStat(rpcClient.Context(), &pb.UploadServerHTTPRequestStatRequest{ + _, err = rpcClient.ServerRPC.UploadServerHTTPRequestStat(rpcClient.Context(), &pb.UploadServerHTTPRequestStatRequest{ Month: timeutil.Format("Ym"), Day: timeutil.Format("Ymd"), RegionCities: pbCities, @@ -346,7 +346,7 @@ func (this *HTTPRequestStatManager) Upload() error { } // 再次尝试 - _, err = rpcClient.ServerRPC().UploadServerHTTPRequestStat(rpcClient.Context(), &pb.UploadServerHTTPRequestStatRequest{ + _, err = rpcClient.ServerRPC.UploadServerHTTPRequestStat(rpcClient.Context(), &pb.UploadServerHTTPRequestStatRequest{ Month: timeutil.Format("Ym"), Day: timeutil.Format("Ymd"), RegionCities: pbCities, diff --git a/internal/stats/traffic_stat_manager.go b/internal/stats/traffic_stat_manager.go index 652d7f8..4f066ea 100644 --- a/internal/stats/traffic_stat_manager.go +++ b/internal/stats/traffic_stat_manager.go @@ -210,7 +210,7 @@ func (this *TrafficStatManager) Upload() error { }) } - _, err = client.ServerDailyStatRPC().UploadServerDailyStats(client.Context(), &pb.UploadServerDailyStatsRequest{ + _, err = client.ServerDailyStatRPC.UploadServerDailyStats(client.Context(), &pb.UploadServerDailyStatsRequest{ Stats: pbServerStats, DomainStats: pbDomainStats, }) diff --git a/internal/waf/action_notify.go b/internal/waf/action_notify.go index 33a5485..b6b5eb6 100644 --- a/internal/waf/action_notify.go +++ b/internal/waf/action_notify.go @@ -34,7 +34,7 @@ func init() { } for task := range notifyChan { - _, err = rpcClient.FirewallRPC().NotifyHTTPFirewallEvent(rpcClient.Context(), &pb.NotifyHTTPFirewallEventRequest{ + _, err = rpcClient.FirewallRPC.NotifyHTTPFirewallEvent(rpcClient.Context(), &pb.NotifyHTTPFirewallEventRequest{ ServerId: task.ServerId, HttpFirewallPolicyId: task.HttpFirewallPolicyId, HttpFirewallRuleGroupId: task.HttpFirewallRuleGroupId, diff --git a/internal/waf/action_record_ip.go b/internal/waf/action_record_ip.go index ecf2fe8..936f1d5 100644 --- a/internal/waf/action_record_ip.go +++ b/internal/waf/action_record_ip.go @@ -50,7 +50,7 @@ func init() { if len(reason) == 0 { reason = "触发WAF规则自动加入" } - _, err = rpcClient.IPItemRPC().CreateIPItem(rpcClient.Context(), &pb.CreateIPItemRequest{ + _, err = rpcClient.IPItemRPC.CreateIPItem(rpcClient.Context(), &pb.CreateIPItemRequest{ IpListId: task.listId, IpFrom: task.ip, IpTo: "",