github.com/m3db/m3@v1.5.0/src/msg/producer/writer/shard_writer_test.go (about)

     1  // Copyright (c) 2018 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 writer
    22  
    23  import (
    24  	"net"
    25  	"sync"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/m3db/m3/src/cluster/placement"
    30  	"github.com/m3db/m3/src/cluster/shard"
    31  	"github.com/m3db/m3/src/msg/generated/proto/msgpb"
    32  	"github.com/m3db/m3/src/msg/producer"
    33  	"github.com/m3db/m3/src/msg/protocol/proto"
    34  	xtest "github.com/m3db/m3/src/x/test"
    35  
    36  	"github.com/fortytw2/leaktest"
    37  	"github.com/stretchr/testify/require"
    38  )
    39  
    40  func TestSharedShardWriter(t *testing.T) {
    41  	defer leaktest.Check(t)()
    42  
    43  	a := newAckRouter(2)
    44  	opts := testOptions()
    45  	sw := newSharedShardWriter(1, a, testMessagePool(opts), opts, testMessageWriterMetrics())
    46  	defer sw.Close()
    47  
    48  	cw1 := newConsumerWriter("i1", a, opts, testConsumerWriterMetrics())
    49  	cw1.Init()
    50  	defer cw1.Close()
    51  
    52  	lis, err := net.Listen("tcp", "127.0.0.1:0")
    53  	require.NoError(t, err)
    54  	defer lis.Close()
    55  
    56  	addr2 := lis.Addr().String()
    57  	cw2 := newConsumerWriter(addr2, a, opts, testConsumerWriterMetrics())
    58  	cw2.Init()
    59  	defer cw2.Close()
    60  
    61  	cws := make(map[string]consumerWriter)
    62  	cws["i1"] = cw1
    63  	cws[addr2] = cw2
    64  
    65  	i1 := placement.NewInstance().SetEndpoint("i1")
    66  	i2 := placement.NewInstance().SetEndpoint(addr2)
    67  
    68  	var wg sync.WaitGroup
    69  	defer wg.Wait()
    70  
    71  	wg.Add(1)
    72  	go func() {
    73  		testConsumeAndAckOnConnectionListener(t, lis, opts.EncoderOptions(), opts.DecoderOptions())
    74  		wg.Done()
    75  	}()
    76  
    77  	sw.UpdateInstances(
    78  		[]placement.Instance{i1},
    79  		cws,
    80  	)
    81  
    82  	ctrl := xtest.NewController(t)
    83  	defer ctrl.Finish()
    84  
    85  	mm := producer.NewMockMessage(ctrl)
    86  	mm.EXPECT().Bytes().Return([]byte("foo"))
    87  	mm.EXPECT().Finalize(producer.Consumed)
    88  	mm.EXPECT().Size().Return(3)
    89  
    90  	sw.Write(producer.NewRefCountedMessage(mm, nil))
    91  
    92  	mw := sw.(*sharedShardWriter).mw.(*messageWriterImpl)
    93  	mw.RLock()
    94  	require.Equal(t, 1, len(mw.consumerWriters))
    95  	require.Equal(t, 1, mw.queue.Len())
    96  	mw.RUnlock()
    97  
    98  	sw.UpdateInstances(
    99  		[]placement.Instance{i1, i2},
   100  		cws,
   101  	)
   102  	mw.RLock()
   103  	require.Equal(t, 2, len(mw.consumerWriters))
   104  	mw.RUnlock()
   105  	for {
   106  		mw.RLock()
   107  		l := mw.queue.Len()
   108  		mw.RUnlock()
   109  		if l == 0 {
   110  			break
   111  		}
   112  		time.Sleep(200 * time.Millisecond)
   113  	}
   114  }
   115  
   116  func TestReplicatedShardWriter(t *testing.T) {
   117  	defer leaktest.Check(t)()
   118  
   119  	a := newAckRouter(3)
   120  	opts := testOptions()
   121  	sw := newReplicatedShardWriter(1, 200, a, testMessagePool(opts), opts, testMessageWriterMetrics()).(*replicatedShardWriter)
   122  	defer sw.Close()
   123  
   124  	lis1, err := net.Listen("tcp", "127.0.0.1:0")
   125  	require.NoError(t, err)
   126  	defer lis1.Close()
   127  
   128  	lis2, err := net.Listen("tcp", "127.0.0.1:0")
   129  	require.NoError(t, err)
   130  	defer lis2.Close()
   131  
   132  	addr1 := lis1.Addr().String()
   133  	cw1 := newConsumerWriter(addr1, a, opts, testConsumerWriterMetrics())
   134  	cw1.Init()
   135  	defer cw1.Close()
   136  
   137  	addr2 := lis2.Addr().String()
   138  	cw2 := newConsumerWriter(addr2, a, opts, testConsumerWriterMetrics())
   139  	cw2.Init()
   140  	defer cw2.Close()
   141  
   142  	cw3 := newConsumerWriter("i3", a, opts, testConsumerWriterMetrics())
   143  	cw3.Init()
   144  	defer cw3.Close()
   145  
   146  	cws := make(map[string]consumerWriter)
   147  	cws[addr1] = cw1
   148  	cws[addr2] = cw2
   149  	cws["i3"] = cw3
   150  	i1 := placement.NewInstance().
   151  		SetEndpoint(addr1).
   152  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1)}))
   153  	i2 := placement.NewInstance().
   154  		SetEndpoint(addr2).
   155  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1)}))
   156  	i3 := placement.NewInstance().
   157  		SetEndpoint("i3").
   158  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1)}))
   159  
   160  	sw.UpdateInstances(
   161  		[]placement.Instance{i1, i3},
   162  		cws,
   163  	)
   164  	require.Equal(t, 2, len(sw.messageWriters))
   165  
   166  	ctrl := xtest.NewController(t)
   167  	defer ctrl.Finish()
   168  
   169  	mm := producer.NewMockMessage(ctrl)
   170  	mm.EXPECT().Size().Return(3)
   171  	mm.EXPECT().Bytes().Return([]byte("foo")).Times(2)
   172  
   173  	sw.Write(producer.NewRefCountedMessage(mm, nil))
   174  
   175  	mw1 := sw.messageWriters[i1.Endpoint()].(*messageWriterImpl)
   176  	require.Equal(t, 1, mw1.queue.Len())
   177  	mw3 := sw.messageWriters[i3.Endpoint()].(*messageWriterImpl)
   178  	require.Equal(t, 1, mw3.queue.Len())
   179  
   180  	var wg sync.WaitGroup
   181  	defer wg.Wait()
   182  
   183  	wg.Add(1)
   184  	go func() {
   185  		testConsumeAndAckOnConnectionListener(t, lis1, opts.EncoderOptions(), opts.DecoderOptions())
   186  		wg.Done()
   187  	}()
   188  
   189  	for {
   190  		mw1.RLock()
   191  		l := mw1.queue.Len()
   192  		mw1.RUnlock()
   193  		if l == 0 {
   194  			break
   195  		}
   196  		time.Sleep(100 * time.Millisecond)
   197  	}
   198  	require.Equal(t, 1, mw3.queue.Len())
   199  
   200  	mm.EXPECT().Finalize(producer.Consumed)
   201  	sw.UpdateInstances(
   202  		[]placement.Instance{i1, i2},
   203  		cws,
   204  	)
   205  
   206  	wg.Add(1)
   207  	go func() {
   208  		testConsumeAndAckOnConnectionListener(t, lis2, opts.EncoderOptions(), opts.DecoderOptions())
   209  		wg.Done()
   210  	}()
   211  
   212  	for {
   213  		mw3.RLock()
   214  		l := mw3.queue.Len()
   215  		mw3.RUnlock()
   216  		if l == 0 {
   217  			break
   218  		}
   219  		time.Sleep(100 * time.Millisecond)
   220  	}
   221  
   222  	mw2 := sw.messageWriters[i2.Endpoint()].(*messageWriterImpl)
   223  	require.Equal(t, mw3, mw2)
   224  	_, ok := sw.messageWriters[i3.Endpoint()]
   225  	require.False(t, ok)
   226  }
   227  
   228  func TestReplicatedShardWriterRemoveMessageWriter(t *testing.T) {
   229  	defer leaktest.Check(t)()
   230  
   231  	router := newAckRouter(2).(*router)
   232  	opts := testOptions()
   233  	sw := newReplicatedShardWriter(1, 200, router, testMessagePool(opts), opts, testMessageWriterMetrics()).(*replicatedShardWriter)
   234  
   235  	lis1, err := net.Listen("tcp", "127.0.0.1:0")
   236  	require.NoError(t, err)
   237  	defer lis1.Close()
   238  
   239  	lis2, err := net.Listen("tcp", "127.0.0.1:0")
   240  	require.NoError(t, err)
   241  	defer lis2.Close()
   242  
   243  	addr1 := lis1.Addr().String()
   244  	cw1 := newConsumerWriter(addr1, router, opts, testConsumerWriterMetrics())
   245  	cw1.Init()
   246  	defer cw1.Close()
   247  
   248  	addr2 := lis2.Addr().String()
   249  	cw2 := newConsumerWriter(addr2, router, opts, testConsumerWriterMetrics())
   250  	cw2.Init()
   251  	defer cw2.Close()
   252  
   253  	cws := make(map[string]consumerWriter)
   254  	cws[addr1] = cw1
   255  	cws[addr2] = cw2
   256  	i1 := placement.NewInstance().
   257  		SetEndpoint(addr1).
   258  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1)}))
   259  	i2 := placement.NewInstance().
   260  		SetEndpoint(addr2).
   261  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1)}))
   262  
   263  	sw.UpdateInstances(
   264  		[]placement.Instance{i1, i2},
   265  		cws,
   266  	)
   267  
   268  	require.Equal(t, 2, len(sw.messageWriters))
   269  
   270  	mw1 := sw.messageWriters[i1.Endpoint()].(*messageWriterImpl)
   271  	mw2 := sw.messageWriters[i2.Endpoint()].(*messageWriterImpl)
   272  	require.Equal(t, 0, mw1.queue.Len())
   273  	require.Equal(t, 0, mw2.queue.Len())
   274  
   275  	ctrl := xtest.NewController(t)
   276  	defer ctrl.Finish()
   277  
   278  	mm := producer.NewMockMessage(ctrl)
   279  	mm.EXPECT().Size().Return(3)
   280  	mm.EXPECT().Bytes().Return([]byte("foo")).Times(2)
   281  
   282  	sw.Write(producer.NewRefCountedMessage(mm, nil))
   283  	require.Equal(t, 1, mw1.queue.Len())
   284  	require.Equal(t, 1, mw2.queue.Len())
   285  
   286  	var wg sync.WaitGroup
   287  	defer wg.Wait()
   288  
   289  	wg.Add(1)
   290  	go func() {
   291  		testConsumeAndAckOnConnectionListener(t, lis1, opts.EncoderOptions(), opts.DecoderOptions())
   292  		wg.Done()
   293  	}()
   294  
   295  	for {
   296  		mw1.RLock()
   297  		l := mw1.queue.Len()
   298  		mw1.RUnlock()
   299  		if l == 0 {
   300  			break
   301  		}
   302  		time.Sleep(100 * time.Millisecond)
   303  	}
   304  
   305  	require.Equal(t, 1, mw2.queue.Len())
   306  
   307  	conn, err := lis2.Accept()
   308  	require.NoError(t, err)
   309  	defer conn.Close()
   310  
   311  	serverEncoder := proto.NewEncoder(opts.EncoderOptions())
   312  	serverDecoder := proto.NewDecoder(conn, opts.DecoderOptions(), 10)
   313  
   314  	var msg msgpb.Message
   315  	require.NoError(t, serverDecoder.Decode(&msg))
   316  	sw.UpdateInstances(
   317  		[]placement.Instance{i1},
   318  		cws,
   319  	)
   320  
   321  	require.Equal(t, 1, len(sw.messageWriters))
   322  
   323  	mm.EXPECT().Finalize(producer.Consumed)
   324  	require.NoError(t, serverEncoder.Encode(&msgpb.Ack{Metadata: []msgpb.Metadata{msg.Metadata}}))
   325  	_, err = conn.Write(serverEncoder.Bytes())
   326  	require.NoError(t, err)
   327  	// Make sure mw2 is closed and removed from router.
   328  	for {
   329  		router.RLock()
   330  		l := len(router.messageWriters)
   331  		router.RUnlock()
   332  		if l == 1 {
   333  			break
   334  		}
   335  		time.Sleep(100 * time.Millisecond)
   336  	}
   337  	mw2.RLock()
   338  	require.Equal(t, 0, mw2.queue.Len())
   339  	mw2.RUnlock()
   340  
   341  	sw.Close()
   342  }
   343  
   344  func TestReplicatedShardWriterUpdate(t *testing.T) {
   345  	defer leaktest.Check(t)()
   346  
   347  	a := newAckRouter(4)
   348  	opts := testOptions()
   349  	sw := newReplicatedShardWriter(1, 200, a, testMessagePool(opts), opts, testMessageWriterMetrics()).(*replicatedShardWriter)
   350  	defer sw.Close()
   351  
   352  	cw1 := newConsumerWriter("i1", a, opts, testConsumerWriterMetrics())
   353  	cw2 := newConsumerWriter("i2", a, opts, testConsumerWriterMetrics())
   354  	cw3 := newConsumerWriter("i3", a, opts, testConsumerWriterMetrics())
   355  	cw4 := newConsumerWriter("i4", a, opts, testConsumerWriterMetrics())
   356  	cws := make(map[string]consumerWriter)
   357  	cws["i1"] = cw1
   358  	cws["i2"] = cw2
   359  	cws["i3"] = cw3
   360  	cws["i4"] = cw4
   361  
   362  	i1 := placement.NewInstance().
   363  		SetEndpoint("i1").
   364  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1).SetCutoffNanos(801).SetCutoverNanos(401)}))
   365  	i2 := placement.NewInstance().
   366  		SetEndpoint("i2").
   367  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1).SetCutoffNanos(802).SetCutoverNanos(402)}))
   368  	i3 := placement.NewInstance().
   369  		SetEndpoint("i3").
   370  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1).SetCutoffNanos(803).SetCutoverNanos(403)}))
   371  	i4 := placement.NewInstance().
   372  		SetEndpoint("i4").
   373  		SetShards(shard.NewShards([]shard.Shard{shard.NewShard(1).SetCutoffNanos(804).SetCutoverNanos(404)}))
   374  
   375  	sw.UpdateInstances([]placement.Instance{i1, i2}, cws)
   376  	require.Equal(t, 2, int(sw.replicaID))
   377  	require.Equal(t, 2, len(sw.messageWriters))
   378  	mw1 := sw.messageWriters[i1.Endpoint()]
   379  	require.NotNil(t, mw1)
   380  	require.Equal(t, 801, int(mw1.CutoffNanos()))
   381  	require.Equal(t, 401, int(mw1.CutoverNanos()))
   382  	require.NotNil(t, sw.messageWriters[i2.Endpoint()])
   383  	require.Equal(t, 0, int(mw1.MessageTTLNanos()))
   384  
   385  	sw.SetMessageTTLNanos(500)
   386  	require.Equal(t, 500, int(mw1.MessageTTLNanos()))
   387  
   388  	sw.UpdateInstances([]placement.Instance{i2, i3}, cws)
   389  	require.Equal(t, 2, int(sw.replicaID))
   390  	require.Equal(t, 2, len(sw.messageWriters))
   391  	mw2 := sw.messageWriters[i2.Endpoint()]
   392  	require.NotNil(t, mw2)
   393  	mw3 := sw.messageWriters[i3.Endpoint()]
   394  	require.NotNil(t, mw3)
   395  	require.Equal(t, mw1, mw3)
   396  	require.Equal(t, 803, int(mw3.CutoffNanos()))
   397  	require.Equal(t, 403, int(mw3.CutoverNanos()))
   398  	m := make(map[uint64]int, 2)
   399  	m[mw2.ReplicatedShardID()] = 1
   400  	m[mw3.ReplicatedShardID()] = 1
   401  	require.Equal(t, map[uint64]int{1: 1, 201: 1}, m)
   402  	require.Equal(t, 500, int(mw2.MessageTTLNanos()))
   403  	require.Equal(t, 500, int(mw3.MessageTTLNanos()))
   404  
   405  	sw.UpdateInstances([]placement.Instance{i3}, cws)
   406  	require.Equal(t, 2, int(sw.replicaID))
   407  	require.Equal(t, 1, len(sw.messageWriters))
   408  	require.NotNil(t, sw.messageWriters[i3.Endpoint()])
   409  	require.Equal(t, 500, int(mw3.MessageTTLNanos()))
   410  	for {
   411  		mw2.(*messageWriterImpl).RLock()
   412  		isClosed := mw2.(*messageWriterImpl).isClosed
   413  		mw2.(*messageWriterImpl).RUnlock()
   414  		if isClosed {
   415  			break
   416  		}
   417  		time.Sleep(100 * time.Millisecond)
   418  	}
   419  
   420  	sw.SetMessageTTLNanos(800)
   421  	require.Equal(t, 800, int(mw3.MessageTTLNanos()))
   422  	sw.UpdateInstances([]placement.Instance{i1, i2, i3}, cws)
   423  	require.Equal(t, 4, int(sw.replicaID))
   424  	require.Equal(t, 3, len(sw.messageWriters))
   425  	newmw1 := sw.messageWriters[i1.Endpoint()]
   426  	require.NotNil(t, newmw1)
   427  	require.NotEqual(t, &mw1, &newmw1)
   428  	newmw2 := sw.messageWriters[i2.Endpoint()]
   429  	require.NotNil(t, newmw2)
   430  	require.NotEqual(t, &mw2, &newmw2)
   431  	newmw3 := sw.messageWriters[i3.Endpoint()]
   432  	require.NotNil(t, newmw3)
   433  	require.Equal(t, &mw3, &newmw3)
   434  	m = make(map[uint64]int, 3)
   435  	m[newmw1.ReplicatedShardID()] = 1
   436  	m[newmw2.ReplicatedShardID()] = 1
   437  	m[newmw3.ReplicatedShardID()] = 1
   438  	require.Equal(t, map[uint64]int{601: 1, 401: 1, mw3.ReplicatedShardID(): 1}, m)
   439  	require.Equal(t, 800, int(newmw1.MessageTTLNanos()))
   440  	require.Equal(t, 800, int(newmw2.MessageTTLNanos()))
   441  	require.Equal(t, 800, int(newmw3.MessageTTLNanos()))
   442  
   443  	sw.UpdateInstances([]placement.Instance{i2, i4}, cws)
   444  	require.Equal(t, 4, int(sw.replicaID))
   445  	require.Equal(t, 2, len(sw.messageWriters))
   446  	require.NotNil(t, sw.messageWriters[i2.Endpoint()])
   447  	require.NotNil(t, sw.messageWriters[i4.Endpoint()])
   448  
   449  	sw.UpdateInstances([]placement.Instance{i1}, cws)
   450  	require.Equal(t, 4, int(sw.replicaID))
   451  	require.Equal(t, 1, len(sw.messageWriters))
   452  	require.NotNil(t, sw.messageWriters[i1.Endpoint()])
   453  }