github.com/metacubex/mihomo@v1.18.5/common/net/deadline/packet_enhance.go (about)

     1  package deadline
     2  
     3  import (
     4  	"net"
     5  	"os"
     6  	"runtime"
     7  
     8  	"github.com/metacubex/mihomo/common/net/packet"
     9  )
    10  
    11  type EnhancePacketConn struct {
    12  	*NetPacketConn
    13  	enhancePacketConn
    14  }
    15  
    16  var _ packet.EnhancePacketConn = (*EnhancePacketConn)(nil)
    17  
    18  func NewEnhancePacketConn(pc packet.EnhancePacketConn) packet.EnhancePacketConn {
    19  	return NewNetPacketConn(pc).(packet.EnhancePacketConn)
    20  }
    21  
    22  type enhanceReadResult struct {
    23  	data []byte
    24  	put  func()
    25  	addr net.Addr
    26  	err  error
    27  }
    28  
    29  type enhancePacketConn struct {
    30  	netPacketConn     *NetPacketConn
    31  	enhancePacketConn packet.EnhancePacketConn
    32  }
    33  
    34  func (c *enhancePacketConn) WaitReadFrom() (data []byte, put func(), addr net.Addr, err error) {
    35  FOR:
    36  	for {
    37  		select {
    38  		case result := <-c.netPacketConn.resultCh:
    39  			if result != nil {
    40  				if result, ok := result.(*enhanceReadResult); ok {
    41  					data = result.data
    42  					put = result.put
    43  					addr = result.addr
    44  					err = result.err
    45  					c.netPacketConn.resultCh <- nil // finish cache read
    46  					return
    47  				}
    48  				c.netPacketConn.resultCh <- result // another type of read
    49  				runtime.Gosched()                  // allowing other goroutines to run
    50  				continue FOR
    51  			} else {
    52  				c.netPacketConn.resultCh <- nil
    53  				break FOR
    54  			}
    55  		case <-c.netPacketConn.pipeDeadline.wait():
    56  			return nil, nil, nil, os.ErrDeadlineExceeded
    57  		}
    58  	}
    59  
    60  	if c.netPacketConn.disablePipe.Load() {
    61  		return c.enhancePacketConn.WaitReadFrom()
    62  	} else if c.netPacketConn.deadline.Load().IsZero() {
    63  		c.netPacketConn.inRead.Store(true)
    64  		defer c.netPacketConn.inRead.Store(false)
    65  		data, put, addr, err = c.enhancePacketConn.WaitReadFrom()
    66  		return
    67  	}
    68  
    69  	<-c.netPacketConn.resultCh
    70  	go c.pipeWaitReadFrom()
    71  
    72  	return c.WaitReadFrom()
    73  }
    74  
    75  func (c *enhancePacketConn) pipeWaitReadFrom() {
    76  	data, put, addr, err := c.enhancePacketConn.WaitReadFrom()
    77  	result := &enhanceReadResult{}
    78  	result.data = data
    79  	result.put = put
    80  	result.addr = addr
    81  	result.err = err
    82  	c.netPacketConn.resultCh <- result
    83  }