修复 trace 日志不全问题

client/client.go

@@ -16,9 +16,10 @@ import (
 type Client struct {
 	config *model.Configure
 
-	localAddr net.Addr
-	natTunnel getty.Client
-	connStore sync.Map
+	localAddr    net.Addr
+	natTunnel    getty.Client
+	connStore    sync.Map
+	sessionStore sync.Map
 }
 
 func NewClientWithConf(cfg *model.Configure) *Client {

client/event.go

@@ -39,7 +39,7 @@ func (client *Client) handleConnect(packet *protocol.MajoraPacket, session getty
 		log.Error().Errorf("SN:%d -> get sessionId failed", packet.SerialNumber)
 	}
 	traceSession := trace.NewSession(sessionId)
-	session.SetAttribute(trace.MajoraSessionName, traceSession)
+	client.AddSession(packet, traceSession)
 	traceSession.Recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("SN:%d -> start handle connect", packet.SerialNumber))
 
 	if session.IsClosed() {
@@ -128,7 +128,8 @@ func decodeMap(data []byte) map[string]string {
 }
 
 func (client *Client) handleTransfer(packet *protocol.MajoraPacket, session getty.Session) {
-	traceRecorder := trace.GetTraceRecorderFromSession(session)
+
+	traceRecorder := client.GetRecorderFromSession(packet.SerialNumber)
 	traceRecorder.RecordEvent(trace.TransferEvent,
 		fmt.Sprintf("SN:%d - > receive transfer packet from natServer,start to be forward to target", packet.SerialNumber))
 
@@ -165,45 +166,45 @@ func (client *Client) handleTransfer(packet *protocol.MajoraPacket, session gett
 }
 
 func (client *Client) handleUpStream(conn *net.TCPConn, packet *protocol.MajoraPacket, session getty.Session) {
-		traceRecorder := trace.GetTraceRecorderFromSession(session)
-		traceRecorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("SN:%d -> handleUpStream start", packet.SerialNumber))
-		log.Run().Debugf("[handleUpStream] %d-> handleUpStream start...", packet.SerialNumber)
-		for {
-			buf := make([]byte, common.BufSize) // 4k
-			cnt, err := conn.Read(buf)
-			if err != nil {
-				log.Run().Debugf("[handleUpStream] %d->read with error:%+v,l:%s->r:%s",
-					packet.SerialNumber, err, conn.LocalAddr(), conn.RemoteAddr())
-				recorderMessage := fmt.Sprintf("SN: %d -> read with l:%s->r:%s",
-					packet.SerialNumber, conn.LocalAddr(), conn.RemoteAddr())
-				traceRecorder.RecordErrorEvent(trace.UpStreamEvent, recorderMessage, err)
-				client.OnClose(session, conn, packet.SerialNumber)
-				break
-			}
-			traceRecorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("SN: %d read count: %d",
-				packet.SerialNumber, cnt))
-			pack := protocol.TypeTransfer.CreatePacket()
-			pack.Data = buf[:cnt]
-			pack.SerialNumber = packet.SerialNumber
-
-			if _, _, err := session.WritePkg(pack, 0); err != nil {
-				log.Error().Errorf("[handleUpStream] %d-> write to server fail %+v", packet.SerialNumber, err.Error())
-				traceRecorder.RecordErrorEvent(trace.UpStreamEvent,
-					fmt.Sprintf("SN: %d -> write to natServer failed", packet.SerialNumber), err)
-				client.OnClose(session, conn, packet.SerialNumber)
-				break
-			} else {
-				log.Run().Debugf("[handleUpStream] %d->success dataLen:%d", packet.SerialNumber, len(packet.Data))
-				traceRecorder.RecordEvent(trace.UpStreamEvent,
-					fmt.Sprintf("SN: %d write to natServer success", packet.SerialNumber))
-			}
+	traceRecorder := client.GetRecorderFromSession(packet.SerialNumber)
+	traceRecorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("SN:%d -> handleUpStream start", packet.SerialNumber))
+	log.Run().Debugf("[handleUpStream] %d-> handleUpStream start...", packet.SerialNumber)
+	for {
+		buf := make([]byte, common.BufSize) // 4k
+		cnt, err := conn.Read(buf)
+		if err != nil {
+			log.Run().Debugf("[handleUpStream] %d->read with error:%+v,l:%s->r:%s",
+				packet.SerialNumber, err, conn.LocalAddr(), conn.RemoteAddr())
+			recorderMessage := fmt.Sprintf("SN: %d -> read with l:%s->r:%s",
+				packet.SerialNumber, conn.LocalAddr(), conn.RemoteAddr())
+			traceRecorder.RecordErrorEvent(trace.UpStreamEvent, recorderMessage, err)
+			client.OnClose(session, conn, packet.SerialNumber)
+			break
+		}
+		traceRecorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("SN: %d read count: %d",
+			packet.SerialNumber, cnt))
+		pack := protocol.TypeTransfer.CreatePacket()
+		pack.Data = buf[:cnt]
+		pack.SerialNumber = packet.SerialNumber
+
+		if _, _, err := session.WritePkg(pack, 0); err != nil {
+			log.Error().Errorf("[handleUpStream] %d-> write to server fail %+v", packet.SerialNumber, err.Error())
+			traceRecorder.RecordErrorEvent(trace.UpStreamEvent,
+				fmt.Sprintf("SN: %d -> write to natServer failed", packet.SerialNumber), err)
+			client.OnClose(session, conn, packet.SerialNumber)
+			break
+		} else {
+			log.Run().Debugf("[handleUpStream] %d->success dataLen:%d", packet.SerialNumber, len(packet.Data))
+			traceRecorder.RecordEvent(trace.UpStreamEvent,
+				fmt.Sprintf("SN: %d write to natServer success", packet.SerialNumber))
 		}
+	}
 }
 
 func (client *Client) handleDisconnectMessage(session getty.Session, packet *protocol.MajoraPacket) {
-	traceRecorder := trace.GetTraceRecorderFromSession(session)
-	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("SN: %d start close session", packet.SerialNumber))
-	log.Run().Debugf("[handleDisconnectMessage] %d->session closed %v", packet.SerialNumber, session.IsClosed())
+	traceRecorder := client.GetRecorderFromSession(packet.SerialNumber)
+	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("SN: %d start close session extra:%s", packet.SerialNumber, packet.Extra))
+	log.Run().Debugf("[handleDisconnectMessage] %d->session closed %v extra:%s", packet.SerialNumber, session.IsClosed())
 	if conn, ok := client.connStore.Load(packet.SerialNumber); ok {
 		client.OnClose(session, conn.(net.Conn), packet.SerialNumber)
 	}
