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

     1  package writer
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/Jeffail/benthos/v3/internal/codec"
    11  	"github.com/Jeffail/benthos/v3/lib/log"
    12  	"github.com/Jeffail/benthos/v3/lib/metrics"
    13  	"github.com/Jeffail/benthos/v3/lib/types"
    14  )
    15  
    16  //------------------------------------------------------------------------------
    17  
    18  // SocketConfig contains configuration fields for the Socket output type.
    19  type SocketConfig struct {
    20  	Network string `json:"network" yaml:"network"`
    21  	Address string `json:"address" yaml:"address"`
    22  	Codec   string `json:"codec" yaml:"codec"`
    23  }
    24  
    25  // NewSocketConfig creates a new SocketConfig with default values.
    26  func NewSocketConfig() SocketConfig {
    27  	return SocketConfig{
    28  		Network: "unix",
    29  		Address: "/tmp/benthos.sock",
    30  		Codec:   "lines",
    31  	}
    32  }
    33  
    34  //------------------------------------------------------------------------------
    35  
    36  // Socket is an output type that sends messages as a continuous steam of line
    37  // delimied messages over socket.
    38  type Socket struct {
    39  	network   string
    40  	address   string
    41  	codec     codec.WriterConstructor
    42  	codecConf codec.WriterConfig
    43  
    44  	stats metrics.Type
    45  	log   log.Modular
    46  
    47  	writer    codec.Writer
    48  	writerMut sync.Mutex
    49  }
    50  
    51  // NewSocket creates a new Socket writer type.
    52  func NewSocket(
    53  	conf SocketConfig,
    54  	mgr types.Manager,
    55  	log log.Modular,
    56  	stats metrics.Type,
    57  ) (*Socket, error) {
    58  	switch conf.Network {
    59  	case "tcp", "udp", "unix":
    60  	default:
    61  		return nil, fmt.Errorf("socket network '%v' is not supported by this output", conf.Network)
    62  	}
    63  	codec, codecConf, err := codec.GetWriter(conf.Codec)
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  	t := Socket{
    68  		network:   conf.Network,
    69  		address:   conf.Address,
    70  		codec:     codec,
    71  		codecConf: codecConf,
    72  		stats:     stats,
    73  		log:       log,
    74  	}
    75  	return &t, nil
    76  }
    77  
    78  //------------------------------------------------------------------------------
    79  
    80  // Connect establises a connection to the target socket server.
    81  func (s *Socket) Connect() error {
    82  	return s.ConnectWithContext(context.Background())
    83  }
    84  
    85  // ConnectWithContext establises a connection to the target socket server.
    86  func (s *Socket) ConnectWithContext(ctx context.Context) error {
    87  	s.writerMut.Lock()
    88  	defer s.writerMut.Unlock()
    89  	if s.writer != nil {
    90  		return nil
    91  	}
    92  
    93  	conn, err := net.Dial(s.network, s.address)
    94  	if err != nil {
    95  		return err
    96  	}
    97  
    98  	s.writer, err = s.codec(conn)
    99  	if err != nil {
   100  		conn.Close()
   101  		return err
   102  	}
   103  
   104  	s.log.Infof("Sending messages over %v socket to: %s\n", s.network, s.address)
   105  	return nil
   106  }
   107  
   108  // Write attempts to write a message.
   109  func (s *Socket) Write(msg types.Message) error {
   110  	return s.WriteWithContext(context.Background(), msg)
   111  }
   112  
   113  // WriteWithContext attempts to write a message.
   114  func (s *Socket) WriteWithContext(ctx context.Context, msg types.Message) error {
   115  	s.writerMut.Lock()
   116  	w := s.writer
   117  	s.writerMut.Unlock()
   118  
   119  	if w == nil {
   120  		return types.ErrNotConnected
   121  	}
   122  
   123  	err := msg.Iter(func(i int, part types.Part) error {
   124  		serr := w.Write(ctx, part)
   125  		if serr != nil || s.codecConf.CloseAfter {
   126  			s.writerMut.Lock()
   127  			s.writer.Close(ctx)
   128  			s.writer = nil
   129  			s.writerMut.Unlock()
   130  		}
   131  		return serr
   132  	})
   133  	if err == nil && msg.Len() > 1 {
   134  		if err = w.EndBatch(); err != nil {
   135  			s.writerMut.Lock()
   136  			s.writer.Close(ctx)
   137  			s.writer = nil
   138  			s.writerMut.Unlock()
   139  		}
   140  	}
   141  	return err
   142  }
   143  
   144  // CloseAsync shuts down the socket output and stops processing messages.
   145  func (s *Socket) CloseAsync() {
   146  	s.writerMut.Lock()
   147  	if s.writer != nil {
   148  		s.writer.Close(context.Background())
   149  		s.writer = nil
   150  	}
   151  	s.writerMut.Unlock()
   152  }
   153  
   154  // WaitForClose blocks until the socket output has closed down.
   155  func (s *Socket) WaitForClose(timeout time.Duration) error {
   156  	return nil
   157  }
   158  
   159  //------------------------------------------------------------------------------