trpc.group/trpc-go/trpc-go@v1.0.3/transport/tnet/server_transport.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  //go:build linux || freebsd || dragonfly || darwin
    15  // +build linux freebsd dragonfly darwin
    16  
    17  // Package tnet provides tRPC-Go transport implementation for tnet networking framework.
    18  package tnet
    19  
    20  import (
    21  	"context"
    22  	"errors"
    23  	"fmt"
    24  	"strings"
    25  	"sync"
    26  
    27  	"trpc.group/trpc-go/tnet"
    28  
    29  	"trpc.group/trpc-go/trpc-go/codec"
    30  	"trpc.group/trpc-go/trpc-go/errs"
    31  	"trpc.group/trpc-go/trpc-go/internal/addrutil"
    32  	"trpc.group/trpc-go/trpc-go/log"
    33  	"trpc.group/trpc-go/trpc-go/transport"
    34  )
    35  
    36  const transportName = "tnet"
    37  
    38  func init() {
    39  	transport.RegisterServerTransport(transportName, DefaultServerTransport)
    40  }
    41  
    42  type serverTransport struct {
    43  	addrToConn map[string]*tcpConn
    44  	m          sync.RWMutex
    45  	opts       *ServerTransportOptions
    46  }
    47  
    48  // DefaultServerTransport is the default implementation of tnet server transport.
    49  var DefaultServerTransport = NewServerTransport(WithReusePort(true))
    50  
    51  // NewServerTransport creates tnet server transport.
    52  func NewServerTransport(opts ...ServerTransportOption) transport.ServerTransport {
    53  	option := &ServerTransportOptions{}
    54  	for _, o := range opts {
    55  		o(option)
    56  	}
    57  	return &serverTransport{addrToConn: make(map[string]*tcpConn), opts: option}
    58  }
    59  
    60  // ListenAndServe begins listen and serve.
    61  func (s *serverTransport) ListenAndServe(ctx context.Context, opts ...transport.ListenServeOption) error {
    62  	lsOpts, err := buildListenServeOptions(opts...)
    63  	if err != nil {
    64  		return err
    65  	}
    66  	log.Infof("service:%s is using tnet transport, current number of pollers: %d",
    67  		lsOpts.ServiceName, tnet.NumPollers())
    68  	networks := strings.Split(lsOpts.Network, ",")
    69  	for _, network := range networks {
    70  		lsOpts.Network = network
    71  		if err := s.switchNetworkToServe(ctx, lsOpts); err != nil {
    72  			log.Error("switch to gonet default transport, ", err)
    73  			opts = append(opts, transport.WithListenNetwork(network))
    74  			return transport.DefaultServerTransport.ListenAndServe(ctx, opts...)
    75  		}
    76  	}
    77  	return nil
    78  }
    79  
    80  // Send implements ServerStreamTransport, sends stream messages.
    81  func (s *serverTransport) Send(ctx context.Context, req []byte) error {
    82  	msg := codec.Message(ctx)
    83  	raddr := msg.RemoteAddr()
    84  	laddr := msg.LocalAddr()
    85  	if raddr == nil || laddr == nil {
    86  		return errs.NewFrameError(errs.RetServerSystemErr,
    87  			fmt.Sprintf("Address is invalid, local: %s, remote: %s", laddr, raddr))
    88  	}
    89  	tc, ok := s.loadConn(addrutil.AddrToKey(laddr, raddr))
    90  	if !ok {
    91  		return errs.NewFrameError(errs.RetServerSystemErr, "can't find conn by addr")
    92  	}
    93  	if _, err := tc.rawConn.Write(req); err != nil {
    94  		tc.close()
    95  		s.Close(ctx)
    96  		return err
    97  	}
    98  	return nil
    99  }
   100  
   101  // Close closes transport, and cleans up cached connections.
   102  func (s *serverTransport) Close(ctx context.Context) {
   103  	msg := codec.Message(ctx)
   104  	raddr := msg.RemoteAddr()
   105  	laddr := msg.LocalAddr()
   106  	s.deleteConn(addrutil.AddrToKey(laddr, raddr))
   107  }
   108  
   109  func (s *serverTransport) switchNetworkToServe(ctx context.Context, opts *transport.ListenServeOptions) error {
   110  	switch opts.Network {
   111  	case "tcp", "tcp4", "tcp6":
   112  		if err := s.listenAndServeTCP(ctx, opts); err != nil {
   113  			return err
   114  		}
   115  	default:
   116  		return fmt.Errorf("tnet server transport doesn't support network type [%s]", opts.Network)
   117  	}
   118  	return nil
   119  }
   120  
   121  func (s *serverTransport) deleteConn(addr string) {
   122  	s.m.Lock()
   123  	defer s.m.Unlock()
   124  	delete(s.addrToConn, addr)
   125  }
   126  
   127  func (s *serverTransport) loadConn(addr string) (*tcpConn, bool) {
   128  	s.m.RLock()
   129  	defer s.m.RUnlock()
   130  	tc, ok := s.addrToConn[addr]
   131  	return tc, ok
   132  }
   133  
   134  func (s *serverTransport) storeConn(addr string, tc *tcpConn) {
   135  	s.m.Lock()
   136  	defer s.m.Unlock()
   137  	s.addrToConn[addr] = tc
   138  }
   139  
   140  func buildListenServeOptions(opts ...transport.ListenServeOption) (*transport.ListenServeOptions, error) {
   141  	lsOpts := &transport.ListenServeOptions{}
   142  	for _, o := range opts {
   143  		o(lsOpts)
   144  	}
   145  	if lsOpts.FramerBuilder == nil {
   146  		return nil, errors.New("transport FramerBuilder empty")
   147  	}
   148  	return lsOpts, nil
   149  }