github.com/Jeffail/benthos/v3@v3.65.0/lib/output/writer/tcp.go (about)

     1  package writer
     2  
     3  import (
     4  	"net"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/Jeffail/benthos/v3/lib/log"
     9  	"github.com/Jeffail/benthos/v3/lib/metrics"
    10  	"github.com/Jeffail/benthos/v3/lib/types"
    11  )
    12  
    13  //------------------------------------------------------------------------------
    14  
    15  // TCPConfig contains configuration fields for the TCP output type.
    16  type TCPConfig struct {
    17  	Address string `json:"address" yaml:"address"`
    18  }
    19  
    20  // NewTCPConfig creates a new TCPConfig with default values.
    21  func NewTCPConfig() TCPConfig {
    22  	return TCPConfig{
    23  		Address: "localhost:4194",
    24  	}
    25  }
    26  
    27  //------------------------------------------------------------------------------
    28  
    29  // TCP is an output type that sends messages as a continuous steam of line
    30  // delimied messages over TCP.
    31  type TCP struct {
    32  	connMut sync.Mutex
    33  	conn    net.Conn
    34  
    35  	address string
    36  
    37  	stats metrics.Type
    38  	log   log.Modular
    39  }
    40  
    41  // NewTCP creates a new TCP writer type.
    42  func NewTCP(
    43  	conf TCPConfig,
    44  	mgr types.Manager,
    45  	log log.Modular,
    46  	stats metrics.Type,
    47  ) (*TCP, error) {
    48  	t := TCP{
    49  		address: conf.Address,
    50  		stats:   stats,
    51  		log:     log,
    52  	}
    53  	return &t, nil
    54  }
    55  
    56  //------------------------------------------------------------------------------
    57  
    58  // Connect does nothing.
    59  func (t *TCP) Connect() error {
    60  	t.connMut.Lock()
    61  	defer t.connMut.Unlock()
    62  	if t.conn != nil {
    63  		return nil
    64  	}
    65  
    66  	var err error
    67  	if t.conn, err = net.Dial("tcp", t.address); err != nil {
    68  		return err
    69  	}
    70  
    71  	t.log.Infof("Sending messages over TCP to: %s\n", t.address)
    72  	return nil
    73  }
    74  
    75  // Write attempts to write a message.
    76  func (t *TCP) Write(msg types.Message) error {
    77  	t.connMut.Lock()
    78  	conn := t.conn
    79  	t.connMut.Unlock()
    80  
    81  	if conn == nil {
    82  		return types.ErrNotConnected
    83  	}
    84  
    85  	err := msg.Iter(func(i int, part types.Part) error {
    86  		partBytes := part.Get()
    87  		if partBytes[len(partBytes)-1] != '\n' {
    88  			partBytes = append(partBytes[:len(partBytes):len(partBytes)], []byte("\n")...)
    89  		}
    90  		_, werr := conn.Write(partBytes)
    91  		return werr
    92  	})
    93  	if err == nil && msg.Len() > 1 {
    94  		_, err = conn.Write([]byte("\n"))
    95  	}
    96  	if err != nil {
    97  		t.connMut.Lock()
    98  		t.conn.Close()
    99  		t.conn = nil
   100  		t.connMut.Unlock()
   101  	}
   102  	return err
   103  }
   104  
   105  // CloseAsync shuts down the TCP output and stops processing messages.
   106  func (t *TCP) CloseAsync() {
   107  	t.connMut.Lock()
   108  	if t.conn != nil {
   109  		t.conn.Close()
   110  		t.conn = nil
   111  	}
   112  	t.connMut.Unlock()
   113  }
   114  
   115  // WaitForClose blocks until the TCP output has closed down.
   116  func (t *TCP) WaitForClose(timeout time.Duration) error {
   117  	return nil
   118  }
   119  
   120  //------------------------------------------------------------------------------