github.com/aristanetworks/goarista@v0.0.0-20240514173732-cca2755bbd44/cmd/octsdb/udp.go (about)

     1  // Copyright (c) 2017 Arista Networks, Inc.
     2  // Use of this source code is governed by the Apache License 2.0
     3  // that can be found in the COPYING file.
     4  
     5  package main
     6  
     7  import (
     8  	"math/rand"
     9  	"time"
    10  
    11  	"github.com/aristanetworks/glog"
    12  	kcp "github.com/xtaci/kcp-go"
    13  )
    14  
    15  type udpClient struct {
    16  	addr    string
    17  	conn    *kcp.UDPSession
    18  	parity  int
    19  	timeout time.Duration
    20  }
    21  
    22  func newUDPClient(addr string, parity int, timeout time.Duration) OpenTSDBConn {
    23  	return &udpClient{
    24  		addr:    addr,
    25  		parity:  parity,
    26  		timeout: timeout,
    27  	}
    28  }
    29  
    30  func (c *udpClient) Put(d *DataPoint) error {
    31  	var err error
    32  	if c.conn == nil {
    33  		// Prevent a bunch of clients all disconnecting and attempting to reconnect
    34  		// at nearly the same time.
    35  		time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond)
    36  
    37  		c.conn, err = kcp.DialWithOptions(c.addr, nil, 10, c.parity)
    38  		if err != nil {
    39  			return err
    40  		}
    41  		c.conn.SetNoDelay(1, 40, 1, 1) // Suggested by kcp-go to lower cpu usage
    42  	}
    43  
    44  	dStr := d.String()
    45  	glog.V(3).Info(dStr)
    46  
    47  	c.conn.SetWriteDeadline(time.Now().Add(c.timeout))
    48  	_, err = c.conn.Write([]byte(dStr))
    49  	if err != nil {
    50  		c.conn.Close()
    51  		c.conn = nil
    52  	}
    53  	return err
    54  }
    55  
    56  type udpServer struct {
    57  	lis    *kcp.Listener
    58  	telnet *telnetClient
    59  }
    60  
    61  func newUDPServer(udpAddr, tsdbAddr string, parity int) (*udpServer, error) {
    62  	lis, err := kcp.ListenWithOptions(udpAddr, nil, 10, parity)
    63  	if err != nil {
    64  		return nil, err
    65  	}
    66  	return &udpServer{
    67  		lis:    lis,
    68  		telnet: newTelnetClient(tsdbAddr).(*telnetClient),
    69  	}, nil
    70  }
    71  
    72  func (c *udpServer) Run() error {
    73  	for {
    74  		conn, err := c.lis.AcceptKCP()
    75  		if err != nil {
    76  			return err
    77  		}
    78  		conn.SetNoDelay(1, 40, 1, 1) // Suggested by kcp-go to lower cpu usage
    79  		if glog.V(3) {
    80  			glog.Infof("New connection from %s", conn.RemoteAddr())
    81  		}
    82  		go func() {
    83  			defer conn.Close()
    84  			var buf [4096]byte
    85  			for {
    86  				n, err := conn.Read(buf[:])
    87  				if err != nil {
    88  					if n != 0 { // Not EOF
    89  						glog.Error(err)
    90  					}
    91  					return
    92  				}
    93  				if glog.V(3) {
    94  					glog.Info(string(buf[:n]))
    95  				}
    96  				err = c.telnet.PutBytes(buf[:n])
    97  				if err != nil {
    98  					glog.Error(err)
    99  					return
   100  				}
   101  			}
   102  		}()
   103  	}
   104  }