github.com/m3db/m3@v1.5.0/src/msg/producer/writer/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  	"errors"
    25  	"net"
    26  	"sync"
    27  	"testing"
    28  	"time"
    29  
    30  	"github.com/m3db/m3/src/cluster/client"
    31  	"github.com/m3db/m3/src/cluster/kv"
    32  	"github.com/m3db/m3/src/cluster/kv/mem"
    33  	"github.com/m3db/m3/src/cluster/placement"
    34  	"github.com/m3db/m3/src/cluster/services"
    35  	"github.com/m3db/m3/src/cluster/shard"
    36  	"github.com/m3db/m3/src/msg/producer"
    37  	"github.com/m3db/m3/src/msg/topic"
    38  	xtest "github.com/m3db/m3/src/x/test"
    39  
    40  	"github.com/fortytw2/leaktest"
    41  	"github.com/golang/mock/gomock"
    42  	"github.com/stretchr/testify/require"
    43  	"go.uber.org/atomic"
    44  )
    45  
    46  func TestWriterInitErrorNoTopic(t *testing.T) {
    47  	defer leaktest.Check(t)()
    48  
    49  	ctrl := xtest.NewController(t)
    50  	defer ctrl.Finish()
    51  
    52  	store := mem.NewStore()
    53  	cs := client.NewMockClient(ctrl)
    54  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
    55  
    56  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
    57  	require.NoError(t, err)
    58  
    59  	opts := testOptions().SetTopicService(ts)
    60  	w := NewWriter(opts)
    61  	require.Error(t, w.Init())
    62  	w.Close()
    63  }
    64  
    65  func TestWriterWriteAfterClosed(t *testing.T) {
    66  	defer leaktest.Check(t)()
    67  
    68  	ctrl := xtest.NewController(t)
    69  	defer ctrl.Finish()
    70  
    71  	store := mem.NewStore()
    72  	cs := client.NewMockClient(ctrl)
    73  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
    74  
    75  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
    76  	require.NoError(t, err)
    77  
    78  	opts := testOptions().SetTopicService(ts)
    79  	w := NewWriter(opts)
    80  	w.Init()
    81  	w.Close()
    82  
    83  	mm := producer.NewMockMessage(ctrl)
    84  	mm.EXPECT().Finalize(producer.Dropped)
    85  	mm.EXPECT().Size().Return(3)
    86  	rm := producer.NewRefCountedMessage(mm, nil)
    87  	err = w.Write(rm)
    88  	require.Error(t, err)
    89  	require.Equal(t, errWriterClosed, err)
    90  }
    91  
    92  func TestWriterWriteWithInvalidShard(t *testing.T) {
    93  	defer leaktest.Check(t)()
    94  
    95  	ctrl := xtest.NewController(t)
    96  	defer ctrl.Finish()
    97  
    98  	store := mem.NewStore()
    99  	cs := client.NewMockClient(ctrl)
   100  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   101  
   102  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   103  	require.NoError(t, err)
   104  
   105  	opts := testOptions().SetTopicService(ts)
   106  	w := NewWriter(opts).(*writer)
   107  	w.numShards = 2
   108  
   109  	mm := producer.NewMockMessage(ctrl)
   110  	mm.EXPECT().Shard().Return(uint32(2))
   111  	mm.EXPECT().Finalize(producer.Dropped)
   112  	mm.EXPECT().Size().Return(3)
   113  	rm := producer.NewRefCountedMessage(mm, nil)
   114  	err = w.Write(rm)
   115  	require.Error(t, err)
   116  
   117  	mm.EXPECT().Shard().Return(uint32(100))
   118  	mm.EXPECT().Finalize(producer.Dropped)
   119  	mm.EXPECT().Size().Return(3)
   120  	rm = producer.NewRefCountedMessage(mm, nil)
   121  	err = w.Write(rm)
   122  	require.Error(t, err)
   123  }
   124  
   125  func TestWriterInvalidTopicUpdate(t *testing.T) {
   126  	defer leaktest.Check(t)()
   127  
   128  	ctrl := xtest.NewController(t)
   129  	defer ctrl.Finish()
   130  
   131  	store := mem.NewStore()
   132  	cs := client.NewMockClient(ctrl)
   133  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   134  
   135  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   136  	require.NoError(t, err)
   137  
   138  	opts := testOptions().SetTopicService(ts)
   139  	sid1 := services.NewServiceID().SetName("s1")
   140  	cs1 := topic.NewConsumerService().SetConsumptionType(topic.Replicated).SetServiceID(sid1)
   141  	testTopic := topic.NewTopic().
   142  		SetName(opts.TopicName()).
   143  		SetNumberOfShards(2).
   144  		SetConsumerServices([]topic.ConsumerService{cs1})
   145  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   146  	require.NoError(t, err)
   147  
   148  	sd := services.NewMockServices(ctrl)
   149  	opts = opts.SetServiceDiscovery(sd)
   150  	ps1 := testPlacementService(store, sid1)
   151  	sd.EXPECT().PlacementService(sid1, gomock.Any()).Return(ps1, nil)
   152  
   153  	p1 := placement.NewPlacement().
   154  		SetInstances([]placement.Instance{
   155  			placement.NewInstance().
   156  				SetID("i1").
   157  				SetEndpoint("addr1").
   158  				SetShards(shard.NewShards([]shard.Shard{
   159  					shard.NewShard(0).SetState(shard.Available),
   160  					shard.NewShard(1).SetState(shard.Available),
   161  				})),
   162  		}).
   163  		SetShards([]uint32{0, 1}).
   164  		SetReplicaFactor(1).
   165  		SetIsSharded(true)
   166  	_, err = ps1.Set(p1)
   167  	require.NoError(t, err)
   168  
   169  	w := NewWriter(opts).(*writer)
   170  	var wg sync.WaitGroup
   171  	w.processFn = func(i interface{}) error {
   172  		defer wg.Done()
   173  		return w.process(i)
   174  	}
   175  	wg.Add(1)
   176  	require.NoError(t, w.Init())
   177  	wg.Wait()
   178  	defer w.Close()
   179  
   180  	require.Equal(t, 2, int(w.numShards))
   181  	testTopic = topic.NewTopic().
   182  		SetName(opts.TopicName()).
   183  		SetNumberOfShards(3).
   184  		SetConsumerServices([]topic.ConsumerService{cs1}).
   185  		SetVersion(1)
   186  	wg.Add(1)
   187  	_, err = ts.CheckAndSet(testTopic, 1)
   188  	require.NoError(t, err)
   189  	wg.Wait()
   190  
   191  	require.Equal(t, 2, int(w.numShards))
   192  }
   193  
   194  func TestWriterRegisterFilter(t *testing.T) {
   195  	defer leaktest.Check(t)()
   196  
   197  	ctrl := xtest.NewController(t)
   198  	defer ctrl.Finish()
   199  
   200  	store := mem.NewStore()
   201  	cs := client.NewMockClient(ctrl)
   202  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   203  
   204  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   205  	require.NoError(t, err)
   206  
   207  	opts := testOptions().SetTopicService(ts)
   208  
   209  	sid1 := services.NewServiceID().SetName("s1")
   210  	cs1 := topic.NewConsumerService().SetConsumptionType(topic.Replicated).SetServiceID(sid1)
   211  	csw1 := NewMockconsumerServiceWriter(ctrl)
   212  
   213  	sid2 := services.NewServiceID().SetName("s2")
   214  	filter := func(producer.Message) bool { return false }
   215  
   216  	w := NewWriter(opts).(*writer)
   217  	w.consumerServiceWriters[cs1.ServiceID().String()] = csw1
   218  
   219  	csw1.EXPECT().UnregisterFilter()
   220  	w.UnregisterFilter(sid1)
   221  
   222  	// Wrong service id triggers nothing.
   223  	w.RegisterFilter(sid2, filter)
   224  
   225  	csw1.EXPECT().RegisterFilter(gomock.Any())
   226  	w.RegisterFilter(sid1, filter)
   227  
   228  	csw1.EXPECT().UnregisterFilter()
   229  	w.UnregisterFilter(sid1)
   230  
   231  	csw1.EXPECT().RegisterFilter(gomock.Any())
   232  	w.RegisterFilter(sid1, filter)
   233  
   234  	csw1.EXPECT().RegisterFilter(gomock.Any())
   235  	csw1.EXPECT().SetMessageTTLNanos(int64(0))
   236  	testTopic := topic.NewTopic().
   237  		SetName(opts.TopicName()).
   238  		SetNumberOfShards(6).
   239  		SetConsumerServices([]topic.ConsumerService{cs1})
   240  	w.process(testTopic)
   241  }
   242  
   243  func TestWriterTopicUpdate(t *testing.T) {
   244  	defer leaktest.Check(t)()
   245  
   246  	ctrl := xtest.NewController(t)
   247  	defer ctrl.Finish()
   248  
   249  	store := mem.NewStore()
   250  	cs := client.NewMockClient(ctrl)
   251  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   252  
   253  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   254  	require.NoError(t, err)
   255  
   256  	opts := testOptions().SetTopicService(ts)
   257  
   258  	sid1 := services.NewServiceID().SetName("s1")
   259  	cs1 := topic.NewConsumerService().SetConsumptionType(topic.Replicated).SetServiceID(sid1)
   260  	testTopic := topic.NewTopic().
   261  		SetName(opts.TopicName()).
   262  		SetNumberOfShards(2).
   263  		SetConsumerServices([]topic.ConsumerService{cs1})
   264  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   265  	require.NoError(t, err)
   266  
   267  	sd := services.NewMockServices(ctrl)
   268  	opts = opts.SetServiceDiscovery(sd)
   269  	ps1 := testPlacementService(store, sid1)
   270  	sd.EXPECT().PlacementService(sid1, gomock.Any()).Return(ps1, nil)
   271  
   272  	p1 := placement.NewPlacement().
   273  		SetInstances([]placement.Instance{
   274  			placement.NewInstance().
   275  				SetID("i1").
   276  				SetEndpoint("addr1").
   277  				SetShards(shard.NewShards([]shard.Shard{
   278  					shard.NewShard(0).SetState(shard.Available),
   279  					shard.NewShard(1).SetState(shard.Available),
   280  				})),
   281  		}).
   282  		SetShards([]uint32{0, 1}).
   283  		SetReplicaFactor(1).
   284  		SetIsSharded(true)
   285  	_, err = ps1.Set(p1)
   286  	require.NoError(t, err)
   287  
   288  	w := NewWriter(opts).(*writer)
   289  	require.NoError(t, w.Init())
   290  	defer w.Close()
   291  
   292  	require.Equal(t, 1, len(w.consumerServiceWriters))
   293  
   294  	sid2 := services.NewServiceID().SetName("s2")
   295  	cs2 := topic.NewConsumerService().SetConsumptionType(topic.Shared).SetServiceID(sid2)
   296  	ps2 := testPlacementService(store, sid2)
   297  	sd.EXPECT().PlacementService(sid2, gomock.Any()).Return(ps2, nil)
   298  
   299  	sid3 := services.NewServiceID().SetName("s3")
   300  	cs3 := topic.NewConsumerService().SetConsumptionType(topic.Replicated).SetServiceID(sid3)
   301  	sd.EXPECT().PlacementService(sid3, gomock.Any()).Return(nil, errors.New("test error"))
   302  
   303  	sid4 := services.NewServiceID().SetName("s4")
   304  	cs4 := topic.NewConsumerService().SetConsumptionType(topic.Shared).SetServiceID(sid4)
   305  	ps4 := placement.NewMockService(ctrl)
   306  	sd.EXPECT().PlacementService(sid4, gomock.Any()).Return(ps4, nil)
   307  	ps4.EXPECT().Watch().Return(nil, errors.New("watch error"))
   308  
   309  	testTopic = testTopic.
   310  		SetConsumerServices([]topic.ConsumerService{
   311  			cs1, cs2,
   312  			cs3, // Could not create consumer service write for cs3.
   313  			cs4, // Could not init cs4.
   314  		}).
   315  		SetVersion(1)
   316  	_, err = ts.CheckAndSet(testTopic, 1)
   317  	require.NoError(t, err)
   318  
   319  	for {
   320  		w.RLock()
   321  		l := len(w.consumerServiceWriters)
   322  		w.RUnlock()
   323  		if l == 2 {
   324  			break
   325  		}
   326  		time.Sleep(100 * time.Millisecond)
   327  	}
   328  
   329  	store.Delete(opts.TopicName())
   330  	time.Sleep(100 * time.Millisecond)
   331  	require.Equal(t, 2, len(w.consumerServiceWriters))
   332  
   333  	testTopic = testTopic.
   334  		SetConsumerServices([]topic.ConsumerService{cs2}).
   335  		SetVersion(0)
   336  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   337  	require.NoError(t, err)
   338  
   339  	for {
   340  		w.RLock()
   341  		l := len(w.consumerServiceWriters)
   342  		w.RUnlock()
   343  		if l == 1 {
   344  			break
   345  		}
   346  		time.Sleep(100 * time.Millisecond)
   347  	}
   348  	w.Close()
   349  
   350  	testTopic = testTopic.
   351  		SetConsumerServices([]topic.ConsumerService{cs1, cs2}).
   352  		SetVersion(1)
   353  	_, err = ts.CheckAndSet(testTopic, 1)
   354  	require.NoError(t, err)
   355  
   356  	// Not going to process topic update anymore.
   357  	time.Sleep(100 * time.Millisecond)
   358  	require.Equal(t, 1, len(w.consumerServiceWriters))
   359  }
   360  
   361  func TestTopicUpdateWithSameConsumerServicesButDifferentOrder(t *testing.T) {
   362  	defer leaktest.Check(t)()
   363  
   364  	ctrl := xtest.NewController(t)
   365  	defer ctrl.Finish()
   366  
   367  	store := mem.NewStore()
   368  	cs := client.NewMockClient(ctrl)
   369  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   370  
   371  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   372  	require.NoError(t, err)
   373  
   374  	opts := testOptions().SetTopicService(ts)
   375  
   376  	sid1 := services.NewServiceID().SetName("s1")
   377  	cs1 := topic.NewConsumerService().SetConsumptionType(topic.Replicated).SetServiceID(sid1)
   378  	sid2 := services.NewServiceID().SetName("s2")
   379  	cs2 := topic.NewConsumerService().SetConsumptionType(topic.Shared).SetServiceID(sid2).SetMessageTTLNanos(500)
   380  	testTopic := topic.NewTopic().
   381  		SetName(opts.TopicName()).
   382  		SetNumberOfShards(1).
   383  		SetConsumerServices([]topic.ConsumerService{cs1, cs2})
   384  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   385  	require.NoError(t, err)
   386  
   387  	sd := services.NewMockServices(ctrl)
   388  	opts = opts.SetServiceDiscovery(sd)
   389  	ps1 := testPlacementService(store, sid1)
   390  	sd.EXPECT().PlacementService(sid1, gomock.Any()).Return(ps1, nil)
   391  	ps2 := testPlacementService(store, sid2)
   392  	sd.EXPECT().PlacementService(sid2, gomock.Any()).Return(ps2, nil)
   393  
   394  	p1 := placement.NewPlacement().
   395  		SetInstances([]placement.Instance{
   396  			placement.NewInstance().
   397  				SetID("i1").
   398  				SetEndpoint("i1").
   399  				SetShards(shard.NewShards([]shard.Shard{
   400  					shard.NewShard(0).SetState(shard.Available),
   401  				})),
   402  		}).
   403  		SetShards([]uint32{0}).
   404  		SetReplicaFactor(1).
   405  		SetIsSharded(true)
   406  	_, err = ps1.Set(p1)
   407  	require.NoError(t, err)
   408  
   409  	p2 := placement.NewPlacement().
   410  		SetInstances([]placement.Instance{
   411  			placement.NewInstance().
   412  				SetID("i2").
   413  				SetEndpoint("i2").
   414  				SetShards(shard.NewShards([]shard.Shard{
   415  					shard.NewShard(0).SetState(shard.Available),
   416  				})),
   417  		}).
   418  		SetShards([]uint32{0}).
   419  		SetReplicaFactor(1).
   420  		SetIsSharded(true)
   421  	_, err = ps2.Set(p2)
   422  	require.NoError(t, err)
   423  
   424  	w := NewWriter(opts).(*writer)
   425  
   426  	called := atomic.NewInt32(0)
   427  	w.processFn = func(update interface{}) error {
   428  		called.Inc()
   429  		return w.process(update)
   430  	}
   431  	require.NoError(t, w.Init())
   432  	require.Equal(t, 1, int(called.Load()))
   433  	require.Equal(t, 2, len(w.consumerServiceWriters))
   434  	csw, ok := w.consumerServiceWriters[cs1.ServiceID().String()]
   435  	require.True(t, ok)
   436  	cswMock1 := NewMockconsumerServiceWriter(ctrl)
   437  	w.consumerServiceWriters[cs1.ServiceID().String()] = cswMock1
   438  	defer csw.Close()
   439  
   440  	csw, ok = w.consumerServiceWriters[cs2.ServiceID().String()]
   441  	require.True(t, ok)
   442  	cswMock2 := NewMockconsumerServiceWriter(ctrl)
   443  	w.consumerServiceWriters[cs2.ServiceID().String()] = cswMock2
   444  	defer csw.Close()
   445  
   446  	cswMock1.EXPECT().SetMessageTTLNanos(int64(0))
   447  	cswMock2.EXPECT().SetMessageTTLNanos(int64(500))
   448  	testTopic = testTopic.
   449  		SetConsumerServices([]topic.ConsumerService{cs2, cs1}).
   450  		SetVersion(1)
   451  	_, err = ts.CheckAndSet(testTopic, 1)
   452  	require.NoError(t, err)
   453  
   454  	// The update will be processed, but nothing will be called on any of the mock writers.
   455  	for called.Load() != 2 {
   456  		time.Sleep(50 * time.Millisecond)
   457  	}
   458  	cswMock1.EXPECT().Close()
   459  	cswMock2.EXPECT().Close()
   460  	w.Close()
   461  }
   462  
   463  func TestWriterWrite(t *testing.T) {
   464  	defer leaktest.Check(t)()
   465  
   466  	ctrl := xtest.NewController(t)
   467  	defer ctrl.Finish()
   468  
   469  	store := mem.NewStore()
   470  	cs := client.NewMockClient(ctrl)
   471  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   472  
   473  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   474  	require.NoError(t, err)
   475  
   476  	opts := testOptions().SetTopicService(ts)
   477  
   478  	sid1 := services.NewServiceID().SetName("s1")
   479  	cs1 := topic.NewConsumerService().SetConsumptionType(topic.Replicated).SetServiceID(sid1)
   480  	sid2 := services.NewServiceID().SetName("s2")
   481  	cs2 := topic.NewConsumerService().SetConsumptionType(topic.Shared).SetServiceID(sid2)
   482  	testTopic := topic.NewTopic().
   483  		SetName(opts.TopicName()).
   484  		SetNumberOfShards(1).
   485  		SetConsumerServices([]topic.ConsumerService{cs1, cs2})
   486  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   487  	require.NoError(t, err)
   488  
   489  	sd := services.NewMockServices(ctrl)
   490  	opts = opts.SetServiceDiscovery(sd)
   491  	ps1 := testPlacementService(store, sid1)
   492  	sd.EXPECT().PlacementService(sid1, gomock.Any()).Return(ps1, nil)
   493  	ps2 := testPlacementService(store, sid2)
   494  	sd.EXPECT().PlacementService(sid2, gomock.Any()).Return(ps2, nil)
   495  
   496  	lis1, err := net.Listen("tcp", "127.0.0.1:0")
   497  	require.NoError(t, err)
   498  	defer lis1.Close()
   499  
   500  	lis2, err := net.Listen("tcp", "127.0.0.1:0")
   501  	require.NoError(t, err)
   502  	defer lis2.Close()
   503  
   504  	lis3, err := net.Listen("tcp", "127.0.0.1:0")
   505  	require.NoError(t, err)
   506  	defer lis3.Close()
   507  
   508  	p1 := placement.NewPlacement().
   509  		SetInstances([]placement.Instance{
   510  			placement.NewInstance().
   511  				SetID("i1").
   512  				SetEndpoint(lis1.Addr().String()).
   513  				SetShards(shard.NewShards([]shard.Shard{
   514  					shard.NewShard(0).SetState(shard.Available),
   515  				})),
   516  			placement.NewInstance().
   517  				SetID("i2").
   518  				SetEndpoint(lis2.Addr().String()).
   519  				SetShards(shard.NewShards([]shard.Shard{
   520  					shard.NewShard(0).SetState(shard.Available),
   521  				})),
   522  		}).
   523  		SetShards([]uint32{0}).
   524  		SetReplicaFactor(2).
   525  		SetIsSharded(true)
   526  	_, err = ps1.Set(p1)
   527  	require.NoError(t, err)
   528  
   529  	p2 := placement.NewPlacement().
   530  		SetInstances([]placement.Instance{
   531  			placement.NewInstance().
   532  				SetID("i13").
   533  				SetEndpoint(lis3.Addr().String()).
   534  				SetShards(shard.NewShards([]shard.Shard{
   535  					shard.NewShard(0).SetState(shard.Available),
   536  				})),
   537  			placement.NewInstance().
   538  				SetID("i4").
   539  				SetEndpoint("addr4").
   540  				SetShards(shard.NewShards([]shard.Shard{
   541  					shard.NewShard(0).SetState(shard.Available),
   542  				})),
   543  		}).
   544  		SetShards([]uint32{0}).
   545  		SetReplicaFactor(2).
   546  		SetIsSharded(true)
   547  	_, err = ps2.Set(p2)
   548  	require.NoError(t, err)
   549  
   550  	w := NewWriter(opts).(*writer)
   551  	require.NoError(t, w.Init())
   552  	defer w.Close()
   553  
   554  	require.Equal(t, 2, len(w.consumerServiceWriters))
   555  
   556  	var wg sync.WaitGroup
   557  	mm := producer.NewMockMessage(ctrl)
   558  	mm.EXPECT().Shard().Return(uint32(0)).Times(3)
   559  	mm.EXPECT().Size().Return(3)
   560  	mm.EXPECT().Bytes().Return([]byte("foo")).Times(3)
   561  	mm.EXPECT().Finalize(producer.Consumed).Do(func(interface{}) { wg.Done() })
   562  	rm := producer.NewRefCountedMessage(mm, nil)
   563  	wg.Add(1)
   564  	require.NoError(t, w.Write(rm))
   565  
   566  	wg.Add(1)
   567  	go func() {
   568  		testConsumeAndAckOnConnectionListener(t, lis1, opts.EncoderOptions(), opts.DecoderOptions())
   569  		wg.Done()
   570  	}()
   571  
   572  	wg.Add(1)
   573  	go func() {
   574  		testConsumeAndAckOnConnectionListener(t, lis2, opts.EncoderOptions(), opts.DecoderOptions())
   575  		wg.Done()
   576  	}()
   577  
   578  	wg.Add(1)
   579  	go func() {
   580  		testConsumeAndAckOnConnectionListener(t, lis3, opts.EncoderOptions(), opts.DecoderOptions())
   581  		wg.Done()
   582  	}()
   583  
   584  	wg.Wait()
   585  	w.Close()
   586  }
   587  
   588  func TestWriterCloseBlocking(t *testing.T) {
   589  	defer leaktest.Check(t)()
   590  
   591  	ctrl := xtest.NewController(t)
   592  	defer ctrl.Finish()
   593  
   594  	store := mem.NewStore()
   595  	cs := client.NewMockClient(ctrl)
   596  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   597  
   598  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   599  	require.NoError(t, err)
   600  
   601  	opts := testOptions().SetTopicService(ts)
   602  
   603  	sid1 := services.NewServiceID().SetName("s1")
   604  	cs1 := topic.NewConsumerService().SetConsumptionType(topic.Shared).SetServiceID(sid1)
   605  	testTopic := topic.NewTopic().
   606  		SetName(opts.TopicName()).
   607  		SetNumberOfShards(1).
   608  		SetConsumerServices([]topic.ConsumerService{cs1})
   609  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   610  	require.NoError(t, err)
   611  
   612  	sd := services.NewMockServices(ctrl)
   613  	opts = opts.SetServiceDiscovery(sd)
   614  	ps1 := testPlacementService(store, sid1)
   615  	sd.EXPECT().PlacementService(sid1, gomock.Any()).Return(ps1, nil)
   616  
   617  	p1 := placement.NewPlacement().
   618  		SetInstances([]placement.Instance{
   619  			placement.NewInstance().
   620  				SetID("i1").
   621  				SetEndpoint("addr1").
   622  				SetShards(shard.NewShards([]shard.Shard{
   623  					shard.NewShard(0).SetState(shard.Available),
   624  				})),
   625  		}).
   626  		SetShards([]uint32{0}).
   627  		SetReplicaFactor(1).
   628  		SetIsSharded(true)
   629  	_, err = ps1.Set(p1)
   630  	require.NoError(t, err)
   631  
   632  	w := NewWriter(opts).(*writer)
   633  	require.NoError(t, w.Init())
   634  	require.Equal(t, 1, len(w.consumerServiceWriters))
   635  
   636  	mm := producer.NewMockMessage(ctrl)
   637  	mm.EXPECT().Size().Return(3)
   638  	mm.EXPECT().Shard().Return(uint32(0)).Times(2)
   639  	mm.EXPECT().Bytes().Return([]byte("foo")).Times(1)
   640  	mm.EXPECT().Finalize(producer.Dropped)
   641  	rm := producer.NewRefCountedMessage(mm, nil)
   642  	require.NoError(t, w.Write(rm))
   643  
   644  	doneCh := make(chan struct{})
   645  	go func() {
   646  		w.Close()
   647  		close(doneCh)
   648  	}()
   649  
   650  	select {
   651  	case <-doneCh:
   652  		require.FailNow(t, "writer.Close() should block until all messages dropped or consumed")
   653  	default:
   654  	}
   655  
   656  	rm.Drop()
   657  	<-doneCh
   658  }
   659  
   660  func TestWriterSetMessageTTLNanosDropMetric(t *testing.T) {
   661  	defer leaktest.Check(t)()
   662  
   663  	ctrl := xtest.NewController(t)
   664  	defer ctrl.Finish()
   665  
   666  	store := mem.NewStore()
   667  	cs := client.NewMockClient(ctrl)
   668  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   669  
   670  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   671  	require.NoError(t, err)
   672  
   673  	opts := testOptions().SetTopicService(ts)
   674  
   675  	sid1 := services.NewServiceID().SetName("s1")
   676  	cs1 := topic.NewConsumerService().SetConsumptionType(topic.Replicated).SetServiceID(sid1)
   677  	sid2 := services.NewServiceID().SetName("s2")
   678  	cs2 := topic.NewConsumerService().SetConsumptionType(topic.Shared).SetServiceID(sid2)
   679  	testTopic := topic.NewTopic().
   680  		SetName(opts.TopicName()).
   681  		SetNumberOfShards(1).
   682  		SetConsumerServices([]topic.ConsumerService{cs1, cs2})
   683  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   684  	require.NoError(t, err)
   685  
   686  	sd := services.NewMockServices(ctrl)
   687  	opts = opts.SetServiceDiscovery(sd)
   688  	ps1 := testPlacementService(store, sid1)
   689  	sd.EXPECT().PlacementService(sid1, gomock.Any()).Return(ps1, nil)
   690  	ps2 := testPlacementService(store, sid2)
   691  	sd.EXPECT().PlacementService(sid2, gomock.Any()).Return(ps2, nil)
   692  
   693  	lis1, err := net.Listen("tcp", "127.0.0.1:0")
   694  	require.NoError(t, err)
   695  	defer lis1.Close()
   696  
   697  	p1 := placement.NewPlacement().
   698  		SetInstances([]placement.Instance{
   699  			placement.NewInstance().
   700  				SetID("i1").
   701  				SetEndpoint(lis1.Addr().String()).
   702  				SetShards(shard.NewShards([]shard.Shard{
   703  					shard.NewShard(0).SetState(shard.Available),
   704  				})),
   705  		}).
   706  		SetShards([]uint32{0}).
   707  		SetReplicaFactor(1).
   708  		SetIsSharded(true)
   709  	_, err = ps1.Set(p1)
   710  	require.NoError(t, err)
   711  
   712  	p2 := placement.NewPlacement().
   713  		SetInstances([]placement.Instance{
   714  			placement.NewInstance().
   715  				SetID("i2").
   716  				SetEndpoint("i2").
   717  				SetShards(shard.NewShards([]shard.Shard{
   718  					shard.NewShard(0).SetState(shard.Available),
   719  				})),
   720  		}).
   721  		SetShards([]uint32{0}).
   722  		SetReplicaFactor(1).
   723  		SetIsSharded(true)
   724  	_, err = ps2.Set(p2)
   725  	require.NoError(t, err)
   726  
   727  	w := NewWriter(opts).(*writer)
   728  	require.NoError(t, w.Init())
   729  	defer w.Close()
   730  
   731  	require.Equal(t, 2, len(w.consumerServiceWriters))
   732  
   733  	var called int
   734  	var wg sync.WaitGroup
   735  	mm := producer.NewMockMessage(ctrl)
   736  	mm.EXPECT().Shard().Return(uint32(0)).AnyTimes()
   737  	mm.EXPECT().Size().Return(3).AnyTimes()
   738  	mm.EXPECT().Bytes().Return([]byte("foo")).AnyTimes()
   739  	mm.EXPECT().Finalize(producer.Consumed).Do(func(interface{}) { called++; wg.Done() })
   740  	require.NoError(t, w.Write(producer.NewRefCountedMessage(mm, nil)))
   741  
   742  	wg.Add(1)
   743  	go func() {
   744  		testConsumeAndAckOnConnectionListener(t, lis1, opts.EncoderOptions(), opts.DecoderOptions())
   745  		wg.Done()
   746  	}()
   747  	wg.Wait()
   748  	require.Equal(t, 0, called)
   749  
   750  	// Wait for the message ttl update to trigger finalize.
   751  	wg.Add(1)
   752  	testTopic = topic.NewTopic().
   753  		SetName(opts.TopicName()).
   754  		SetNumberOfShards(1).
   755  		SetConsumerServices([]topic.ConsumerService{cs1, cs2.SetMessageTTLNanos(int64(50 * time.Millisecond))})
   756  	_, err = ts.CheckAndSet(testTopic, 1)
   757  	require.NoError(t, err)
   758  	wg.Wait()
   759  	require.Equal(t, 1, called)
   760  
   761  	lis2, err := net.Listen("tcp", "127.0.0.1:0")
   762  	require.NoError(t, err)
   763  	defer lis2.Close()
   764  
   765  	p2 = placement.NewPlacement().
   766  		SetInstances([]placement.Instance{
   767  			placement.NewInstance().
   768  				SetID("i2").
   769  				SetEndpoint(lis2.Addr().String()).
   770  				SetShards(shard.NewShards([]shard.Shard{
   771  					shard.NewShard(0).SetState(shard.Available),
   772  				})),
   773  		}).
   774  		SetShards([]uint32{0}).
   775  		SetReplicaFactor(1).
   776  		SetIsSharded(true)
   777  	_, err = ps2.Set(p2)
   778  	require.NoError(t, err)
   779  
   780  	testTopic = topic.NewTopic().
   781  		SetName(opts.TopicName()).
   782  		SetNumberOfShards(1).
   783  		SetConsumerServices([]topic.ConsumerService{cs1, cs2.SetMessageTTLNanos(0)})
   784  	_, err = ts.CheckAndSet(testTopic, 2)
   785  	require.NoError(t, err)
   786  
   787  	require.NoError(t, w.Write(producer.NewRefCountedMessage(mm, nil)))
   788  
   789  	called = 0
   790  	mm.EXPECT().Finalize(producer.Consumed).Do(func(interface{}) { called++; wg.Done() })
   791  	wg.Add(1)
   792  	go func() {
   793  		testConsumeAndAckOnConnectionListener(t, lis1, opts.EncoderOptions(), opts.DecoderOptions())
   794  		wg.Done()
   795  	}()
   796  	wg.Wait()
   797  	require.Equal(t, 0, called)
   798  
   799  	time.Sleep(200 * time.Millisecond)
   800  	require.Equal(t, 0, called)
   801  
   802  	// Wait for the consumer to trigger finalize because there is no more message ttl.
   803  	wg.Add(1)
   804  	go func() {
   805  		testConsumeAndAckOnConnectionListener(t, lis2, opts.EncoderOptions(), opts.DecoderOptions())
   806  	}()
   807  	wg.Wait()
   808  	require.Equal(t, 1, called)
   809  
   810  	w.Close()
   811  }
   812  
   813  func TestWriterNumShards(t *testing.T) {
   814  	defer leaktest.Check(t)()
   815  
   816  	ctrl := xtest.NewController(t)
   817  	defer ctrl.Finish()
   818  
   819  	store := mem.NewStore()
   820  	cs := client.NewMockClient(ctrl)
   821  	cs.EXPECT().Store(gomock.Any()).Return(store, nil)
   822  
   823  	ts, err := topic.NewService(topic.NewServiceOptions().SetConfigService(cs))
   824  	require.NoError(t, err)
   825  
   826  	opts := testOptions().SetTopicService(ts)
   827  
   828  	testTopic := topic.NewTopic().
   829  		SetName(opts.TopicName()).
   830  		SetNumberOfShards(2)
   831  	_, err = ts.CheckAndSet(testTopic, kv.UninitializedVersion)
   832  	require.NoError(t, err)
   833  
   834  	w := NewWriter(opts).(*writer)
   835  	defer w.Close()
   836  
   837  	require.Equal(t, 0, int(w.NumShards()))
   838  
   839  	require.NoError(t, w.Init())
   840  	require.Equal(t, 2, int(w.NumShards()))
   841  }