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

     1  package input
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"fmt"
     7  	"net"
     8  	"sync"
     9  
    10  	"github.com/observiq/carbon/operator"
    11  	"github.com/observiq/carbon/operator/helper"
    12  	"go.uber.org/zap"
    13  )
    14  
    15  func init() {
    16  	operator.Register("tcp_input", func() operator.Builder { return NewTCPInputConfig("") })
    17  }
    18  
    19  func NewTCPInputConfig(operatorID string) *TCPInputConfig {
    20  	return &TCPInputConfig{
    21  		InputConfig: helper.NewInputConfig(operatorID, "tcp_input"),
    22  	}
    23  }
    24  
    25  // TCPInputConfig is the configuration of a tcp input operator.
    26  type TCPInputConfig struct {
    27  	helper.InputConfig `yaml:",inline"`
    28  
    29  	ListenAddress string `json:"listen_address,omitempty" yaml:"listen_address,omitempty"`
    30  }
    31  
    32  // Build will build a tcp input operator.
    33  func (c TCPInputConfig) Build(context operator.BuildContext) (operator.Operator, error) {
    34  	inputOperator, err := c.InputConfig.Build(context)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  
    39  	if c.ListenAddress == "" {
    40  		return nil, fmt.Errorf("missing required parameter 'listen_address'")
    41  	}
    42  
    43  	address, err := net.ResolveTCPAddr("tcp", c.ListenAddress)
    44  	if err != nil {
    45  		return nil, fmt.Errorf("failed to resolve listen_address: %s", err)
    46  	}
    47  
    48  	tcpInput := &TCPInput{
    49  		InputOperator: inputOperator,
    50  		address:       address,
    51  	}
    52  	return tcpInput, nil
    53  }
    54  
    55  // TCPInput is an operator that listens for log entries over tcp.
    56  type TCPInput struct {
    57  	helper.InputOperator
    58  	address *net.TCPAddr
    59  
    60  	listener  *net.TCPListener
    61  	cancel    context.CancelFunc
    62  	waitGroup *sync.WaitGroup
    63  }
    64  
    65  // Start will start listening for log entries over tcp.
    66  func (t *TCPInput) Start() error {
    67  	listener, err := net.ListenTCP("tcp", t.address)
    68  	if err != nil {
    69  		return fmt.Errorf("failed to listen on interface: %w", err)
    70  	}
    71  
    72  	t.listener = listener
    73  	ctx, cancel := context.WithCancel(context.Background())
    74  	t.cancel = cancel
    75  	t.waitGroup = &sync.WaitGroup{}
    76  	t.goListen(ctx)
    77  	return nil
    78  }
    79  
    80  // goListenn will listen for tcp connections.
    81  func (t *TCPInput) goListen(ctx context.Context) {
    82  	t.waitGroup.Add(1)
    83  
    84  	go func() {
    85  		defer t.waitGroup.Done()
    86  
    87  		for {
    88  			conn, err := t.listener.AcceptTCP()
    89  			if err != nil {
    90  				select {
    91  				case <-ctx.Done():
    92  					return
    93  				default:
    94  					t.Debugw("Listener accept error", zap.Error(err))
    95  				}
    96  			}
    97  
    98  			t.Debugf("Received connection: %s", conn.RemoteAddr().String())
    99  			subctx, cancel := context.WithCancel(ctx)
   100  			t.goHandleClose(subctx, conn)
   101  			t.goHandleMessages(subctx, conn, cancel)
   102  		}
   103  	}()
   104  }
   105  
   106  // goHandleClose will wait for the context to finish before closing a connection.
   107  func (t *TCPInput) goHandleClose(ctx context.Context, conn net.Conn) {
   108  	t.waitGroup.Add(1)
   109  
   110  	go func() {
   111  		defer t.waitGroup.Done()
   112  		<-ctx.Done()
   113  		t.Debugf("Closing connection: %s", conn.RemoteAddr().String())
   114  		if err := conn.Close(); err != nil {
   115  			t.Errorf("Failed to close connection: %s", err)
   116  		}
   117  	}()
   118  }
   119  
   120  // goHandleMessages will handles messages from a tcp connection.
   121  func (t *TCPInput) goHandleMessages(ctx context.Context, conn net.Conn, cancel context.CancelFunc) {
   122  	t.waitGroup.Add(1)
   123  
   124  	go func() {
   125  		defer t.waitGroup.Done()
   126  		defer cancel()
   127  
   128  		scanner := bufio.NewScanner(conn)
   129  		for scanner.Scan() {
   130  			entry, err := t.NewEntry(scanner.Text())
   131  			if err != nil {
   132  				t.Errorw("Failed to create entry", zap.Error(err))
   133  				continue
   134  			}
   135  			t.Write(ctx, entry)
   136  		}
   137  		if err := scanner.Err(); err != nil {
   138  			t.Errorw("Scanner error", zap.Error(err))
   139  		}
   140  	}()
   141  }
   142  
   143  // Stop will stop listening for log entries over TCP.
   144  func (t *TCPInput) Stop() error {
   145  	t.cancel()
   146  
   147  	if err := t.listener.Close(); err != nil {
   148  		return err
   149  	}
   150  
   151  	t.waitGroup.Wait()
   152  	return nil
   153  }