github.com/psiphon-inc/goarista@v0.0.0-20160825065156-d002785f4c67/lanz/client.go (about)

     1  // Copyright (C) 2016  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 lanz implements a LANZ client that will listen to notofications from LANZ streaming
     6  // server and will decode them and send them as a protobuf over a channel to a receiver.
     7  package lanz
     8  
     9  import (
    10  	"bufio"
    11  	"encoding/binary"
    12  	"io"
    13  	"net"
    14  	"time"
    15  
    16  	pb "github.com/aristanetworks/goarista/lanz/proto"
    17  
    18  	"github.com/aristanetworks/glog"
    19  	"github.com/golang/protobuf/proto"
    20  )
    21  
    22  const (
    23  	defaultConnectTimeout = 10 * time.Second
    24  	defaultConnectBackoff = 30 * time.Second
    25  )
    26  
    27  // Client is the LANZ client interface.
    28  type Client interface {
    29  	// Run is the main loop of the client.
    30  	// It connects to the LANZ server and reads the notifications, decodes them
    31  	// and sends them to the channel.
    32  	// In case of disconnect, it will reconnect automatically.
    33  	Run(ch chan<- *pb.LanzRecord)
    34  	// Stops the client.
    35  	Stop()
    36  }
    37  
    38  // ConnectReadCloser extends the io.ReadCloser interface with a Connect method.
    39  type ConnectReadCloser interface {
    40  	io.ReadCloser
    41  	// Connect connects to the address, returning an error if it fails.
    42  	Connect() error
    43  }
    44  
    45  type client struct {
    46  	addr     string
    47  	stopping bool
    48  	timeout  time.Duration
    49  	backoff  time.Duration
    50  	conn     ConnectReadCloser
    51  }
    52  
    53  // New creates a new client with default TCP connection to the LANZ server.
    54  func New(opts ...Option) Client {
    55  	c := &client{
    56  		stopping: false,
    57  		timeout:  defaultConnectTimeout,
    58  		backoff:  defaultConnectBackoff,
    59  	}
    60  
    61  	for _, opt := range opts {
    62  		opt(c)
    63  	}
    64  
    65  	if c.conn == nil {
    66  		if c.addr == "" {
    67  			panic("Neither address, nor connector specified")
    68  		}
    69  		c.conn = &netConnector{
    70  			addr:    c.addr,
    71  			timeout: c.timeout,
    72  			backoff: c.backoff,
    73  		}
    74  	}
    75  
    76  	return c
    77  }
    78  
    79  func (c *client) Run(ch chan<- *pb.LanzRecord) {
    80  	for !c.stopping {
    81  		if err := c.conn.Connect(); err != nil && !c.stopping {
    82  			glog.V(1).Infof("Can't connect to LANZ server: %v", err)
    83  			time.Sleep(c.backoff)
    84  			continue
    85  		}
    86  		glog.V(1).Infof("Connected successfully to LANZ server: %v", c.addr)
    87  		if err := c.read(bufio.NewReader(c.conn), ch); err != nil && !c.stopping {
    88  			if err != io.EOF && err != io.ErrUnexpectedEOF {
    89  				glog.Errorf("Error receiving LANZ events: %v", err)
    90  			}
    91  			c.conn.Close()
    92  			time.Sleep(c.backoff)
    93  		}
    94  	}
    95  
    96  	close(ch)
    97  }
    98  
    99  func (c *client) read(r *bufio.Reader, ch chan<- *pb.LanzRecord) error {
   100  	for !c.stopping {
   101  		len, err := binary.ReadUvarint(r)
   102  		if err != nil {
   103  			return err
   104  		}
   105  
   106  		buf := make([]byte, len)
   107  		if _, err = io.ReadFull(r, buf); err != nil {
   108  			return err
   109  		}
   110  
   111  		rec := &pb.LanzRecord{}
   112  		if err = proto.Unmarshal(buf, rec); err != nil {
   113  			return err
   114  		}
   115  
   116  		ch <- rec
   117  	}
   118  
   119  	return nil
   120  }
   121  
   122  func (c *client) Stop() {
   123  	if c.stopping {
   124  		return
   125  	}
   126  
   127  	c.stopping = true
   128  	c.conn.Close()
   129  }
   130  
   131  type netConnector struct {
   132  	net.Conn
   133  	addr    string
   134  	timeout time.Duration
   135  	backoff time.Duration
   136  }
   137  
   138  func (c *netConnector) Connect() (err error) {
   139  	c.Conn, err = net.DialTimeout("tcp", c.addr, c.timeout)
   140  	if err != nil {
   141  	}
   142  	return
   143  }