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

   1 package conn
   2
   3 import (
   4         "errors"
   5         "net"
   6         "sync/atomic"
   7         "testing"
   8         "time"
   9
  10         "github.com/go-kit/kit/log"
  11 )
  12
  13 func TestManager(t *testing.T) {
  14         var (
  15                 tickc    = make(chan time.Time)
  16                 after    = func(time.Duration) <-chan time.Time { return tickc }
  17                 dialconn = &mockConn{}
  18                 dialerr  = error(nil)
  19                 dialer   = func(string, string) (net.Conn, error) { return dialconn, dialerr }
  20                 mgr      = NewManager(dialer, "netw", "addr", after, log.NewNopLogger())
  21         )
  22
  23         // First conn should be fine.
  24         conn := mgr.Take()
  25         if conn == nil {
  26                 t.Fatal("nil conn")
  27         }
  28
  29         // Write and check it went through.
  30         if _, err := conn.Write([]byte{1, 2, 3}); err != nil {
  31                 t.Fatal(err)
  32         }
  33         if want, have := uint64(3), atomic.LoadUint64(&dialconn.wr); want != have {
  34                 t.Errorf("want %d, have %d", want, have)
  35         }
  36
  37         // Put an error to kill the conn.
  38         mgr.Put(errors.New("should kill the connection"))
  39
  40         // First takes should fail.
  41         for i := 0; i < 10; i++ {
  42                 if conn = mgr.Take(); conn != nil {
  43                         t.Fatalf("iteration %d: want nil conn, got real conn", i)
  44                 }
  45         }
  46
  47         // Trigger the reconnect.
  48         tickc <- time.Now()
  49
  50         // The dial should eventually succeed and yield a good conn.
  51         if !within(100*time.Millisecond, func() bool {
  52                 conn = mgr.Take()
  53                 return conn != nil
  54         }) {
  55                 t.Fatal("conn remained nil")
  56         }
  57
  58         // Write and check it went through.
  59         if _, err := conn.Write([]byte{4, 5}); err != nil {
  60                 t.Fatal(err)
  61         }
  62         if want, have := uint64(5), atomic.LoadUint64(&dialconn.wr); want != have {
  63                 t.Errorf("want %d, have %d", want, have)
  64         }
  65
  66         // Dial starts failing.
  67         dialconn, dialerr = nil, errors.New("oh noes")
  68         mgr.Put(errors.New("trigger that reconnect y'all"))
  69         if conn = mgr.Take(); conn != nil {
  70                 t.Fatalf("want nil conn, got real conn")
  71         }
  72
  73         // As many reconnects as they want.
  74         go func() {
  75                 done := time.After(100 * time.Millisecond)
  76                 for {
  77                         select {
  78                         case tickc <- time.Now():
  79                         case <-done:
  80                                 return
  81                         }
  82                 }
  83         }()
  84
  85         // The dial should never succeed.
  86         if within(100*time.Millisecond, func() bool {
  87                 conn = mgr.Take()
  88                 return conn != nil
  89         }) {
  90                 t.Fatal("eventually got a good conn, despite failing dialer")
  91         }
  92 }
  93
  94 func TestIssue292(t *testing.T) {
  95         // The util/conn.Manager won't attempt to reconnect to the provided endpoint
  96         // if the endpoint is initially unavailable (e.g. dial tcp :8080:
  97         // getsockopt: connection refused). If the endpoint is up when
  98         // conn.NewManager is called and then goes down/up, it reconnects just fine.
  99
 100         var (
 101                 tickc    = make(chan time.Time)
 102                 after    = func(time.Duration) <-chan time.Time { return tickc }
 103                 dialconn = net.Conn(nil)
 104                 dialerr  = errors.New("fail")
 105                 dialer   = func(string, string) (net.Conn, error) { return dialconn, dialerr }
 106                 mgr      = NewManager(dialer, "netw", "addr", after, log.NewNopLogger())
 107         )
 108
 109         if conn := mgr.Take(); conn != nil {
 110                 t.Fatal("first Take should have yielded nil conn, but didn't")
 111         }
 112
 113         dialconn, dialerr = &mockConn{}, nil
 114         select {
 115         case tickc <- time.Now():
 116         case <-time.After(time.Second):
 117                 t.Fatal("manager isn't listening for a tick, despite a failed dial")
 118         }
 119
 120         if !within(time.Second, func() bool {
 121                 return mgr.Take() != nil
 122         }) {
 123                 t.Fatal("second Take should have yielded good conn, but didn't")
 124         }
 125 }
 126
 127 type mockConn struct {
 128         rd, wr uint64
 129 }
 130
 131 func (c *mockConn) Read(b []byte) (n int, err error) {
 132         atomic.AddUint64(&c.rd, uint64(len(b)))
 133         return len(b), nil
 134 }
 135
 136 func (c *mockConn) Write(b []byte) (n int, err error) {
 137         atomic.AddUint64(&c.wr, uint64(len(b)))
 138         return len(b), nil
 139 }
 140
 141 func (c *mockConn) Close() error                       { return nil }
 142 func (c *mockConn) LocalAddr() net.Addr                { return nil }
 143 func (c *mockConn) RemoteAddr() net.Addr               { return nil }
 144 func (c *mockConn) SetDeadline(t time.Time) error      { return nil }
 145 func (c *mockConn) SetReadDeadline(t time.Time) error  { return nil }
 146 func (c *mockConn) SetWriteDeadline(t time.Time) error { return nil }
 147
 148 func within(d time.Duration, f func() bool) bool {
 149         deadline := time.Now().Add(d)
 150         for {
 151                 if time.Now().After(deadline) {
 152                         return false
 153                 }
 154                 if f() {
 155                         return true
 156                 }
 157                 time.Sleep(d / 10)
 158         }
 159 }