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 }