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 //------------------------------------------------------------------------------