修复 trace 日志不全问题

client/client.go

@@ -19,6 +19,7 @@ type Client struct {
 	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,7 +166,7 @@ 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 := 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 {
@@ -201,9 +202,9 @@ func (client *Client) handleUpStream(conn *net.TCPConn, packet *protocol.MajoraP
 }
 
 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"
@@ -151,12 +150,3 @@ func NewSession(sessionId string) *Session {
 		Recorder: acquireRecorder(sessionId),
 	}
 }
\ No newline at end of file
-
-func GetTraceRecorderFromSession(session getty.Session) Recorder {
-	majoraSession := session.GetAttribute(MajoraSessionName)
-	traceSession, ok := majoraSession.(*Session)
-	if !ok {
-		return NewSession(sessionIdNop).Recorder
-	}
-	return traceSession.Recorder
-}