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

     1  //go:build ZMQ4
     2  // +build ZMQ4
     3  
     4  package writer
     5  
     6  import (
     7  	"fmt"
     8  	"strings"
     9  	"time"
    10  
    11  	"github.com/Jeffail/benthos/v3/lib/log"
    12  	"github.com/Jeffail/benthos/v3/lib/message"
    13  	"github.com/Jeffail/benthos/v3/lib/metrics"
    14  	"github.com/Jeffail/benthos/v3/lib/types"
    15  	"github.com/pebbe/zmq4"
    16  )
    17  
    18  //------------------------------------------------------------------------------
    19  
    20  // ZMQ4 is an output type that writes ZMQ4 messages.
    21  type ZMQ4 struct {
    22  	log   log.Modular
    23  	stats metrics.Type
    24  
    25  	urls []string
    26  	conf *ZMQ4Config
    27  
    28  	pollTimeout time.Duration
    29  	poller      *zmq4.Poller
    30  	socket      *zmq4.Socket
    31  }
    32  
    33  // NewZMQ4 creates a new ZMQ4 output type.
    34  func NewZMQ4(conf *ZMQ4Config, log log.Modular, stats metrics.Type) (*ZMQ4, error) {
    35  	z := ZMQ4{
    36  		log:   log,
    37  		stats: stats,
    38  		conf:  conf,
    39  	}
    40  
    41  	_, err := getZMQType(conf.SocketType)
    42  	if nil != err {
    43  		return nil, err
    44  	}
    45  
    46  	if tout := conf.PollTimeout; len(tout) > 0 {
    47  		var err error
    48  		if z.pollTimeout, err = time.ParseDuration(tout); err != nil {
    49  			return nil, fmt.Errorf("failed to parse poll timeout string: %v", err)
    50  		}
    51  	}
    52  
    53  	for _, u := range conf.URLs {
    54  		for _, splitU := range strings.Split(u, ",") {
    55  			if len(splitU) > 0 {
    56  				z.urls = append(z.urls, splitU)
    57  			}
    58  		}
    59  	}
    60  
    61  	return &z, nil
    62  }
    63  
    64  //------------------------------------------------------------------------------
    65  
    66  func getZMQType(t string) (zmq4.Type, error) {
    67  	switch t {
    68  	case "PUB":
    69  		return zmq4.PUB, nil
    70  	case "PUSH":
    71  		return zmq4.PUSH, nil
    72  	}
    73  	return zmq4.PULL, types.ErrInvalidZMQType
    74  }
    75  
    76  //------------------------------------------------------------------------------
    77  
    78  // Connect attempts to establish a connection to a ZMQ4 socket.
    79  func (z *ZMQ4) Connect() error {
    80  	if z.socket != nil {
    81  		return nil
    82  	}
    83  
    84  	t, err := getZMQType(z.conf.SocketType)
    85  	if nil != err {
    86  		return err
    87  	}
    88  
    89  	ctx, err := zmq4.NewContext()
    90  	if nil != err {
    91  		return err
    92  	}
    93  
    94  	var socket *zmq4.Socket
    95  	if socket, err = ctx.NewSocket(t); nil != err {
    96  		return err
    97  	}
    98  
    99  	defer func() {
   100  		if err != nil && socket != nil {
   101  			socket.Close()
   102  		}
   103  	}()
   104  
   105  	socket.SetSndhwm(z.conf.HighWaterMark)
   106  
   107  	for _, address := range z.urls {
   108  		if z.conf.Bind {
   109  			err = socket.Bind(address)
   110  		} else {
   111  			err = socket.Connect(address)
   112  		}
   113  		if err != nil {
   114  			return err
   115  		}
   116  	}
   117  
   118  	z.socket = socket
   119  	z.poller = zmq4.NewPoller()
   120  	z.poller.Add(z.socket, zmq4.POLLOUT)
   121  
   122  	z.log.Infof("Sending ZMQ4 messages to URLs: %s\n", z.urls)
   123  	return nil
   124  }
   125  
   126  // Write will attempt to write a message to the ZMQ4 socket.
   127  func (z *ZMQ4) Write(msg types.Message) error {
   128  	if z.socket == nil {
   129  		return types.ErrNotConnected
   130  	}
   131  	_, err := z.socket.SendMessageDontwait(message.GetAllBytes(msg))
   132  	if err != nil {
   133  		var polled []zmq4.Polled
   134  		if polled, err = z.poller.Poll(z.pollTimeout); len(polled) == 1 {
   135  			_, err = z.socket.SendMessage(message.GetAllBytes(msg))
   136  		} else if err == nil {
   137  			return types.ErrTimeout
   138  		}
   139  	}
   140  	return err
   141  }
   142  
   143  // CloseAsync shuts down the ZMQ4 output and stops processing messages.
   144  func (z *ZMQ4) CloseAsync() {
   145  	if z.socket != nil {
   146  		z.socket.Close()
   147  		z.socket = nil
   148  	}
   149  }
   150  
   151  // WaitForClose blocks until the ZMQ4 output has closed down.
   152  func (z *ZMQ4) WaitForClose(timeout time.Duration) error {
   153  	return nil
   154  }
   155  
   156  //------------------------------------------------------------------------------