github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/operator/builtin/input/udp.go (about)

     1  package input
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"sync"
     8  
     9  	"github.com/observiq/carbon/operator"
    10  	"github.com/observiq/carbon/operator/helper"
    11  	"go.uber.org/zap"
    12  )
    13  
    14  func init() {
    15  	operator.Register("udp_input", func() operator.Builder { return NewUDPInputConfig("") })
    16  }
    17  
    18  func NewUDPInputConfig(operatorID string) *UDPInputConfig {
    19  	return &UDPInputConfig{
    20  		InputConfig: helper.NewInputConfig(operatorID, "udp_input"),
    21  	}
    22  }
    23  
    24  // UDPInputConfig is the configuration of a udp input operator.
    25  type UDPInputConfig struct {
    26  	helper.InputConfig `yaml:",inline"`
    27  
    28  	ListenAddress string `json:"listen_address,omitempty" yaml:"listen_address,omitempty"`
    29  }
    30  
    31  // Build will build a udp input operator.
    32  func (c UDPInputConfig) Build(context operator.BuildContext) (operator.Operator, error) {
    33  	inputOperator, err := c.InputConfig.Build(context)
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  
    38  	if c.ListenAddress == "" {
    39  		return nil, fmt.Errorf("missing required parameter 'listen_address'")
    40  	}
    41  
    42  	address, err := net.ResolveUDPAddr("udp", c.ListenAddress)
    43  	if err != nil {
    44  		return nil, fmt.Errorf("failed to resolve listen_address: %s", err)
    45  	}
    46  
    47  	udpInput := &UDPInput{
    48  		InputOperator: inputOperator,
    49  		address:       address,
    50  		buffer:        make([]byte, 8192),
    51  	}
    52  	return udpInput, nil
    53  }
    54  
    55  // UDPInput is an operator that listens to a socket for log entries.
    56  type UDPInput struct {
    57  	buffer []byte
    58  	helper.InputOperator
    59  	address *net.UDPAddr
    60  
    61  	connection net.PacketConn
    62  	cancel     context.CancelFunc
    63  	waitGroup  *sync.WaitGroup
    64  }
    65  
    66  // Start will start listening for messages on a socket.
    67  func (u *UDPInput) Start() error {
    68  	ctx, cancel := context.WithCancel(context.Background())
    69  	u.cancel = cancel
    70  	u.waitGroup = &sync.WaitGroup{}
    71  
    72  	conn, err := net.ListenUDP("udp", u.address)
    73  	if err != nil {
    74  		return fmt.Errorf("failed to open connection: %s", err)
    75  	}
    76  	u.connection = conn
    77  
    78  	u.goHandleMessages(ctx)
    79  	return nil
    80  }
    81  
    82  // goHandleMessages will handle messages from a udp connection.
    83  func (u *UDPInput) goHandleMessages(ctx context.Context) {
    84  	u.waitGroup.Add(1)
    85  
    86  	go func() {
    87  		defer u.waitGroup.Done()
    88  
    89  		for {
    90  			message, err := u.readMessage()
    91  			if err != nil {
    92  				select {
    93  				case <-ctx.Done():
    94  					return
    95  				default:
    96  					u.Errorw("Failed reading messages", zap.Error(err))
    97  				}
    98  				break
    99  			}
   100  
   101  			entry, err := u.NewEntry(message)
   102  			if err != nil {
   103  				u.Errorw("Failed to create entry", zap.Error(err))
   104  				continue
   105  			}
   106  
   107  			u.Write(ctx, entry)
   108  		}
   109  	}()
   110  }
   111  
   112  // readMessage will read log messages from the connection.
   113  func (u *UDPInput) readMessage() (string, error) {
   114  	n, _, err := u.connection.ReadFrom(u.buffer)
   115  	if err != nil {
   116  		return "", err
   117  	}
   118  
   119  	// Remove trailing characters and NULs
   120  	for ; (n > 0) && (u.buffer[n-1] < 32); n-- {
   121  	}
   122  
   123  	return string(u.buffer[:n]), nil
   124  }
   125  
   126  // Stop will stop listening for udp messages.
   127  func (u *UDPInput) Stop() error {
   128  	u.cancel()
   129  	u.connection.Close()
   130  	u.waitGroup.Wait()
   131  	return nil
   132  }