github.com/m3db/m3@v1.5.0/src/x/server/server_test.go (about)

     1  // Copyright (c) 2017 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package server
    22  
    23  import (
    24  	"fmt"
    25  	"net"
    26  	"sort"
    27  	"sync"
    28  	"sync/atomic"
    29  	"testing"
    30  	"time"
    31  
    32  	"github.com/m3db/m3/src/x/retry"
    33  
    34  	"github.com/stretchr/testify/require"
    35  )
    36  
    37  const (
    38  	testListenAddress = "127.0.0.1:0"
    39  )
    40  
    41  // nolint: unparam
    42  func testServer(addr string) (*server, *mockHandler, *int32, *int32) {
    43  	var (
    44  		numAdded   int32
    45  		numRemoved int32
    46  	)
    47  
    48  	opts := NewOptions().SetRetryOptions(retry.NewOptions().SetMaxRetries(2))
    49  	opts = opts.SetInstrumentOptions(opts.InstrumentOptions().SetReportInterval(time.Second))
    50  
    51  	h := newMockHandler()
    52  	s := NewServer(addr, h, opts).(*server)
    53  
    54  	s.addConnectionFn = func(conn net.Conn) bool {
    55  		atomic.AddInt32(&numAdded, 1)
    56  		ret := s.addConnection(conn)
    57  		return ret
    58  	}
    59  
    60  	s.removeConnectionFn = func(conn net.Conn) {
    61  		atomic.AddInt32(&numRemoved, 1)
    62  		s.removeConnection(conn)
    63  	}
    64  
    65  	return s, h, &numAdded, &numRemoved
    66  }
    67  
    68  func TestServerListenAndClose(t *testing.T) {
    69  	s, h, numAdded, numRemoved := testServer(testListenAddress)
    70  
    71  	var (
    72  		numClients  = 9
    73  		expectedRes []string
    74  	)
    75  
    76  	err := s.ListenAndServe()
    77  	require.NoError(t, err)
    78  	listenAddr := s.listener.Addr().String()
    79  
    80  	for i := 0; i < numClients; i++ {
    81  		conn, err := net.Dial("tcp", listenAddr)
    82  		require.NoError(t, err)
    83  
    84  		msg := fmt.Sprintf("msg%d", i)
    85  		expectedRes = append(expectedRes, msg)
    86  
    87  		_, err = conn.Write([]byte(msg))
    88  		require.NoError(t, err)
    89  	}
    90  
    91  	for h.called() < numClients {
    92  		time.Sleep(100 * time.Millisecond)
    93  	}
    94  
    95  	require.False(t, h.isClosed())
    96  
    97  	s.Close()
    98  
    99  	require.True(t, h.isClosed())
   100  	require.Equal(t, int32(numClients), atomic.LoadInt32(numAdded))
   101  	require.Equal(t, int32(numClients), atomic.LoadInt32(numRemoved))
   102  	require.Equal(t, numClients, h.called())
   103  	require.Equal(t, expectedRes, h.res())
   104  }
   105  
   106  func TestServe(t *testing.T) {
   107  	s, _, _, _ := testServer(testListenAddress)
   108  
   109  	l, err := net.Listen("tcp", testListenAddress)
   110  	require.NoError(t, err)
   111  
   112  	err = s.Serve(l)
   113  	require.NoError(t, err)
   114  	require.Equal(t, l, s.listener)
   115  	require.Equal(t, l.Addr().String(), s.address)
   116  
   117  	s.Close()
   118  }
   119  
   120  type mockHandler struct {
   121  	sync.Mutex
   122  
   123  	n        int
   124  	closed   bool
   125  	received []string
   126  }
   127  
   128  func newMockHandler() *mockHandler { return &mockHandler{} }
   129  
   130  func (h *mockHandler) Handle(conn net.Conn) {
   131  	h.Lock()
   132  	b := make([]byte, 16)
   133  
   134  	n, _ := conn.Read(b)
   135  	h.n++
   136  	h.received = append(h.received, string(b[:n]))
   137  	h.Unlock()
   138  }
   139  
   140  func (h *mockHandler) Close() {
   141  	h.Lock()
   142  	h.closed = true
   143  	h.Unlock()
   144  }
   145  
   146  func (h *mockHandler) isClosed() bool {
   147  	h.Lock()
   148  	defer h.Unlock()
   149  
   150  	return h.closed
   151  }
   152  
   153  func (h *mockHandler) called() int {
   154  	h.Lock()
   155  	defer h.Unlock()
   156  
   157  	return h.n
   158  }
   159  
   160  func (h *mockHandler) res() []string {
   161  	h.Lock()
   162  	defer h.Unlock()
   163  
   164  	sort.Strings(h.received)
   165  	return h.received
   166  }