1
package chord
2
 
3
import (
4
	"context"
5
	"errors"
6
	"net"
7
	"net/http"
8
 
9
	"go.miragespace.co/specter/spec/chord"
10
	"go.miragespace.co/specter/spec/protocol"
11
	"go.miragespace.co/specter/spec/rpc"
12
	"go.miragespace.co/specter/spec/transport"
13
	"go.miragespace.co/specter/util"
14
	"go.miragespace.co/specter/util/ratecounter"
15
 
16
	"github.com/go-chi/chi/v5"
17
	"github.com/go-chi/chi/v5/middleware"
18
	"github.com/twitchtv/twirp"
19
	"go.uber.org/zap"
20
)
21
 
22
func (n *LocalNode) logError(ctx context.Context, err twirp.Error) context.Context {
23
	if err.Code() == twirp.FailedPrecondition || errors.Is(err, chord.ErrNodeGone) {
24
		return ctx
25
	}
26
	delegation := rpc.GetDelegation(ctx)
27
	if delegation != nil {
28
		service, _ := twirp.ServiceName(ctx)
29
		method, _ := twirp.MethodName(ctx)
30
		l := n.logger.With(
31
			zap.String("component", "rpc_server"),
32
			zap.Object("peer", delegation.Identity),
33
			zap.String("service", service),
34
			zap.String("method", method),
35
		)
36
		cause, key := rpc.GetErrorMeta(err)
37
		if cause != "" {
38
			l = l.With(zap.String("cause", cause))
39
		}
40
		if key != "" {
41
			l = l.With(zap.String("kv-key", key))
42
		}
43
		l.Error("Error handling RPC request", zap.Error(err))
44
		n.rpcErrorCount.Inc()
45
	}
46
	return ctx
47
}
48
 
49
func (n *LocalNode) getIncrementor(rate *ratecounter.Rate) func(ctx context.Context) (context.Context, error) {
50
	return func(ctx context.Context) (context.Context, error) {
51
		rate.Increment()
52
		return ctx, nil
53
	}
54
}
55
 
56
func (n *LocalNode) getRPCHandler(ctx context.Context) http.Handler {
57
	n.rpcHandlerOnce.Do(func() {
58
		r := &Server{
59
			LocalNode: n,
60
			Factory: func(node *protocol.Node) (chord.VNode, error) {
61
				return NewRemoteNode(ctx, n.BaseLogger, n.ChordClient, node)
62
			},
63
		}
64
		nsTwirp := protocol.NewVNodeServiceServer(
65
			r,
66
			twirp.WithServerHooks(&twirp.ServerHooks{
67
				RequestReceived: n.getIncrementor(n.chordRate),
68
				Error:           n.logError,
69
			}),
70
		)
71
		ksTwirp := protocol.NewKVServiceServer(
72
			r,
73
			twirp.WithServerHooks(&twirp.ServerHooks{
74
				RequestReceived: n.getIncrementor(n.kvRate),
75
				Error:           n.logError,
76
			}),
77
		)
78
 
79
		rpcHandler := chi.NewRouter()
80
		rpcHandler.Use(middleware.Recoverer)
81
		rpcHandler.Mount(nsTwirp.PathPrefix(), rpc.ExtractContext(nsTwirp))
82
		rpcHandler.Mount(ksTwirp.PathPrefix(), rpc.ExtractContext(ksTwirp))
83
		n.rpcHandler = rpcHandler
84
	})
85
 
86
	return n.rpcHandler
87
}
88
 
89
func (n *LocalNode) AttachRouter(ctx context.Context, router *transport.StreamRouter) {
90
	n.stopWg.Add(1)
91
 
92
	rpcHandler := n.getRPCHandler(ctx)
93
	srv := &http.Server{
94
		BaseContext: func(l net.Listener) context.Context {
95
			return ctx
96
		},
97
		ConnContext: func(ctx context.Context, c net.Conn) context.Context {
98
			return rpc.WithDelegation(ctx, c.(*transport.StreamDelegate))
99
		},
100
		ReadHeaderTimeout: rpcTimeout,
101
		Handler:           rpcHandler,
102
		ErrorLog:          util.GetStdLogger(n.logger, "rpc_server"),
103
	}
104
 
105
	go srv.Serve(n.rpcAcceptor)
106
	go func() {
107
		defer n.stopWg.Done()
108
 
109
		<-n.stopCh
110
		n.rpcAcceptor.Close()
111
	}()
112
 
113
	router.HandleChord(protocol.Stream_RPC, n.Identity(), func(delegate *transport.StreamDelegate) {
114
		n.rpcAcceptor.Handle(delegate)
115
	})
116
}
117
 
118
func (n *LocalNode) AttachRoot(ctx context.Context, router *transport.StreamRouter) {
119
	router.HandleChord(protocol.Stream_RPC, nil, func(delegate *transport.StreamDelegate) {
120
		n.rpcAcceptor.Handle(delegate)
121
	})
122
}
123
 
124
func (n *LocalNode) AttachExternal(ctx context.Context, listener net.Listener) {
125
	rpcHandler := n.getRPCHandler(ctx)
126
	srv := &http.Server{
127
		BaseContext: func(l net.Listener) context.Context {
128
			return ctx
129
		},
130
		ReadHeaderTimeout: rpcTimeout,
131
		Handler:           rpcHandler,
132
		ErrorLog:          util.GetStdLogger(n.logger, "rpc_external"),
133
	}
134
	go srv.Serve(listener)
135
}