// Publisher returns a LogPublisher for the given subscription
// as well as a cancel function that should be called when the log stream
// is completed.
func (a *Agent) Publisher(ctx context.Context, subscriptionID string) (exec.LogPublisher, func(), error) {
// TODO(stevvooe): The level of coordination here is WAY too much for logs.
// These should only be best effort and really just buffer until a session is
// ready. Ideally, they would use a separate connection completely.
var (
err error
publisher api.LogBroker_PublishLogsClient
)
err = a.withSession(ctx, func(session *session) error {
publisher, err = api.NewLogBrokerClient(session.conn.ClientConn).PublishLogs(ctx)
return err
})
if err != nil {
return nil, nil, err
}
return exec.LogPublisherFunc(func(ctx context.Context, message api.LogMessage) error {
select {
case <-ctx.Done():
publisher.CloseSend()
return ctx.Err()
default:
}
return publisher.Send(&api.PublishLogsMessage{
SubscriptionID: subscriptionID,
Messages: []api.LogMessage{message},
})
}), func() {
publisher.CloseSend()
}, nil
}
func (s *session) logSubscriptions(ctx context.Context) error {
log := log.G(ctx).WithFields(logrus.Fields{"method": "(*session).logSubscriptions"})
log.Debugf("")
client := api.NewLogBrokerClient(s.conn)
subscriptions, err := client.ListenSubscriptions(ctx, &api.ListenSubscriptionsRequest{})
if err != nil {
return err
}
defer subscriptions.CloseSend()
for {
resp, err := subscriptions.Recv()
if err != nil {
return err
}
select {
case s.subscriptions <- resp:
case <-s.closed:
return errSessionClosed
case <-ctx.Done():
return ctx.Err()
}
}
}
func testBrokerClient(t *testing.T, tca *testutils.TestCA, addr string) (api.LogBrokerClient, *ca.SecurityConfig, func()) {
securityConfig, err := tca.NewNodeConfig(ca.WorkerRole)
if err != nil {
t.Fatal(err)
}
opts := []grpc.DialOption{grpc.WithTimeout(10 * time.Second), grpc.WithTransportCredentials(securityConfig.ClientTLSCreds)}
cc, err := grpc.Dial(addr, opts...)
if err != nil {
t.Fatalf("error dialing local server: %v", err)
}
return api.NewLogBrokerClient(cc), securityConfig, func() {
cc.Close()
}
}
func (s *session) logSubscriptions(ctx context.Context) error {
log := log.G(ctx).WithFields(logrus.Fields{"method": "(*session).logSubscriptions"})
log.Debugf("")
client := api.NewLogBrokerClient(s.conn)
subscriptions, err := client.ListenSubscriptions(ctx, &api.ListenSubscriptionsRequest{})
if err != nil {
return err
}
defer subscriptions.CloseSend()
for {
resp, err := subscriptions.Recv()
if grpc.Code(err) == codes.Unimplemented {
log.Warning("manager does not support log subscriptions")
// Don't return, because returning would bounce the session
select {
case <-s.closed:
return errSessionClosed
case <-ctx.Done():
return ctx.Err()
}
}
if err != nil {
return err
}
select {
case s.subscriptions <- resp:
case <-s.closed:
return errSessionClosed
case <-ctx.Done():
return ctx.Err()
}
}
}