github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/common/transport/nbio_tcp.go (about)

     1  package transport
     2  
     3  import (
     4  	"errors"
     5  	"github.com/lesismal/llib/std/crypto/tls"
     6  	"github.com/lesismal/nbio"
     7  	ntls "github.com/lesismal/nbio/extension/tls"
     8  	"net"
     9  	"runtime"
    10  	"sync/atomic"
    11  	"unsafe"
    12  )
    13  
    14  type NBioTcpEngine struct {
    15  	started int32
    16  	closed  int32
    17  	tlsC    *tls.Config
    18  	server  *nbio.Engine
    19  	onRead  func(conn ConnAdapter)
    20  	onMsg   func(conn ConnAdapter, bytes []byte)
    21  	onClose func(conn ConnAdapter, err error)
    22  	onOpen  func(conn ConnAdapter)
    23  }
    24  
    25  func NewNBioTcpClient() ClientBuilder {
    26  	return &NBioTcpEngine{
    27  		server: nbio.NewEngine(nbio.Config{
    28  			Name:    "LittleRpc-TCP-Client",
    29  			NPoller: runtime.NumCPU(),
    30  		}),
    31  		onOpen: func(conn ConnAdapter) {
    32  			return
    33  		},
    34  		onMsg: func(conn ConnAdapter, bytes []byte) {
    35  			return
    36  		},
    37  		onClose: func(conn ConnAdapter, err error) {
    38  			return
    39  		},
    40  	}
    41  }
    42  
    43  func NewNBioTcpServer(config NetworkServerConfig) ServerBuilder {
    44  	nConfig := nbio.Config{}
    45  	nConfig.Name = "LittleRpc-Server-Tcp"
    46  	nConfig.Network = "tcp"
    47  	nConfig.ReadBufferSize = ReadBufferSize
    48  	nConfig.MaxWriteBufferSize = MaxWriteBufferSize
    49  	nConfig.Addrs = config.Addrs
    50  	eng := nbio.NewEngine(nConfig)
    51  	server := &NBioTcpEngine{}
    52  	server.server = eng
    53  	// set default function
    54  	server.onMsg = func(conn ConnAdapter, bytes []byte) {
    55  		return
    56  	}
    57  	server.onOpen = func(conn ConnAdapter) {
    58  		return
    59  	}
    60  	server.onClose = func(conn ConnAdapter, err error) {
    61  		return
    62  	}
    63  	return server
    64  }
    65  
    66  func (engine *NBioTcpEngine) NewConn(config NetworkClientConfig) (ConnAdapter, error) {
    67  	netConn, err := net.Dial("tcp", config.ServerAddr)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  	convConn, err := engine.server.AddConn(netConn)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  	return (*nTcpConn)(unsafe.Pointer(convConn)), nil
    76  }
    77  
    78  func (engine *NBioTcpEngine) EventDriveInter() EventDriveInter {
    79  	return engine
    80  }
    81  
    82  func (engine *NBioTcpEngine) Client() ClientEngine {
    83  	return engine
    84  }
    85  
    86  func (engine *NBioTcpEngine) Server() ServerEngine {
    87  	return engine
    88  }
    89  
    90  func (engine *NBioTcpEngine) OnRead(f func(conn ConnAdapter)) {
    91  	engine.onRead = f
    92  }
    93  
    94  func (engine *NBioTcpEngine) OnMessage(f func(conn ConnAdapter, data []byte)) {
    95  	engine.onMsg = f
    96  }
    97  
    98  func (engine *NBioTcpEngine) OnOpen(f func(conn ConnAdapter)) {
    99  	engine.onOpen = f
   100  }
   101  
   102  func (engine *NBioTcpEngine) OnClose(f func(conn ConnAdapter, err error)) {
   103  	engine.onClose = f
   104  }
   105  
   106  func (engine *NBioTcpEngine) Start() error {
   107  	if !atomic.CompareAndSwapInt32(&engine.started, 0, 1) {
   108  		return errors.New("wsEngine already started")
   109  	}
   110  	server := engine.server
   111  	engine.bind()
   112  	return server.Start()
   113  }
   114  
   115  func (engine *NBioTcpEngine) bind() {
   116  	server := engine.server
   117  	if engine.tlsC == nil {
   118  		server.OnOpen(func(c *nbio.Conn) {
   119  			engine.onOpen((*nTcpConn)(unsafe.Pointer(c)))
   120  		})
   121  		if engine.onRead == nil {
   122  			server.OnData(func(c *nbio.Conn, data []byte) {
   123  				engine.onMsg((*nTcpConn)(unsafe.Pointer(c)), data)
   124  			})
   125  		} else {
   126  			server.OnRead(func(c *nbio.Conn) {
   127  				engine.onRead((*nTcpConn)(unsafe.Pointer(c)))
   128  			})
   129  		}
   130  		server.OnClose(func(c *nbio.Conn, err error) {
   131  			engine.onClose((*nTcpConn)(unsafe.Pointer(c)), err)
   132  		})
   133  	} else {
   134  		engine.tlsC.BuildNameToCertificate()
   135  		server.OnClose(ntls.WrapClose(func(c *nbio.Conn, tlsConn *ntls.Conn, err error) {
   136  			engine.onClose(nil, err)
   137  		}))
   138  		server.OnOpen(ntls.WrapOpen(engine.tlsC, false,
   139  			func(c *nbio.Conn, tlsConn *ntls.Conn) {
   140  				engine.onOpen(nil)
   141  			}),
   142  		)
   143  		server.OnData(ntls.WrapData(func(c *nbio.Conn, tlsConn *ntls.Conn, data []byte) {
   144  			engine.onMsg(nil, data)
   145  		}))
   146  	}
   147  }
   148  
   149  func (engine *NBioTcpEngine) Stop() error {
   150  	if !atomic.CompareAndSwapInt32(&engine.closed, 0, 1) {
   151  		return errors.New("wsEngine already closed")
   152  	}
   153  	engine.server.Stop()
   154  	return nil
   155  }
   156  
   157  type nTcpConn struct {
   158  	nbio.Conn
   159  }
   160  
   161  func (n *nTcpConn) SetSource(s interface{}) {
   162  	n.SetSession(s)
   163  }
   164  
   165  func (n *nTcpConn) Source() interface{} {
   166  	return n.Session()
   167  }