优化代码

internal/rpc/services/service_node_stream.go

@@ -2,15 +2,13 @@ package services
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
-	"github.com/TeaOSLab/EdgeAPI/internal/configs"
+	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
 	"github.com/TeaOSLab/EdgeAPI/internal/errors"
 	"github.com/TeaOSLab/EdgeAPI/internal/goman"
 	"github.com/TeaOSLab/EdgeAPI/internal/remotelogs"
 	rpcutils "github.com/TeaOSLab/EdgeAPI/internal/rpc/utils"
-	"github.com/TeaOSLab/EdgeCommon/pkg/messageconfigs"
 	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
 	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
 	"github.com/iwind/TeaGo/logs"
@@ -106,24 +104,10 @@ func (this *NodeService) NodeStream(server pb.NodeService_NodeStreamServer) erro
 		}
 	}()
 
-	// 返回连接成功
-	{
-		apiConfig, err := configs.SharedAPIConfig()
-		if err != nil {
-			return err
-		}
-		connectedMessage := &messageconfigs.ConnectedAPINodeMessage{APINodeId: apiConfig.NumberId()}
-		connectedMessageJSON, err := json.Marshal(connectedMessage)
-		if err != nil {
-			return errors.Wrap(err)
-		}
-		err = server.Send(&pb.NodeStreamMessage{
-			Code:     messageconfigs.MessageCodeConnectedAPINode,
-			DataJSON: connectedMessageJSON,
-		})
-		if err != nil {
-			return err
-		}
+	// 设置API节点
+	err = models.SharedNodeDAO.UpdateNodeConnectedAPINodes(nil, nodeId, []int64{teaconst.NodeId})
+	if err != nil {
+		return err
 	}
 
 	//logs.Println("[RPC]accepted node '" + numberutils.FormatInt64(nodeId) + "' connection")