vendor/github.com/go-kit/kit/util/conn/manager.go

   1 package conn
   2
   3 import (
   4         "errors"
   5         "net"
   6         "time"
   7
   8         "github.com/go-kit/kit/log"
   9 )
  10
  11 // Dialer imitates net.Dial. Dialer is assumed to yield connections that are
  12 // safe for use by multiple concurrent goroutines.
  13 type Dialer func(network, address string) (net.Conn, error)
  14
  15 // AfterFunc imitates time.After.
  16 type AfterFunc func(time.Duration) <-chan time.Time
  17
  18 // Manager manages a net.Conn.
  19 //
  20 // Clients provide a way to create the connection with a Dialer, network, and
  21 // address. Clients should Take the connection when they want to use it, and Put
  22 // back whatever error they receive from its use. When a non-nil error is Put,
  23 // the connection is invalidated, and a new connection is established.
  24 // Connection failures are retried after an exponential backoff.
  25 type Manager struct {
  26         dialer  Dialer
  27         network string
  28         address string
  29         after   AfterFunc
  30         logger  log.Logger
  31
  32         takec chan net.Conn
  33         putc  chan error
  34 }
  35
  36 // NewManager returns a connection manager using the passed Dialer, network, and
  37 // address. The AfterFunc is used to control exponential backoff and retries.
  38 // The logger is used to log errors; pass a log.NopLogger if you don't care to
  39 // receive them. For normal use, prefer NewDefaultManager.
  40 func NewManager(d Dialer, network, address string, after AfterFunc, logger log.Logger) *Manager {
  41         m := &Manager{
  42                 dialer:  d,
  43                 network: network,
  44                 address: address,
  45                 after:   after,
  46                 logger:  logger,
  47
  48                 takec: make(chan net.Conn),
  49                 putc:  make(chan error),
  50         }
  51         go m.loop()
  52         return m
  53 }
  54
  55 // NewDefaultManager is a helper constructor, suitable for most normal use in
  56 // real (non-test) code. It uses the real net.Dial and time.After functions.
  57 func NewDefaultManager(network, address string, logger log.Logger) *Manager {
  58         return NewManager(net.Dial, network, address, time.After, logger)
  59 }
  60
  61 // Take yields the current connection. It may be nil.
  62 func (m *Manager) Take() net.Conn {
  63         return <-m.takec
  64 }
  65
  66 // Put accepts an error that came from a previously yielded connection. If the
  67 // error is non-nil, the manager will invalidate the current connection and try
  68 // to reconnect, with exponential backoff. Putting a nil error is a no-op.
  69 func (m *Manager) Put(err error) {
  70         m.putc <- err
  71 }
  72
  73 // Write writes the passed data to the connection in a single Take/Put cycle.
  74 func (m *Manager) Write(b []byte) (int, error) {
  75         conn := m.Take()
  76         if conn == nil {
  77                 return 0, ErrConnectionUnavailable
  78         }
  79         n, err := conn.Write(b)
  80         defer m.Put(err)
  81         return n, err
  82 }
  83
  84 func (m *Manager) loop() {
  85         var (
  86                 conn       = dial(m.dialer, m.network, m.address, m.logger) // may block slightly
  87                 connc      = make(chan net.Conn, 1)
  88                 reconnectc <-chan time.Time // initially nil
  89                 backoff    = time.Second
  90         )
  91
  92         // If the initial dial fails, we need to trigger a reconnect via the loop
  93         // body, below. If we did this in a goroutine, we would race on the conn
  94         // variable. So we use a buffered chan instead.
  95         connc <- conn
  96
  97         for {
  98                 select {
  99                 case <-reconnectc:
 100                         reconnectc = nil // one-shot
 101                         go func() { connc <- dial(m.dialer, m.network, m.address, m.logger) }()
 102
 103                 case conn = <-connc:
 104                         if conn == nil {
 105                                 // didn't work
 106                                 backoff = exponential(backoff) // wait longer
 107                                 reconnectc = m.after(backoff)  // try again
 108                         } else {
 109                                 // worked!
 110                                 backoff = time.Second // reset wait time
 111                                 reconnectc = nil      // no retry necessary
 112                         }
 113
 114                 case m.takec <- conn:
 115
 116                 case err := <-m.putc:
 117                         if err != nil && conn != nil {
 118                                 m.logger.Log("err", err)
 119                                 conn = nil                            // connection is bad
 120                                 reconnectc = m.after(time.Nanosecond) // trigger immediately
 121                         }
 122                 }
 123         }
 124 }
 125
 126 func dial(d Dialer, network, address string, logger log.Logger) net.Conn {
 127         conn, err := d(network, address)
 128         if err != nil {
 129                 logger.Log("err", err)
 130                 conn = nil // just to be sure
 131         }
 132         return conn
 133 }
 134
 135 func exponential(d time.Duration) time.Duration {
 136         d *= 2
 137         if d > time.Minute {
 138                 d = time.Minute
 139         }
 140         return d
 141 }
 142
 143 // ErrConnectionUnavailable is returned by the Manager's Write method when the
 144 // manager cannot yield a good connection.
 145 var ErrConnectionUnavailable = errors.New("connection unavailable")