func NewRaftProxyHealthServer(local HealthServer, connSelector raftpicker.Interface, cluster raftpicker.RaftCluster, ctxMod func(context.Context) (context.Context, error)) HealthServer {
redirectChecker := func(ctx context.Context) (context.Context, error) {
s, ok := transport.StreamFromContext(ctx)
if !ok {
return ctx, grpc.Errorf(codes.InvalidArgument, "remote addr is not found in context")
}
addr := s.ServerTransport().RemoteAddr().String()
md, ok := metadata.FromContext(ctx)
if ok && len(md["redirect"]) != 0 {
return ctx, grpc.Errorf(codes.ResourceExhausted, "more than one redirect to leader from: %s", md["redirect"])
}
if !ok {
md = metadata.New(map[string]string{})
}
md["redirect"] = append(md["redirect"], addr)
return metadata.NewContext(ctx, md), nil
}
mods := []func(context.Context) (context.Context, error){redirectChecker}
mods = append(mods, ctxMod)
return &raftProxyHealthServer{
local: local,
cluster: cluster,
connSelector: connSelector,
ctxMods: mods,
}
}
func NewRaftProxyResourceAllocatorServer(local ResourceAllocatorServer, connSelector raftselector.ConnProvider, localCtxMod, remoteCtxMod func(context.Context) (context.Context, error)) ResourceAllocatorServer {
redirectChecker := func(ctx context.Context) (context.Context, error) {
s, ok := transport.StreamFromContext(ctx)
if !ok {
return ctx, grpc.Errorf(codes.InvalidArgument, "remote addr is not found in context")
}
addr := s.ServerTransport().RemoteAddr().String()
md, ok := metadata.FromContext(ctx)
if ok && len(md["redirect"]) != 0 {
return ctx, grpc.Errorf(codes.ResourceExhausted, "more than one redirect to leader from: %s", md["redirect"])
}
if !ok {
md = metadata.New(map[string]string{})
}
md["redirect"] = append(md["redirect"], addr)
return metadata.NewContext(ctx, md), nil
}
remoteMods := []func(context.Context) (context.Context, error){redirectChecker}
remoteMods = append(remoteMods, remoteCtxMod)
var localMods []func(context.Context) (context.Context, error)
if localCtxMod != nil {
localMods = []func(context.Context) (context.Context, error){localCtxMod}
}
return &raftProxyResourceAllocatorServer{
local: local,
connSelector: connSelector,
localCtxMods: localMods,
remoteCtxMods: remoteMods,
}
}
// GRPCCallInfo returns an augmented context with a CallInfo structure,
// only for gRPC contexts.
func GRPCCallInfo(ctx context.Context) context.Context {
stream, ok := transport.StreamFromContext(ctx)
if !ok {
return ctx
}
return NewContext(ctx, &gRPCCallInfoImpl{
method: stream.Method(),
})
}
// SetTrailer sets the trailer metadata that will be sent when an RPC returns.
// It may be called at most once from a unary RPC handler. The ctx is the RPC
// handler's Context or one derived from it.
func SetTrailer(ctx context.Context, md metadata.MD) error {
if md.Len() == 0 {
return nil
}
stream, ok := transport.StreamFromContext(ctx)
if !ok {
return fmt.Errorf("grpc: failed to fetch the stream from the context %v", ctx)
}
return stream.SetTrailer(md)
}
// SendHeader sends header metadata. It may be called at most once from a unary
// RPC handler. The ctx is the RPC handler's Context or one derived from it.
func SendHeader(ctx context.Context, md metadata.MD) error {
if md.Len() == 0 {
return nil
}
stream, ok := transport.StreamFromContext(ctx)
if !ok {
return fmt.Errorf("grpc: failed to fetch the stream from the context %v", ctx)
}
t := stream.ServerTransport()
if t == nil {
grpclog.Fatalf("grpc: SendHeader: %v has no ServerTransport to send header metadata.", stream)
}
return t.WriteHeader(stream, md)
}
// SendHeader sends header metadata. It may be called at most once.
// The provided md and headers set by SetHeader() will be sent.
func SendHeader(ctx context.Context, md metadata.MD) error {
stream, ok := transport.StreamFromContext(ctx)
if !ok {
return Errorf(codes.Internal, "grpc: failed to fetch the stream from the context %v", ctx)
}
t := stream.ServerTransport()
if t == nil {
grpclog.Fatalf("grpc: SendHeader: %v has no ServerTransport to send header metadata.", stream)
}
if err := t.WriteHeader(stream, md); err != nil {
return toRPCErr(err)
}
return nil
}
func GetPeerInfo(s grpc.Stream) PeerInfo {
var pi PeerInfo
ctx := s.Context()
trs, ok := transport.StreamFromContext(ctx)
if ok {
pi.addr = trs.ServerTransport().RemoteAddr().String()
}
p, _ := peer.FromContext(ctx)
switch creds := p.AuthInfo.(type) {
case credentials.TLSInfo:
state := creds.State
if len(state.PeerCertificates) > 0 {
pi.cert = state.PeerCertificates[0]
}
}
return pi
}
// backendTransportStream picks and establishes a Stream to the backend.
func backendTransportStream(director StreamDirector, ctx context.Context) (transport.ClientTransport, *transport.Stream, error) {
grpcConn, err := director(ctx)
if err != nil {
if grpc.Code(err) != codes.Unknown { // rpcError check
return nil, nil, err
} else {
return nil, nil, grpc.Errorf(codes.Aborted, "cant dial to backend: %v", err)
}
}
// TODO(michal): ClientConn.GetTransport() IS NOT IN UPSTREAM GRPC!
// To make this work, copy patch/get_transport.go to google.golang.org/grpc/
backendTrans, _, err := grpcConn.GetTransport(ctx)
frontendStream, _ := transport.StreamFromContext(ctx)
callHdr := &transport.CallHdr{
Method: frontendStream.Method(),
Host: "TODOFIXTLS", // TODO(michal): This can fail if the backend server is using TLS Hostname verification. Use conn.authority, once it's public?
}
backendStream, err := backendTrans.NewStream(ctx, callHdr)
if err != nil {
return nil, nil, grpc.Errorf(codes.Unknown, "cant establish stream to backend: %v", err)
}
return backendTrans, backendStream, nil
}
// Session is a stream which controls agent connection.
// Each message contains list of backup Managers with weights. Also there is
// a special boolean field Disconnect which if true indicates that node should
// reconnect to another Manager immediately.
func (d *Dispatcher) Session(r *api.SessionRequest, stream api.Dispatcher_SessionServer) error {
ctx := stream.Context()
nodeInfo, err := ca.RemoteNode(ctx)
if err != nil {
return err
}
nodeID := nodeInfo.NodeID
if err := d.isRunningLocked(); err != nil {
return err
}
// register the node.
sessionID, err := d.register(stream.Context(), nodeID, r.Description)
if err != nil {
return err
}
fields := logrus.Fields{
"node.id": nodeID,
"node.session": sessionID,
"method": "(*Dispatcher).Session",
}
if nodeInfo.ForwardedBy != nil {
fields["forwarder.id"] = nodeInfo.ForwardedBy.NodeID
}
log := log.G(ctx).WithFields(fields)
var nodeObj *api.Node
nodeUpdates, cancel, err := store.ViewAndWatch(d.store, func(readTx store.ReadTx) error {
nodeObj = store.GetNode(readTx, nodeID)
return nil
}, state.EventUpdateNode{Node: &api.Node{ID: nodeID},
Checks: []state.NodeCheckFunc{state.NodeCheckID}},
)
if cancel != nil {
defer cancel()
}
if err != nil {
log.WithError(err).Error("ViewAndWatch Node failed")
}
if _, err = d.nodes.GetWithSession(nodeID, sessionID); err != nil {
return err
}
if err := stream.Send(&api.SessionMessage{
SessionID: sessionID,
Node: nodeObj,
Managers: d.getManagers(),
NetworkBootstrapKeys: d.networkBootstrapKeys,
}); err != nil {
return err
}
managerUpdates, mgrCancel := d.mgrQueue.Watch()
defer mgrCancel()
keyMgrUpdates, keyMgrCancel := d.keyMgrQueue.Watch()
defer keyMgrCancel()
// disconnectNode is a helper forcibly shutdown connection
disconnectNode := func() error {
// force disconnect by shutting down the stream.
transportStream, ok := transport.StreamFromContext(stream.Context())
if ok {
// if we have the transport stream, we can signal a disconnect
// in the client.
if err := transportStream.ServerTransport().Close(); err != nil {
log.WithError(err).Error("session end")
}
}
nodeStatus := api.NodeStatus{State: api.NodeStatus_DISCONNECTED, Message: "node is currently trying to find new manager"}
if err := d.nodeRemove(nodeID, nodeStatus); err != nil {
log.WithError(err).Error("failed to remove node")
}
// still return an abort if the transport closure was ineffective.
return grpc.Errorf(codes.Aborted, "node must disconnect")
}
for {
// After each message send, we need to check the nodes sessionID hasn't
// changed. If it has, we will the stream and make the node
// re-register.
node, err := d.nodes.GetWithSession(nodeID, sessionID)
if err != nil {
return err
}
var mgrs []*api.WeightedPeer
var disconnect bool
select {
case ev := <-managerUpdates:
mgrs = ev.([]*api.WeightedPeer)
case ev := <-nodeUpdates:
nodeObj = ev.(state.EventUpdateNode).Node
case <-stream.Context().Done():
return stream.Context().Err()
case <-node.Disconnect:
disconnect = true
case <-d.ctx.Done():
disconnect = true
case <-keyMgrUpdates:
}
if mgrs == nil {
mgrs = d.getManagers()
}
if err := stream.Send(&api.SessionMessage{
SessionID: sessionID,
Node: nodeObj,
Managers: mgrs,
NetworkBootstrapKeys: d.networkBootstrapKeys,
}); err != nil {
return err
}
if disconnect {
return disconnectNode()
}
}
}