forked from hashicorp/consul
-
Notifications
You must be signed in to change notification settings - Fork 0
/
watch_servers.go
136 lines (116 loc) · 4.45 KB
/
watch_servers.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
package serverdiscovery
import (
"context"
"errors"
"github.com/hashicorp/go-hclog"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/hernad/consul/agent/consul/autopilotevents"
"github.com/hernad/consul/agent/consul/stream"
external "github.com/hernad/consul/agent/grpc-external"
"github.com/hernad/consul/agent/structs"
"github.com/hernad/consul/proto-public/pbserverdiscovery"
)
// WatchServers provides a stream on which you can receive the list of servers
// that are ready to receive incoming requests including stale queries. The
// current set of ready servers are sent immediately at the start of the
// stream and new updates will be sent whenver the set of ready servers changes.
func (s *Server) WatchServers(req *pbserverdiscovery.WatchServersRequest, serverStream pbserverdiscovery.ServerDiscoveryService_WatchServersServer) error {
logger := s.Logger.Named("watch-servers").With("request_id", external.TraceID())
logger.Debug("starting stream")
defer logger.Trace("stream closed")
options, err := external.QueryOptionsFromContext(serverStream.Context())
if err != nil {
return err
}
// Serve the ready servers from an EventPublisher subscription. If the subscription is
// closed due to an ACL change, we'll attempt to re-authorize and resume it to
// prevent unnecessarily terminating the stream.
var idx uint64
for {
var err error
idx, err = s.serveReadyServers(options.Token, idx, req, serverStream, logger)
if errors.Is(err, stream.ErrSubForceClosed) {
logger.Trace("subscription force-closed due to an ACL change or snapshot restore, will attempt to re-auth and resume")
} else {
return err
}
}
}
func (s *Server) serveReadyServers(token string, index uint64, req *pbserverdiscovery.WatchServersRequest, serverStream pbserverdiscovery.ServerDiscoveryService_WatchServersServer, logger hclog.Logger) (uint64, error) {
if err := external.RequireAnyValidACLToken(s.ACLResolver, token); err != nil {
return 0, err
}
// Start the subscription.
sub, err := s.Publisher.Subscribe(&stream.SubscribeRequest{
Topic: autopilotevents.EventTopicReadyServers,
Subject: stream.SubjectNone,
Token: token,
Index: index,
})
if err != nil {
logger.Error("failed to subscribe to server discovery events", "error", err)
return 0, status.Error(codes.Internal, "failed to subscribe to server discovery events")
}
defer sub.Unsubscribe()
for {
event, err := sub.Next(serverStream.Context())
switch {
case errors.Is(err, stream.ErrSubForceClosed):
return index, err
case errors.Is(err, context.Canceled):
return 0, nil
case err != nil:
logger.Error("failed to read next event", "error", err)
return index, status.Error(codes.Internal, err.Error())
}
// We do not send framing events (e.g. EndOfSnapshot, NewSnapshotToFollow)
// because we send a full list of ready servers on every event, rather than expecting
// clients to maintain a state-machine in the way they do for service health.
if event.IsFramingEvent() {
continue
}
// Note: this check isn't strictly necessary because the event publishing
// machinery will ensure the index increases monotonically, but it can be
// tricky to faithfully reproduce this in tests (e.g. the EventPublisher
// garbage collects topic buffers and snapshots aggressively when streams
// disconnect) so this avoids a bunch of confusing setup code.
if event.Index <= index {
continue
}
index = event.Index
rsp, err := eventToResponse(req, event)
if err != nil {
logger.Error("failed to convert event to response", "error", err)
return index, status.Error(codes.Internal, err.Error())
}
if err := serverStream.Send(rsp); err != nil {
logger.Error("failed to send response", "error", err)
return index, err
}
}
}
func eventToResponse(req *pbserverdiscovery.WatchServersRequest, event stream.Event) (*pbserverdiscovery.WatchServersResponse, error) {
readyServers, err := autopilotevents.ExtractEventPayload(event)
if err != nil {
return nil, err
}
var servers []*pbserverdiscovery.Server
for _, srv := range readyServers {
addr := srv.Address
wanAddr, ok := srv.TaggedAddresses[structs.TaggedAddressWAN]
if req.Wan && ok {
addr = wanAddr
}
servers = append(servers, &pbserverdiscovery.Server{
Id: srv.ID,
Version: srv.Version,
Address: addr,
})
}
return &pbserverdiscovery.WatchServersResponse{
Servers: servers,
}, nil
}