addjust code

Author: Geapefurit

api/getconnectioninfo.go

@@ -0,0 +1,19 @@
+//go:build !codeanalysis
+// +build !codeanalysis
+
+package api
+
+import (
+	"context"
+
+	"github.com/NpoolPlatform/message/npool/sphinxproxy"
+	"github.com/NpoolPlatform/sphinx-proxy/pkg/connection"
+)
+
+func (s *Server) GetConnectionInfos(ctx context.Context, in *sphinxproxy.GetConnectionInfosRequest) (*sphinxproxy.GetConnectionInfosResponse, error) {
+	infos := connection.GetConnectionMGR().GetConnectionInfos()
+	return &sphinxproxy.GetConnectionInfosResponse{
+		Infos: infos,
+		Total: uint32(len(infos)),
+	}, nil
+}

pkg/connection/connection.go

@@ -0,0 +1,111 @@
+package connection
+
+import (
+	"context"
+	"encoding/json"
+	"sync"
+	"time"
+
+	"github.com/NpoolPlatform/go-service-framework/pkg/logger"
+	"github.com/NpoolPlatform/go-service-framework/pkg/wlog"
+	"github.com/NpoolPlatform/message/npool/sphinxproxy"
+)
+
+type Connection struct {
+	sphinxproxy.SphinxProxyStream_ProxyConnnectionServer
+	*sphinxproxy.ConnectionInfo
+	ctx          context.Context
+	cancel       context.CancelFunc
+	onCloseFuncs []func(conn *Connection)
+	recvHandlers []func(data *sphinxproxy.DataElement)
+	closeOnce    sync.Once
+}
+
+func RegisterConnection(stream sphinxproxy.SphinxProxyStream_ProxyConnnectionServer) (*Connection, error) {
+	select {
+	case <-time.NewTicker(time.Second * 3).C:
+		return nil, wlog.Errorf("timeout for register connection")
+	default:
+		data, err := stream.Recv()
+		if err != nil {
+			return nil, wlog.WrapError(err)
+		}
+
+		statusCode := sphinxproxy.StatusCode_StatusCodeSuccess
+		statusMsg := ""
+
+		connInfo := &sphinxproxy.ConnectionInfo{}
+		err = json.Unmarshal(data.Payload, connInfo)
+		if err != nil {
+			statusCode = sphinxproxy.StatusCode_StatusCodeFailed
+			statusMsg = err.Error()
+		}
+
+		if statusCode == sphinxproxy.StatusCode_StatusCodeSuccess && data.ConnectID != connInfo.ID {
+			statusCode = sphinxproxy.StatusCode_StatusCodeFailed
+			statusMsg = err.Error()
+		}
+
+		err = stream.Send(&sphinxproxy.DataElement{
+			ConnectID:  data.ConnectID,
+			MsgID:      data.MsgID,
+			StatusCode: statusCode,
+			StatusMsg:  &statusMsg,
+		})
+		if err != nil {
+			return nil, wlog.WrapError(err)
+		}
+
+		if statusCode != sphinxproxy.StatusCode_StatusCodeSuccess {
+			return nil, wlog.Errorf(statusMsg)
+		}
+
+		ctx, cancel := context.WithCancel(stream.Context())
+		return &Connection{
+			SphinxProxyStream_ProxyConnnectionServer: stream,
+			ConnectionInfo:                           connInfo,
+			ctx:                                      ctx,
+			cancel:                                   cancel,
+		}, nil
+	}
+}
+
+func (conn *Connection) WatchClose(onClose func(conn *Connection)) {
+	conn.onCloseFuncs = append(conn.onCloseFuncs, onClose)
+}
+
+func (conn *Connection) WatchRecv(onRecv func(data *sphinxproxy.DataElement)) {
+	conn.recvHandlers = append(conn.recvHandlers, onRecv)
+}
+
+func (conn *Connection) Close() {
+	conn.closeOnce.Do(func() {
+		for _, onClose := range conn.onCloseFuncs {
+			onClose(conn)
+		}
+		conn.cancel()
+		logger.Sugar().Warnf(
+			"connection is closed, siderType: %v, ID: %v, Position: %v",
+			conn.SiderType,
+			conn.Position,
+			conn.Position)
+	})
+}
+
+func (conn *Connection) OnRecv() {
+	go func() {
+		defer conn.Close()
+		for {
+			data, err := conn.Recv()
+			if err != nil {
+				logger.Sugar().Error(err)
+				return
+			}
+			for _, recvHandler := range conn.recvHandlers {
+				go recvHandler(data)
+			}
+		}
+	}()
+
+	<-conn.ctx.Done()
+}