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  }