github.com/haraldrudell/parl@v0.4.176/pnet/udp.go (about)

     1  /*
     2  © 2021–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
     3  ISC License
     4  */
     5  
     6  package pnet
     7  
     8  import (
     9  	"net"
    10  	"sync"
    11  	"sync/atomic"
    12  
    13  	"github.com/haraldrudell/parl"
    14  	"github.com/haraldrudell/parl/perrors"
    15  )
    16  
    17  type UDP struct {
    18  	Network        string
    19  	F              UDPFunc
    20  	MaxSize        int
    21  	net.UDPAddr    // struct IP Port Zone
    22  	ListenInvoked  atomic.Bool
    23  	StartingListen sync.WaitGroup
    24  	ErrCh          chan<- error
    25  	IsListening    atomic.Bool
    26  	NetUDPConn     *net.UDPConn
    27  	connMutex      sync.RWMutex
    28  	Addr           net.Addr
    29  	IsShutdown     atomic.Bool
    30  }
    31  
    32  type UDPFunc func(b []byte, oob []byte, flags int, addr *net.UDPAddr)
    33  
    34  // NewUDP network: "udp" "udp4" "udp6" address: "host:port"
    35  func NewUDP(network, address string, udpFunc UDPFunc, maxSize int) (udp *UDP) {
    36  	if maxSize < 1 {
    37  		maxSize = udpDefaultMaxSize
    38  	}
    39  	if netUDPAddr, err := net.ResolveUDPAddr(network, address); err != nil {
    40  		panic(perrors.Errorf("net.ResolveUDPAddr: '%w'", err))
    41  	} else {
    42  		udp = &UDP{UDPAddr: *netUDPAddr, Network: network, F: udpFunc, MaxSize: maxSize}
    43  	}
    44  	return
    45  }
    46  
    47  const (
    48  	udpDefaultMaxSize = 65507 // max for ipv4
    49  	oobSize           = 40
    50  	netReadOperation  = "read"
    51  	useOfClosed       = "use of closed network connection"
    52  )
    53  
    54  func (udp *UDP) Listen() (errCh <-chan error) {
    55  	udp.StartingListen.Add(1)
    56  	if !udp.ListenInvoked.CompareAndSwap(false, true) {
    57  		udp.StartingListen.Done()
    58  		panic(perrors.New("multiple udp.Listen invocations"))
    59  	}
    60  	if udp.IsShutdown.Load() {
    61  		udp.StartingListen.Done()
    62  		panic(perrors.New("udp.Listen after Shutdown"))
    63  	}
    64  	errChan := make(chan error)
    65  	errCh = errChan
    66  	udp.ErrCh = errChan
    67  	go udp.listenThread()
    68  	return
    69  }
    70  
    71  func (udp *UDP) listenThread() {
    72  	errCh := udp.ErrCh
    73  	defer close(errCh)
    74  	var FInvocations sync.WaitGroup
    75  	defer FInvocations.Wait()
    76  	var startingDone bool
    77  	defer func() {
    78  		if !startingDone {
    79  			udp.StartingListen.Done()
    80  		}
    81  	}()
    82  	defer parl.Recover2(func() parl.DA { return parl.A() }, nil, func(e error) { errCh <- e }) // capture panics
    83  
    84  	// listen
    85  	var netUDPConn *net.UDPConn // represents a network file descriptor
    86  	var err error
    87  	if netUDPConn, err = net.ListenUDP(udp.Network, &udp.UDPAddr); err != nil {
    88  		errCh <- perrors.Errorf("net.ListenUDP: '%w'", err)
    89  		return
    90  	}
    91  	if udp.setConn(netUDPConn) { // isShutdown
    92  		if err = netUDPConn.Close(); err != nil {
    93  			errCh <- perrors.Errorf("netUDPConn.Close: '%w'", err)
    94  		}
    95  		return
    96  	}
    97  	udp.Addr = netUDPConn.LocalAddr()
    98  	udp.IsListening.Store(true)
    99  	udp.StartingListen.Done()
   100  	startingDone = true
   101  	defer func() {
   102  		if !udp.IsShutdown.Load() {
   103  			if err := netUDPConn.Close(); err != nil {
   104  				errCh <- err
   105  			}
   106  		}
   107  		udp.IsListening.Store(false)
   108  	}()
   109  
   110  	// read datagrams
   111  	for {
   112  		b := make([]byte, udp.MaxSize)
   113  		oob := make([]byte, oobSize)
   114  		var n int
   115  		var oobn int
   116  		var flags int
   117  		var addr *net.UDPAddr
   118  		var err error
   119  		n, oobn, flags, addr, err = netUDPConn.ReadMsgUDP(b, oob)
   120  		if err != nil {
   121  			if udp.IsShutdown.Load() && udp.isClosedErr(err) {
   122  				return // we are shutdown
   123  			}
   124  			errCh <- perrors.Errorf("ReadMsgUDP: '%w'", err)
   125  			return
   126  		}
   127  		FInvocations.Add(1)
   128  		go func() {
   129  			defer FInvocations.Done()
   130  			udp.F(b[:n], oob[:oobn], flags, addr)
   131  		}()
   132  	}
   133  }
   134  
   135  func (udp *UDP) WaitForUp() (isUp bool, addr net.Addr) {
   136  	if !udp.ListenInvoked.Load() {
   137  		return // Listen has not been invoked
   138  	}
   139  	udp.StartingListen.Wait()
   140  	if isUp = udp.IsListening.Load(); isUp {
   141  		addr = udp.Addr
   142  	}
   143  	return
   144  }
   145  
   146  func (udp *UDP) isClosedErr(err error) (isClose bool) {
   147  	// read udp 127.0.0.1:50050: use of closed network connection
   148  	// &net.OpError{Op:"read", Net:"udp", Source:(*net.UDPAddr)(0xc00007a030), Addr:net.Addr(nil), Err:(*errors.errorString)(0xc000098160)}
   149  	opErr, ok := err.(*net.OpError)
   150  	if !ok {
   151  		return
   152  	}
   153  	if opErr.Op != netReadOperation {
   154  		return
   155  	}
   156  	e := opErr.Err
   157  	// &errors.errorString{s:"use of closed network connection"}
   158  	if e.Error() != useOfClosed {
   159  		return // some other error
   160  	}
   161  	isClose = true
   162  	return
   163  }
   164  
   165  func (udp *UDP) setConn(conn *net.UDPConn) (isShutdown bool) {
   166  	udp.connMutex.Lock()
   167  	defer udp.connMutex.Unlock()
   168  	isShutdown = udp.IsShutdown.Load()
   169  	if !isShutdown {
   170  		udp.NetUDPConn = conn
   171  	}
   172  	return
   173  }
   174  
   175  func (udp *UDP) Shutdown() {
   176  	udp.connMutex.RLock()
   177  	defer udp.connMutex.RUnlock()
   178  	if !udp.IsShutdown.CompareAndSwap(false, true) {
   179  		return // it was already shutdown
   180  	}
   181  	conn := udp.NetUDPConn
   182  	if conn == nil {
   183  		return // no need to close connection
   184  	}
   185  	err := conn.Close()
   186  	if err == nil {
   187  		return // conn successfully closed
   188  	}
   189  	udp.ErrCh <- err
   190  }