github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/swarmkit/manager/logbroker/broker_test.go (about) 1 package logbroker 2 3 import ( 4 "context" 5 "fmt" 6 "io" 7 "net" 8 "sync" 9 "testing" 10 "time" 11 12 "google.golang.org/grpc" 13 14 "github.com/docker/swarmkit/api" 15 "github.com/docker/swarmkit/ca" 16 "github.com/docker/swarmkit/ca/testutils" 17 "github.com/docker/swarmkit/manager/state/store" 18 "github.com/docker/swarmkit/protobuf/ptypes" 19 gogotypes "github.com/gogo/protobuf/types" 20 "github.com/stretchr/testify/require" 21 ) 22 23 func TestLogBrokerLogs(t *testing.T) { 24 ctx, ca, broker, serverAddr, brokerAddr, done := testLogBrokerEnv(t) 25 defer done() 26 27 client, clientDone := testLogClient(t, serverAddr) 28 defer clientDone() 29 brokerClient, agentSecurity, brokerClientDone := testBrokerClient(t, ca, brokerAddr) 30 defer brokerClientDone() 31 32 var ( 33 wg sync.WaitGroup 34 hold = make(chan struct{}) // coordinates pubsub start 35 messagesExpected int 36 ) 37 38 subStream, err := brokerClient.ListenSubscriptions(ctx, &api.ListenSubscriptionsRequest{}) 39 if err != nil { 40 t.Fatal(err) 41 } 42 43 stream, err := client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 44 Options: &api.LogSubscriptionOptions{ 45 Follow: true, 46 }, 47 Selector: &api.LogSelector{ 48 NodeIDs: []string{agentSecurity.ServerTLSCreds.NodeID()}, 49 }, 50 }) 51 if err != nil { 52 t.Fatalf("error subscribing: %v", err) 53 } 54 55 subscription, err := subStream.Recv() 56 if err != nil { 57 t.Fatal(err) 58 } 59 60 // spread some services across nodes with a bunch of tasks. 61 const ( 62 nNodes = 5 63 nServices = 20 64 nTasksPerService = 20 65 nLogMessagesPerTask = 5 66 ) 67 68 for service := 0; service < nServices; service++ { 69 serviceID := fmt.Sprintf("service-%v", service) 70 71 for task := 0; task < nTasksPerService; task++ { 72 taskID := fmt.Sprintf("%v.task-%v", serviceID, task) 73 74 for node := 0; node < nNodes; node++ { 75 nodeID := fmt.Sprintf("node-%v", node) 76 77 if (task+1)%(node+1) != 0 { 78 continue 79 } 80 messagesExpected += nLogMessagesPerTask 81 82 wg.Add(1) 83 go func(nodeID, serviceID, taskID string) { 84 <-hold 85 86 // Each goroutine gets its own publisher 87 publisher, err := brokerClient.PublishLogs(ctx) 88 require.NoError(t, err) 89 90 defer func() { 91 _, err := publisher.CloseAndRecv() 92 require.NoError(t, err) 93 wg.Done() 94 }() 95 96 msgctx := api.LogContext{ 97 NodeID: agentSecurity.ClientTLSCreds.NodeID(), 98 ServiceID: serviceID, 99 TaskID: taskID, 100 } 101 for i := 0; i < nLogMessagesPerTask; i++ { 102 require.NoError(t, publisher.Send(&api.PublishLogsMessage{ 103 SubscriptionID: subscription.ID, 104 Messages: []api.LogMessage{newLogMessage(msgctx, "log message number %d", i)}, 105 })) 106 } 107 }(nodeID, serviceID, taskID) 108 } 109 } 110 } 111 112 t.Logf("expected %v messages", messagesExpected) 113 close(hold) 114 var messages int 115 for messages < messagesExpected { 116 msgs, err := stream.Recv() 117 require.NoError(t, err) 118 for range msgs.Messages { 119 messages++ 120 if messages%100 == 0 { 121 fmt.Println(messages, "received") 122 } 123 } 124 } 125 t.Logf("received %v messages", messages) 126 127 wg.Wait() 128 129 // Make sure double Start throws an error 130 require.EqualError(t, broker.Start(ctx), errAlreadyRunning.Error()) 131 // Stop should work 132 require.NoError(t, broker.Stop()) 133 // Double stopping should fail 134 require.EqualError(t, broker.Stop(), errNotRunning.Error()) 135 } 136 137 func listenSubscriptions(ctx context.Context, t *testing.T, client api.LogBrokerClient) <-chan *api.SubscriptionMessage { 138 subscriptions, err := client.ListenSubscriptions(ctx, &api.ListenSubscriptionsRequest{}) 139 require.NoError(t, err) 140 141 ch := make(chan *api.SubscriptionMessage) 142 go func() { 143 defer close(ch) 144 145 for { 146 select { 147 case <-ctx.Done(): 148 return 149 default: 150 } 151 sub, err := subscriptions.Recv() 152 if err != nil { 153 return 154 } 155 ch <- sub 156 } 157 }() 158 159 return ch 160 } 161 162 func ensureSubscription(t *testing.T, subscriptions <-chan *api.SubscriptionMessage) *api.SubscriptionMessage { 163 select { 164 case s := <-subscriptions: 165 require.NotNil(t, s) 166 return s 167 case <-time.After(5 * time.Second): 168 require.FailNow(t, "subscription expected") 169 } 170 return nil 171 } 172 173 func ensureNoSubscription(t *testing.T, subscriptions <-chan *api.SubscriptionMessage) { 174 select { 175 case s := <-subscriptions: 176 require.FailNow(t, fmt.Sprintf("unexpected subscription: %v", s)) 177 case <-time.After(10 * time.Millisecond): 178 return 179 } 180 } 181 182 func TestLogBrokerSubscriptions(t *testing.T) { 183 ctx, ca, _, serverAddr, brokerAddr, done := testLogBrokerEnv(t) 184 defer done() 185 186 client, clientDone := testLogClient(t, serverAddr) 187 defer clientDone() 188 189 agent1, agent1Security, agent1Done := testBrokerClient(t, ca, brokerAddr) 190 defer agent1Done() 191 192 agent2, agent2Security, agent2Done := testBrokerClient(t, ca, brokerAddr) 193 defer agent2Done() 194 195 // Have an agent listen to subscriptions before anyone has subscribed. 196 subscriptions1 := listenSubscriptions(ctx, t, agent1) 197 198 // Send two subscriptions - one will match both agent1 and agent2 while 199 // the other only agent1 200 _, err := client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 201 Options: &api.LogSubscriptionOptions{ 202 Follow: true, 203 }, 204 Selector: &api.LogSelector{ 205 NodeIDs: []string{ 206 agent1Security.ServerTLSCreds.NodeID(), 207 }, 208 }, 209 }) 210 require.NoError(t, err) 211 _, err = client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 212 Options: &api.LogSubscriptionOptions{ 213 Follow: true, 214 }, 215 Selector: &api.LogSelector{ 216 NodeIDs: []string{ 217 agent1Security.ServerTLSCreds.NodeID(), 218 agent2Security.ServerTLSCreds.NodeID(), 219 }, 220 }, 221 }) 222 require.NoError(t, err) 223 224 // Make sure we received two subscriptions on agent 1 (already joined). 225 { 226 s1 := ensureSubscription(t, subscriptions1) 227 require.False(t, s1.Close) 228 require.Contains(t, s1.Selector.NodeIDs, agent1Security.ServerTLSCreds.NodeID()) 229 230 s2 := ensureSubscription(t, subscriptions1) 231 require.False(t, s2.Close) 232 require.Contains(t, s2.Selector.NodeIDs, agent1Security.ServerTLSCreds.NodeID()) 233 234 // Ensure we received two different subscriptions. 235 require.NotEqual(t, s1.ID, s2.ID) 236 } 237 238 // Join a second agent. 239 subscriptions2 := listenSubscriptions(ctx, t, agent2) 240 241 // Make sure we receive past subscriptions. 242 // Make sure we receive *only* the right one. 243 { 244 s := ensureSubscription(t, subscriptions2) 245 require.False(t, s.Close) 246 require.Equal(t, []string{agent1Security.ServerTLSCreds.NodeID(), agent2Security.ServerTLSCreds.NodeID()}, s.Selector.NodeIDs) 247 248 ensureNoSubscription(t, subscriptions2) 249 } 250 } 251 252 func TestLogBrokerSelector(t *testing.T) { 253 ctx, ca, _, serverAddr, brokerAddr, done := testLogBrokerEnv(t) 254 defer done() 255 256 client, clientDone := testLogClient(t, serverAddr) 257 defer clientDone() 258 259 agent1, agent1Security, agent1Done := testBrokerClient(t, ca, brokerAddr) 260 defer agent1Done() 261 agent1subscriptions := listenSubscriptions(ctx, t, agent1) 262 263 agent2, agent2Security, agent2Done := testBrokerClient(t, ca, brokerAddr) 264 defer agent2Done() 265 266 agent2subscriptions := listenSubscriptions(ctx, t, agent2) 267 268 // Subscribe to a task. 269 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 270 return store.CreateTask(tx, &api.Task{ 271 ID: "task", 272 }) 273 })) 274 _, err := client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 275 Options: &api.LogSubscriptionOptions{ 276 Follow: true, 277 }, 278 Selector: &api.LogSelector{ 279 TaskIDs: []string{"task"}, 280 }, 281 }) 282 require.NoError(t, err) 283 284 // Since it's not assigned to any agent, nobody should receive it. 285 ensureNoSubscription(t, agent1subscriptions) 286 ensureNoSubscription(t, agent2subscriptions) 287 288 // Assign the task to agent-1. Make sure it's received by agent-1 but *not* 289 // agent-2. 290 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 291 task := store.GetTask(tx, "task") 292 require.NotNil(t, task) 293 task.NodeID = agent1Security.ServerTLSCreds.NodeID() 294 return store.UpdateTask(tx, task) 295 })) 296 297 ensureSubscription(t, agent1subscriptions) 298 ensureNoSubscription(t, agent2subscriptions) 299 300 // Subscribe to a service. 301 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 302 return store.CreateService(tx, &api.Service{ 303 ID: "service", 304 }) 305 })) 306 _, err = client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 307 Options: &api.LogSubscriptionOptions{ 308 Follow: true, 309 }, 310 Selector: &api.LogSelector{ 311 ServiceIDs: []string{"service"}, 312 }, 313 }) 314 require.NoError(t, err) 315 316 // Since there are no corresponding tasks, nobody should receive it. 317 ensureNoSubscription(t, agent1subscriptions) 318 ensureNoSubscription(t, agent2subscriptions) 319 320 // Create a task that does *NOT* belong to our service and assign it to node-1. 321 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 322 return store.CreateTask(tx, &api.Task{ 323 ID: "wrong-task", 324 ServiceID: "wrong-service", 325 NodeID: agent1Security.ServerTLSCreds.NodeID(), 326 }) 327 })) 328 329 // Ensure agent-1 doesn't receive it. 330 ensureNoSubscription(t, agent1subscriptions) 331 332 // Now create another task that does belong to our service and assign it to node-1. 333 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 334 return store.CreateTask(tx, &api.Task{ 335 ID: "service-task-1", 336 ServiceID: "service", 337 NodeID: agent1Security.ServerTLSCreds.NodeID(), 338 }) 339 })) 340 341 // Make sure agent-1 receives it... 342 ensureSubscription(t, agent1subscriptions) 343 // ...and agent-2 does not. 344 ensureNoSubscription(t, agent2subscriptions) 345 346 // Create another task, same as above. 347 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 348 return store.CreateTask(tx, &api.Task{ 349 ID: "service-task-2", 350 ServiceID: "service", 351 NodeID: agent1Security.ServerTLSCreds.NodeID(), 352 }) 353 })) 354 355 // agent-1 should *not* receive it anymore since the subscription was already delivered. 356 // agent-2 should still not get it. 357 ensureNoSubscription(t, agent1subscriptions) 358 ensureNoSubscription(t, agent2subscriptions) 359 360 // Now, create another one and assign it to agent-2. 361 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 362 return store.CreateTask(tx, &api.Task{ 363 ID: "service-task-3", 364 ServiceID: "service", 365 NodeID: agent2Security.ServerTLSCreds.NodeID(), 366 }) 367 })) 368 369 // Make sure it's delivered to agent-2. 370 ensureSubscription(t, agent2subscriptions) 371 // it shouldn't do anything for agent-1. 372 ensureNoSubscription(t, agent1subscriptions) 373 } 374 375 func TestLogBrokerNoFollow(t *testing.T) { 376 t.Parallel() 377 378 ctx, ca, _, serverAddr, brokerAddr, done := testLogBrokerEnv(t) 379 defer done() 380 381 client, clientDone := testLogClient(t, serverAddr) 382 defer clientDone() 383 384 agent1, agent1Security, agent1Done := testBrokerClient(t, ca, brokerAddr) 385 defer agent1Done() 386 agent1subscriptions := listenSubscriptions(ctx, t, agent1) 387 388 agent2, agent2Security, agent2Done := testBrokerClient(t, ca, brokerAddr) 389 defer agent2Done() 390 agent2subscriptions := listenSubscriptions(ctx, t, agent2) 391 392 // Create fake environment. 393 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 394 if err := store.CreateTask(tx, &api.Task{ 395 ID: "task1", 396 ServiceID: "service", 397 Status: api.TaskStatus{ 398 State: api.TaskStateRunning, 399 }, 400 NodeID: agent1Security.ServerTLSCreds.NodeID(), 401 }); err != nil { 402 return err 403 } 404 405 return store.CreateTask(tx, &api.Task{ 406 ID: "task2", 407 ServiceID: "service", 408 Status: api.TaskStatus{ 409 State: api.TaskStateRunning, 410 }, 411 NodeID: agent2Security.ServerTLSCreds.NodeID(), 412 }) 413 })) 414 415 // We need to sleep here to give ListenSubscriptions time to call 416 // registerSubscription before SubscribeLogs concludes that one or both 417 // of the agents are not connected, and prematurely calls Done for one 418 // or both nodes. Think of these stream RPC calls as goroutines which 419 // don't have synchronization around anything that happens in the RPC 420 // handler before a send or receive. It would be nice if we had a way 421 // of confirming that a node was listening for subscriptions before 422 // calling SubscribeLogs, but the current API doesn't provide this. 423 time.Sleep(time.Second) 424 425 // Subscribe to logs in no follow mode 426 logs, err := client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 427 Options: &api.LogSubscriptionOptions{ 428 Follow: false, 429 }, 430 Selector: &api.LogSelector{ 431 ServiceIDs: []string{"service"}, 432 }, 433 }) 434 require.NoError(t, err) 435 436 // Get the subscriptions from the agents. 437 subscription1 := ensureSubscription(t, agent1subscriptions) 438 require.Equal(t, subscription1.Selector.ServiceIDs[0], "service") 439 subscription2 := ensureSubscription(t, agent2subscriptions) 440 require.Equal(t, subscription2.Selector.ServiceIDs[0], "service") 441 442 require.Equal(t, subscription1.ID, subscription2.ID) 443 444 // Publish a log message from agent-1 and close the publisher 445 publisher, err := agent1.PublishLogs(ctx) 446 require.NoError(t, err) 447 require.NoError(t, 448 publisher.Send(&api.PublishLogsMessage{ 449 SubscriptionID: subscription1.ID, 450 Messages: []api.LogMessage{ 451 newLogMessage(api.LogContext{ 452 NodeID: agent1Security.ServerTLSCreds.NodeID(), 453 ServiceID: "service", 454 TaskID: "task1", 455 }, "log message"), 456 }, 457 })) 458 _, err = publisher.CloseAndRecv() 459 require.NoError(t, err) 460 461 // Ensure we get it from the other end 462 log, err := logs.Recv() 463 require.NoError(t, err) 464 require.Len(t, log.Messages, 1) 465 require.Equal(t, log.Messages[0].Context.NodeID, agent1Security.ServerTLSCreds.NodeID()) 466 467 // Now publish a message from the other agent and close the subscription 468 publisher, err = agent2.PublishLogs(ctx) 469 require.NoError(t, err) 470 require.NoError(t, 471 publisher.Send(&api.PublishLogsMessage{ 472 SubscriptionID: subscription2.ID, 473 Messages: []api.LogMessage{ 474 newLogMessage(api.LogContext{ 475 NodeID: agent2Security.ServerTLSCreds.NodeID(), 476 ServiceID: "service", 477 TaskID: "task2", 478 }, "log message"), 479 }, 480 })) 481 _, err = publisher.CloseAndRecv() 482 require.NoError(t, err) 483 484 // Ensure we get it from the other end 485 log, err = logs.Recv() 486 require.NoError(t, err) 487 require.Len(t, log.Messages, 1) 488 require.Equal(t, log.Messages[0].Context.NodeID, agent2Security.ServerTLSCreds.NodeID()) 489 490 // Since we receive both messages the log stream should end 491 _, err = logs.Recv() 492 require.Equal(t, err, io.EOF) 493 } 494 495 func TestLogBrokerNoFollowMissingNode(t *testing.T) { 496 t.Parallel() 497 498 ctx, ca, _, serverAddr, brokerAddr, done := testLogBrokerEnv(t) 499 defer done() 500 501 client, clientDone := testLogClient(t, serverAddr) 502 defer clientDone() 503 504 agent, agentSecurity, agentDone := testBrokerClient(t, ca, brokerAddr) 505 defer agentDone() 506 agentSubscriptions := listenSubscriptions(ctx, t, agent) 507 508 // Create fake environment. 509 // A service with one instance on a genuine node and another instance 510 // and a node that didn't connect to the broker. 511 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 512 if err := store.CreateTask(tx, &api.Task{ 513 ID: "task1", 514 ServiceID: "service", 515 Status: api.TaskStatus{ 516 State: api.TaskStateRunning, 517 }, 518 NodeID: agentSecurity.ServerTLSCreds.NodeID(), 519 }); err != nil { 520 return err 521 } 522 523 return store.CreateTask(tx, &api.Task{ 524 ID: "task2", 525 ServiceID: "service", 526 NodeID: "node-2", 527 Status: api.TaskStatus{ 528 State: api.TaskStateRunning, 529 }, 530 }) 531 })) 532 533 // We need to sleep here to give ListenSubscriptions time to call 534 // registerSubscription before SubscribeLogs concludes that the actual 535 // agent is not connected, and prematurely calls Done for it. Think of 536 // these stream RPC calls as goroutines which don't have synchronization 537 // around anything that happens in the RPC handler before a send or 538 // receive. It would be nice if we had a way of confirming that a node 539 // was listening for subscriptions before calling SubscribeLogs, but 540 // the current API doesn't provide this. 541 time.Sleep(time.Second) 542 543 // Subscribe to logs in no follow mode 544 logs, err := client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 545 Options: &api.LogSubscriptionOptions{ 546 Follow: false, 547 }, 548 Selector: &api.LogSelector{ 549 ServiceIDs: []string{"service"}, 550 }, 551 }) 552 require.NoError(t, err) 553 554 // Grab the subscription and publish a log message from the connected agent. 555 subscription := ensureSubscription(t, agentSubscriptions) 556 require.Equal(t, subscription.Selector.ServiceIDs[0], "service") 557 publisher, err := agent.PublishLogs(ctx) 558 require.NoError(t, err) 559 require.NoError(t, 560 publisher.Send(&api.PublishLogsMessage{ 561 SubscriptionID: subscription.ID, 562 Messages: []api.LogMessage{ 563 newLogMessage(api.LogContext{ 564 NodeID: agentSecurity.ServerTLSCreds.NodeID(), 565 ServiceID: "service", 566 TaskID: "task1", 567 }, "log message"), 568 }, 569 })) 570 _, err = publisher.CloseAndRecv() 571 require.NoError(t, err) 572 573 // Ensure we receive the message that we could grab 574 log, err := logs.Recv() 575 require.NoError(t, err) 576 require.Len(t, log.Messages, 1) 577 require.Equal(t, log.Messages[0].Context.NodeID, agentSecurity.ServerTLSCreds.NodeID()) 578 579 // Ensure the log stream ends with an error complaining about the missing node 580 _, err = logs.Recv() 581 require.Error(t, err) 582 require.Contains(t, err.Error(), "node-2 is not available") 583 } 584 585 func TestLogBrokerNoFollowNotYetRunningTask(t *testing.T) { 586 ctx, ca, _, serverAddr, _, done := testLogBrokerEnv(t) 587 defer done() 588 589 client, clientDone := testLogClient(t, serverAddr) 590 defer clientDone() 591 592 // Create fake environment. 593 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 594 return store.CreateTask(tx, &api.Task{ 595 ID: "task1", 596 ServiceID: "service", 597 Status: api.TaskStatus{ 598 State: api.TaskStateNew, 599 }, 600 }) 601 })) 602 603 // Subscribe to logs in no follow mode 604 logs, err := client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 605 Options: &api.LogSubscriptionOptions{ 606 Follow: false, 607 }, 608 Selector: &api.LogSelector{ 609 ServiceIDs: []string{"service"}, 610 }, 611 }) 612 require.NoError(t, err) 613 614 // The log stream should be empty, because the task was not yet running 615 _, err = logs.Recv() 616 require.Error(t, err) 617 require.Equal(t, err, io.EOF) 618 } 619 620 func TestLogBrokerNoFollowDisconnect(t *testing.T) { 621 t.Parallel() 622 623 ctx, ca, _, serverAddr, brokerAddr, done := testLogBrokerEnv(t) 624 defer done() 625 626 client, clientDone := testLogClient(t, serverAddr) 627 defer clientDone() 628 629 agent1, agent1Security, agent1Done := testBrokerClient(t, ca, brokerAddr) 630 defer agent1Done() 631 agent1subscriptions := listenSubscriptions(ctx, t, agent1) 632 633 agent2, agent2Security, agent2Done := testBrokerClient(t, ca, brokerAddr) 634 defer agent2Done() 635 agent2subscriptions := listenSubscriptions(ctx, t, agent2) 636 637 // Create fake environment. 638 require.NoError(t, ca.MemoryStore.Update(func(tx store.Tx) error { 639 if err := store.CreateTask(tx, &api.Task{ 640 ID: "task1", 641 ServiceID: "service", 642 Status: api.TaskStatus{ 643 State: api.TaskStateRunning, 644 }, 645 NodeID: agent1Security.ServerTLSCreds.NodeID(), 646 }); err != nil { 647 return err 648 } 649 650 return store.CreateTask(tx, &api.Task{ 651 ID: "task2", 652 ServiceID: "service", 653 Status: api.TaskStatus{ 654 State: api.TaskStateRunning, 655 }, 656 NodeID: agent2Security.ServerTLSCreds.NodeID(), 657 }) 658 })) 659 660 // We need to sleep here to give ListenSubscriptions time to call 661 // registerSubscription before SubscribeLogs concludes that one or both 662 // of the agents are not connected, and prematurely calls Done for one 663 // or both nodes. Think of these stream RPC calls as goroutines which 664 // don't have synchronization around anything that happens in the RPC 665 // handler before a send or receive. It would be nice if we had a way 666 // of confirming that a node was listening for subscriptions before 667 // calling SubscribeLogs, but the current API doesn't provide this. 668 time.Sleep(time.Second) 669 670 // Subscribe to logs in no follow mode 671 logs, err := client.SubscribeLogs(ctx, &api.SubscribeLogsRequest{ 672 Options: &api.LogSubscriptionOptions{ 673 Follow: false, 674 }, 675 Selector: &api.LogSelector{ 676 ServiceIDs: []string{"service"}, 677 }, 678 }) 679 require.NoError(t, err) 680 681 // Get the subscriptions from the agents. 682 subscription1 := ensureSubscription(t, agent1subscriptions) 683 require.Equal(t, subscription1.Selector.ServiceIDs[0], "service") 684 subscription2 := ensureSubscription(t, agent2subscriptions) 685 require.Equal(t, subscription2.Selector.ServiceIDs[0], "service") 686 687 require.Equal(t, subscription1.ID, subscription2.ID) 688 689 // Publish a log message from agent-1 and close the publisher 690 publisher, err := agent1.PublishLogs(ctx) 691 require.NoError(t, err) 692 require.NoError(t, 693 publisher.Send(&api.PublishLogsMessage{ 694 SubscriptionID: subscription1.ID, 695 Messages: []api.LogMessage{ 696 newLogMessage(api.LogContext{ 697 NodeID: agent1Security.ServerTLSCreds.NodeID(), 698 ServiceID: "service", 699 TaskID: "task1", 700 }, "log message"), 701 }, 702 })) 703 _, err = publisher.CloseAndRecv() 704 require.NoError(t, err) 705 706 // Now suddenly disconnect agent2... 707 agent2Done() 708 709 // Ensure we get the first message 710 log, err := logs.Recv() 711 require.NoError(t, err) 712 require.Len(t, log.Messages, 1) 713 require.Equal(t, log.Messages[0].Context.NodeID, agent1Security.ServerTLSCreds.NodeID()) 714 715 // ...and then an error 716 _, err = logs.Recv() 717 require.Error(t, err) 718 require.Contains(t, err.Error(), "disconnected unexpectedly") 719 } 720 721 func testLogBrokerEnv(t *testing.T) (context.Context, *testutils.TestCA, *LogBroker, string, string, func()) { 722 ctx, cancel := context.WithCancel(context.Background()) 723 724 tca := testutils.NewTestCA(t) 725 broker := New(tca.MemoryStore) 726 727 // Log Server 728 logListener, err := net.Listen("tcp", "localhost:0") 729 if err != nil { 730 t.Fatalf("error setting up listener: %v", err) 731 } 732 logServer := grpc.NewServer() 733 api.RegisterLogsServer(logServer, broker) 734 735 go func() { 736 if err := logServer.Serve(logListener); err != nil { 737 // SIGH(stevvooe): GRPC won't really shutdown gracefully. 738 // This should be fatal. 739 t.Logf("error serving grpc service: %v", err) 740 } 741 }() 742 743 // Log Broker 744 brokerListener, err := net.Listen("tcp", "localhost:0") 745 if err != nil { 746 t.Fatalf("error setting up listener: %v", err) 747 } 748 749 securityConfig, err := tca.NewNodeConfig(ca.ManagerRole) 750 if err != nil { 751 t.Fatal(err) 752 } 753 serverOpts := []grpc.ServerOption{grpc.Creds(securityConfig.ServerTLSCreds)} 754 brokerServer := grpc.NewServer(serverOpts...) 755 756 authorize := func(ctx context.Context, roles []string) error { 757 _, err := ca.AuthorizeForwardedRoleAndOrg(ctx, roles, []string{ca.ManagerRole}, tca.Organization, nil) 758 return err 759 } 760 authenticatedLogBrokerAPI := api.NewAuthenticatedWrapperLogBrokerServer(broker, authorize) 761 762 api.RegisterLogBrokerServer(brokerServer, authenticatedLogBrokerAPI) 763 go func() { 764 if err := brokerServer.Serve(brokerListener); err != nil { 765 // SIGH(stevvooe): GRPC won't really shutdown gracefully. 766 // This should be fatal. 767 t.Logf("error serving grpc service: %v", err) 768 } 769 }() 770 771 require.NoError(t, broker.Start(ctx)) 772 773 return ctx, tca, broker, logListener.Addr().String(), brokerListener.Addr().String(), func() { 774 broker.Stop() 775 776 logServer.Stop() 777 brokerServer.Stop() 778 779 logListener.Close() 780 brokerListener.Close() 781 782 cancel() 783 } 784 } 785 786 func testLogClient(t *testing.T, addr string) (api.LogsClient, func()) { 787 // Log client 788 logCc, err := grpc.Dial(addr, grpc.WithInsecure()) 789 if err != nil { 790 t.Fatalf("error dialing local server: %v", err) 791 } 792 return api.NewLogsClient(logCc), func() { 793 logCc.Close() 794 } 795 } 796 797 func testBrokerClient(t *testing.T, tca *testutils.TestCA, addr string) (api.LogBrokerClient, *ca.SecurityConfig, func()) { 798 securityConfig, err := tca.NewNodeConfig(ca.WorkerRole) 799 if err != nil { 800 t.Fatal(err) 801 } 802 803 opts := []grpc.DialOption{grpc.WithTimeout(10 * time.Second), grpc.WithTransportCredentials(securityConfig.ClientTLSCreds)} 804 cc, err := grpc.Dial(addr, opts...) 805 if err != nil { 806 t.Fatalf("error dialing local server: %v", err) 807 } 808 809 return api.NewLogBrokerClient(cc), securityConfig, func() { 810 cc.Close() 811 } 812 } 813 814 func printLogMessages(msgs ...api.LogMessage) { 815 for _, msg := range msgs { 816 ts, _ := gogotypes.TimestampFromProto(msg.Timestamp) 817 fmt.Printf("%v %v %s\n", msg.Context, ts, string(msg.Data)) 818 } 819 } 820 821 // newLogMessage is just a helper to build a new log message. 822 func newLogMessage(msgctx api.LogContext, format string, vs ...interface{}) api.LogMessage { 823 return api.LogMessage{ 824 Context: msgctx, 825 Timestamp: ptypes.MustTimestampProto(time.Now()), 826 Data: []byte(fmt.Sprintf(format, vs...)), 827 } 828 }