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

     1  package actor
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"time"
     7  )
     8  
     9  const (
    10  	TCPConnTimeout = 5 * time.Second
    11  )
    12  
    13  type TCPConnBuilder func(*net.TCPAddr, time.Duration) (net.Conn, error)
    14  
    15  var DefaultTCPConnBuilder = func(tcpaddr *net.TCPAddr, timeout time.Duration) (net.Conn, error) {
    16  	return net.DialTimeout("tcp", tcpaddr.String(), timeout)
    17  }
    18  
    19  type SinkHeadTCP struct {
    20  	addr        *net.TCPAddr
    21  	conn        net.Conn
    22  	connbuilder TCPConnBuilder
    23  
    24  	ConnectTimeout time.Duration
    25  	WriteTimeout   time.Duration
    26  }
    27  
    28  var _ (SinkHead) = (*SinkHeadTCP)(nil)
    29  
    30  func NewSinkHeadTCP(tcpaddr *net.TCPAddr) (*SinkHeadTCP, error) {
    31  	return &SinkHeadTCP{
    32  		addr:           tcpaddr,
    33  		connbuilder:    DefaultTCPConnBuilder,
    34  		ConnectTimeout: TCPConnTimeout,
    35  	}, nil
    36  }
    37  
    38  func (h *SinkHeadTCP) 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 *SinkHeadTCP) Start() error {
    49  	return nil
    50  }
    51  
    52  func (h *SinkHeadTCP) Stop() error {
    53  	if h.conn != nil {
    54  		return h.conn.Close()
    55  	}
    56  	return nil
    57  }
    58  
    59  func (h *SinkHeadTCP) Write(data []byte) (int, error, bool) {
    60  	if h.conn == nil {
    61  		return 0, fmt.Errorf("tcp sink head conn is nil"), true
    62  	}
    63  	l := len(data)
    64  	buf := make([]byte, l+2)
    65  	copy(buf, data)
    66  	copy(buf[l:], []byte("\r\n"))
    67  	rec := false
    68  	n, err := h.conn.Write(buf)
    69  	if err != nil {
    70  		rec = true
    71  		h.conn = nil
    72  	}
    73  
    74  	return n, err, rec
    75  }