trpc.group/trpc-go/trpc-go@v1.0.3/transport/server_transport_udp.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  package transport
    15  
    16  import (
    17  	"context"
    18  	"errors"
    19  	"math"
    20  	"net"
    21  	"time"
    22  
    23  	"github.com/panjf2000/ants/v2"
    24  	"trpc.group/trpc-go/trpc-go/codec"
    25  	"trpc.group/trpc-go/trpc-go/errs"
    26  	"trpc.group/trpc-go/trpc-go/internal/packetbuffer"
    27  	"trpc.group/trpc-go/trpc-go/internal/report"
    28  	"trpc.group/trpc-go/trpc-go/log"
    29  )
    30  
    31  func (s *serverTransport) serveUDP(ctx context.Context, rwc *net.UDPConn, pool *ants.PoolWithFunc,
    32  	opts *ListenServeOptions) error {
    33  
    34  	// Sets the size of the operating system's receive buffer associated with the connection.
    35  	if s.opts.RecvUDPRawSocketBufSize > 0 {
    36  		rwc.SetReadBuffer(s.opts.RecvUDPRawSocketBufSize)
    37  	}
    38  
    39  	var tempDelay time.Duration
    40  	buf := packetbuffer.New(rwc, s.opts.RecvUDPPacketBufferSize)
    41  	defer buf.Close()
    42  	fr := opts.FramerBuilder.New(buf)
    43  	copyFrame := !codec.IsSafeFramer(fr)
    44  
    45  	for {
    46  		select {
    47  		case <-ctx.Done():
    48  			return errors.New("recv server close event")
    49  		default:
    50  		}
    51  
    52  		req, err := fr.ReadFrame()
    53  		if err != nil {
    54  			if ne, ok := err.(net.Error); ok && ne.Temporary() {
    55  				if tempDelay == 0 {
    56  					tempDelay = 5 * time.Millisecond
    57  				} else {
    58  					tempDelay *= 2
    59  				}
    60  				if max := 1 * time.Second; tempDelay > max {
    61  					tempDelay = max
    62  				}
    63  				time.Sleep(tempDelay)
    64  				continue
    65  			}
    66  			report.UDPServerTransportReadFail.Incr()
    67  			log.Trace("transport: udpconn serve ReadFrame fail ", err)
    68  			buf.Next()
    69  			continue
    70  		}
    71  		tempDelay = 0
    72  
    73  		report.UDPServerTransportReceiveSize.Set(float64(len(req)))
    74  
    75  		remoteAddr, ok := buf.CurrentPacketAddr().(*net.UDPAddr)
    76  		if !ok {
    77  			log.Trace("transport: udpconn serve address is not udp address")
    78  			buf.Next()
    79  			continue
    80  		}
    81  
    82  		// One packet of udp corresponds to one trpc packet,
    83  		// and after parsing, there should not be any remaining data.
    84  		if err := buf.Next(); err != nil {
    85  			report.UDPServerTransportUnRead.Incr()
    86  			log.Trace("transport: udpconn serve ReadFrame data remaining bytes data, ", err)
    87  			continue
    88  		}
    89  
    90  		c := &udpconn{
    91  			conn:       s.newConn(ctx, opts),
    92  			rwc:        rwc,
    93  			remoteAddr: remoteAddr,
    94  		}
    95  
    96  		if copyFrame {
    97  			c.req = make([]byte, len(req))
    98  			copy(c.req, req)
    99  		} else {
   100  			c.req = req
   101  		}
   102  
   103  		if pool == nil {
   104  			go c.serve()
   105  			continue
   106  		}
   107  		if err := pool.Invoke(c); err != nil {
   108  			report.UDPServerTransportJobQueueFullFail.Incr()
   109  			log.Trace("transport: udpconn serve routine pool put job queue fail ", err)
   110  			go c.serve()
   111  		}
   112  	}
   113  }
   114  
   115  // udpconn is the UDP connection which is established when server receives a client connecting
   116  // request.
   117  type udpconn struct {
   118  	*conn
   119  	req        []byte
   120  	rwc        *net.UDPConn
   121  	remoteAddr *net.UDPAddr
   122  }
   123  
   124  func (c *udpconn) serve() {
   125  	// Generate a new empty message binding to the ctx.
   126  	ctx, msg := codec.WithNewMessage(context.Background())
   127  	defer codec.PutBackMessage(msg)
   128  
   129  	// Set local address and remote address to message.
   130  	msg.WithLocalAddr(c.rwc.LocalAddr())
   131  	msg.WithRemoteAddr(c.remoteAddr)
   132  
   133  	rsp, err := c.handle(ctx, c.req)
   134  	if err != nil {
   135  		if err != errs.ErrServerNoResponse {
   136  			report.UDPServerTransportHandleFail.Incr()
   137  			log.Tracef("udp handle fail:%v", err)
   138  		}
   139  		return
   140  	}
   141  
   142  	report.UDPServerTransportSendSize.Set(float64(len(rsp)))
   143  	if _, err := c.rwc.WriteToUDP(rsp, c.remoteAddr); err != nil {
   144  		report.UDPServerTransportWriteFail.Incr()
   145  		log.Tracef("udp write out fail:%v", err)
   146  		return
   147  	}
   148  }
   149  
   150  func createUDPRoutinePool(size int) *ants.PoolWithFunc {
   151  	if size <= 0 {
   152  		size = math.MaxInt32
   153  	}
   154  	pool, err := ants.NewPoolWithFunc(size, func(args interface{}) {
   155  		c, ok := args.(*udpconn)
   156  		if !ok {
   157  			log.Tracef("routine pool args type error, shouldn't happen!")
   158  			return
   159  		}
   160  		c.serve()
   161  	})
   162  	if err != nil {
   163  		log.Tracef("routine pool create error:%v", err)
   164  		return nil
   165  	}
   166  	return pool
   167  }