@@ -218,6 +219,25 @@ func (client *Client) handleDestroyMessage() {
 	client.natTunnel.Close()
 }
 
+func (client *Client) AddSession(packet *protocol.MajoraPacket, session *trace.Session) {
+	if _, ok := client.sessionStore.Load(packet.SerialNumber); ok {
+		log.Error().Errorf("[AddSession] %d->error, has one", packet.SerialNumber)
+	}
+	client.sessionStore.Store(packet.SerialNumber, session)
+	log.Run().Debugf("[AddSession] %d-> success", packet.SerialNumber)
+}
+
+func (client *Client) GetRecorderFromSession(sn int64) trace.Recorder {
+	session, ok := client.sessionStore.Load(sn)
+	if !ok {
+		log.Error().Errorf("[GetSession] get session failed, has one", sn)
+		session = trace.NewSession("")
+	}
+	traceSession := session.(*trace.Session)
+	return traceSession.Recorder
+
+}
+
 func (client *Client) AddConnection(packet *protocol.MajoraPacket, conn *net.TCPConn, addr string) {
 	if _, ok := client.connStore.Load(packet.SerialNumber); ok {
 		log.Error().Errorf("[AddConnection] %d->error, has one", packet.SerialNumber)
@@ -233,8 +253,9 @@ func (client *Client) OnClose(natSession getty.Session, upStreamSession net.Conn
 			log.Error().Errorf("OnClose %+v", err)
 		}
 	}()
-	traceRecorder := trace.GetTraceRecorderFromSession(natSession)
+	traceRecorder := client.GetRecorderFromSession(serialNumber)
 	client.connStore.Delete(serialNumber)
+	client.sessionStore.Delete(serialNumber)
 	_ = upStreamSession.Close()
 	client.closeVirtualConnection(natSession, serialNumber)
 	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("SN: %d close success", serialNumber))
@@ -242,7 +263,7 @@ func (client *Client) OnClose(natSession getty.Session, upStreamSession net.Conn
 
 //closeVirtualConnection disconnect to server
 func (client *Client) closeVirtualConnection(session getty.Session, serialNumber int64) {
-	traceRecorder := trace.GetTraceRecorderFromSession(session)
+	traceRecorder := client.GetRecorderFromSession(serialNumber)
 
 	log.Run().Debugf("[closeVirtualConnection] %d->session closed %v", serialNumber, session.IsClosed())
 	if session.IsClosed() {

log/logger.go

@@ -28,7 +28,7 @@ const (
 func getLogWriter(path string) zapcore.WriteSyncer {
 	lumberJackLogger := &lumberjack.Logger{
 		Filename:   path,
-		MaxSize:    10,
+		MaxSize:    100,
 		MaxBackups: 5,
 		MaxAge:     30,
 		Compress:   false,
@@ -40,7 +40,7 @@ func getLogWriter(path string) zapcore.WriteSyncer {
 }
 
 func getEncoder() zapcore.Encoder {
-	encoderConfig := zap.NewProductionEncoderConfig()
+	encoderConfig := zap.NewDevelopmentEncoderConfig()
 	encoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
 	encoderConfig.EncodeLevel = zapcore.CapitalLevelEncoder
 	return zapcore.NewConsoleEncoder(encoderConfig)

trace/trace.go

@@ -4,7 +4,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/adamweixuan/getty"
 	"go.uber.org/zap"
 
 	"virjar.com/majora-go/env"
@@ -13,7 +12,7 @@ import (
 )
 
 var (
-	sessionEventChan  = make(chan *sessionEvent, 100)
+	sessionEventChan  = make(chan *sessionEvent, 1000)
 	ConnectEvent      = "ConnectEvent"
 	TransferEvent     = "TransferEvent"
 	MajoraSessionName = "MajoraSessionId"
@@ -150,13 +149,4 @@ func NewSession(sessionId string) *Session {
 	return &Session{
 		Recorder: acquireRecorder(sessionId),
 	}
-}
-
-func GetTraceRecorderFromSession(session getty.Session) Recorder {
-	majoraSession := session.GetAttribute(MajoraSessionName)
-	traceSession, ok := majoraSession.(*Session)
-	if !ok {
-		return NewSession(sessionIdNop).Recorder
-	}
-	return traceSession.Recorder
-}
+}
\ No newline at end of file