github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/server/event_drive.go (about)

     1  package server
     2  
     3  import (
     4  	"context"
     5  	lContext "github.com/nyan233/littlerpc/core/common/context"
     6  	"github.com/nyan233/littlerpc/core/common/errorhandler"
     7  	msgparser2 "github.com/nyan233/littlerpc/core/common/msgparser"
     8  	"github.com/nyan233/littlerpc/core/common/transport"
     9  	"github.com/nyan233/littlerpc/core/middle/plugin"
    10  	lerror "github.com/nyan233/littlerpc/core/protocol/error"
    11  	"github.com/nyan233/littlerpc/core/protocol/message"
    12  	"github.com/nyan233/littlerpc/core/utils/convert"
    13  	"math"
    14  	"runtime"
    15  	"time"
    16  )
    17  
    18  func (s *Server) onClose(conn transport.ConnAdapter, err error) {
    19  	if err != nil {
    20  		s.logger.Warn("LRPC: Close Connection: %s -> %s err: %v", conn.RemoteAddr(), conn.LocalAddr(), err)
    21  	} else {
    22  		s.logger.Info("LRPC: Close Connection: %s -> %s", conn.RemoteAddr(), conn.LocalAddr())
    23  	}
    24  	if !s.pManager.Event4S(plugin.OnClose) {
    25  		s.logger.Warn("LRPC: plugin entry interrupted onClose")
    26  	}
    27  	// 关闭之前的清理工作
    28  	desc, ok := conn.Source().(*connSourceDesc)
    29  	if !ok {
    30  		s.logger.Warn("LRPC: onClose connection source not found")
    31  		return
    32  	}
    33  	desc.ctxManager.CancelAll()
    34  }
    35  
    36  func (s *Server) onRead(c transport.ConnAdapter) {
    37  	s.parseMessageAndHandle(c, nil, false)
    38  }
    39  
    40  func (s *Server) onMessage(c transport.ConnAdapter, data []byte) {
    41  	s.parseMessageAndHandle(c, data, true)
    42  }
    43  
    44  func (s *Server) parseMessageAndHandle(c transport.ConnAdapter, data []byte, prepared bool) {
    45  	if prepared && !s.pManager.Event4S(plugin.OnMessage) {
    46  		s.logger.Info("LRPC: plugin interrupted onMessage")
    47  		return
    48  	} else if !prepared && !s.pManager.Event4S(plugin.OnRead) {
    49  		s.logger.Info("LRPC: plugin interrupted onRead")
    50  		return
    51  	}
    52  	desc, ok := c.Source().(*connSourceDesc)
    53  	if !ok {
    54  		s.logger.Error("LRPC: no register message-parser, remote ip = %s", c.RemoteAddr())
    55  		_ = c.Close()
    56  		return
    57  	}
    58  	// 2023/02/22 : 删除Debug Message相关的代码
    59  	defer s.eventLoopTopRecover(c, desc)
    60  	var traitMsgs []msgparser2.ParserMessage
    61  	var err error
    62  	if prepared {
    63  		traitMsgs, err = desc.Parser.Parse(data)
    64  	} else {
    65  		traitMsgs, err = desc.Parser.ParseOnReader(msgparser2.DefaultReader(c))
    66  	}
    67  	if err != nil {
    68  		// 错误处理过程会在严重错误时关闭连接, 所以msgId == math.MaxUint64也没有关系
    69  		// 设为0有可能和客户端生成的MessageId冲突
    70  		// 在解码消息失败时也不可能拿到正确的msgId
    71  		s.handleError(c, desc.Writer, math.MaxUint64, s.eHandle.LWarpErrorDesc(errorhandler.ErrMessageDecoding, err.Error()))
    72  		s.logger.Warn("LRPC: parse failed %v", err)
    73  		return
    74  	}
    75  	for _, traitMsg := range traitMsgs {
    76  		// init message option
    77  		msgOpt := newConnDesc(s, traitMsg, c, desc)
    78  		msgOpt.SelectCodecAndEncoder()
    79  		msgOpt.setFreeFunc(func(msg *message.Message) {
    80  			desc.Parser.Free(msg)
    81  		})
    82  		switch traitMsg.Message.GetMsgType() {
    83  		case message.Ping:
    84  			s.messageKeepAlive(msgOpt)
    85  		case message.ContextCancel:
    86  			s.messageContextCancel(msgOpt)
    87  		case message.Call:
    88  			s.messageCall(msgOpt, desc)
    89  		default:
    90  			// 释放消息, 这一步所有分支内都不可少
    91  			msgOpt.Free()
    92  			msgOpt.FreePluginCtx()
    93  			s.handleError(c, msgOpt.Desc.Writer, traitMsg.Message.GetMsgId(), lerror.LWarpStdError(errorhandler.ErrServer,
    94  				"Unknown Message Call Type", traitMsg.Message.GetMsgType()))
    95  			continue
    96  		}
    97  	}
    98  }
    99  
   100  func (s *Server) onOpen(conn transport.ConnAdapter) {
   101  	if !s.pManager.Event4S(plugin.OnOpen) {
   102  		_ = conn.Close()
   103  		s.logger.Info("LRPC: plugin interrupted onOpen")
   104  		return
   105  	}
   106  	// 初始化连接的相关数据
   107  	cfg := s.config.Load()
   108  	desc := &connSourceDesc{}
   109  	desc.Parser = cfg.ParserFactory(
   110  		&msgparser2.SimpleAllocTor{SharedPool: s.pool.TakeMessagePool()},
   111  		uint32(s.config.Load().ReadBufferSize),
   112  	)
   113  	desc.Writer = cfg.WriterFactory()
   114  	desc.ctxManager = newContextManager()
   115  	desc.remoteAddr = conn.RemoteAddr()
   116  	desc.localAddr = conn.LocalAddr()
   117  	desc.cacheCtx = lContext.WithLocalAddr(context.Background(), desc.localAddr)
   118  	desc.cacheCtx = lContext.WithRemoteAddr(context.Background(), desc.remoteAddr)
   119  	conn.SetSource(desc)
   120  	// init keepalive
   121  	if cfg.KeepAlive {
   122  		if err := conn.SetDeadline(time.Now().Add(cfg.KeepAliveTimeout)); err != nil {
   123  			s.logger.Error("LRPC: keepalive set deadline failed: %v", err)
   124  			_ = conn.Close()
   125  		}
   126  	}
   127  }
   128  
   129  func (s *Server) eventLoopTopRecover(c transport.ConnAdapter, desc *connSourceDesc) {
   130  	e := recover()
   131  	if e == nil {
   132  		return
   133  	}
   134  	switch e.(type) {
   135  	case lerror.LErrorDesc:
   136  		s.handleError(c, desc.Writer, math.MaxUint64, e.(lerror.LErrorDesc))
   137  	case error:
   138  		s.handleError(c, desc.Writer, math.MaxUint64,
   139  			s.eHandle.LNewErrorDesc(lerror.UnsafeOption, "runtime error", e.(error).Error()))
   140  	default:
   141  		s.handleError(c, desc.Writer, math.MaxUint64,
   142  			s.eHandle.LWarpErrorDesc(errorhandler.ErrServer, "unknown error"))
   143  	}
   144  	var buf [4096]byte
   145  	length := runtime.Stack(buf[:], false)
   146  	s.logger.Panic("LRPC: eventLoopTopRecover panic : %v\n%s", e, convert.BytesToString(buf[:length]))
   147  }