github.com/awesome-flow/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/corev1alpha1/actor/sink_head_udp.go (about)

     1  package actor
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"time"
     7  )
     8  
     9  const (
    10  	UDPConnTimeout = 5 * time.Second
    11  )
    12  
    13  type UDPConnBuilder func(*net.UDPAddr, time.Duration) (net.Conn, error)
    14  
    15  var DefaultUDPConnBuilder = func(udpaddr *net.UDPAddr, timeout time.Duration) (net.Conn, error) {
    16  	return net.DialTimeout("udp", udpaddr.String(), timeout)
    17  }
    18  
    19  type SinkHeadUDP struct {
    20  	addr        *net.UDPAddr
    21  	conn        net.Conn
    22  	connbuilder UDPConnBuilder
    23  
    24  	ConnectTimeout time.Duration
    25  	WriteTimeout   time.Duration
    26  }
    27  
    28  var _ (SinkHead) = (*SinkHeadUDP)(nil)
    29  
    30  func NewSinkHeadUDP(udpaddr *net.UDPAddr) (*SinkHeadUDP, error) {
    31  	return &SinkHeadUDP{
    32  		addr:           udpaddr,
    33  		connbuilder:    DefaultUDPConnBuilder,
    34  		ConnectTimeout: UDPConnTimeout,
    35  	}, nil
    36  }
    37  
    38  func (h *SinkHeadUDP) Connect() error {
    39  	conn, err := h.connbuilder(h.addr, h.ConnectTimeout)
    40  	if err != nil {
    41  		return err
    42  	}
    43  	h.conn = conn
    44  
    45  	return nil
    46  }
    47  
    48  func (h *SinkHeadUDP) Start() error {
    49  	return nil
    50  }
    51  
    52  func (h *SinkHeadUDP) Stop() error {
    53  	if h.conn != nil {
    54  		return h.conn.Close()
    55  	}
    56  	return nil
    57  }
    58  
    59  func (h *SinkHeadUDP) Write(data []byte) (int, error, bool) {
    60  	if h.conn == nil {
    61  		return 0, fmt.Errorf("udp sink head conn is nil"), true
    62  	}
    63  	n, err := h.conn.Write(data)
    64  	rec := false
    65  	if err != nil {
    66  		h.conn = nil
    67  		rec = true
    68  	}
    69  
    70  	return n, err, rec
    71  }