manager_test.go

Documentation: github.com/go-kit/kit/util/conn

     1  package conn
     2  
     3  import (
     4  	"errors"
     5  	"net"
     6  	"sync/atomic"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/go-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  }
   160
View as plain text