-
Notifications
You must be signed in to change notification settings - Fork 298
/
conn.go
212 lines (189 loc) · 5.92 KB
/
conn.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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package jsonrpc2
import (
"context"
"encoding/json"
"fmt"
"sync"
"sync/atomic"
"github.com/rs/zerolog/log"
)
// Conn is the common interface to jsonrpc clients and servers.
// Conn is bidirectional; it does not have a designated server or client end.
// It manages the jsonrpc2 protocol, connecting responses back to their calls.
type Conn interface {
// Call invokes the target method and waits for a response.
// The params will be marshaled to JSON before sending over the wire, and will
// be handed to the method invoked.
// The response will be unmarshaled from JSON into the result.
// The id returned will be unique from this connection, and can be used for
// logging or tracking.
Call(ctx context.Context, method string, params, result interface{}) (ID, error)
// Notify invokes the target method but does not wait for a response.
// The params will be marshaled to JSON before sending over the wire, and will
// be handed to the method invoked.
Notify(ctx context.Context, method string, params interface{}) error
// Go starts a goroutine to handle the connection.
// It must be called exactly once for each Conn.
// It returns immediately.
// You must block on Done() to wait for the connection to shut down.
// This is a temporary measure, this should be started automatically in the
// future.
Go(ctx context.Context, handler Handler)
// Close closes the connection and it's underlying stream.
// It does not wait for the close to complete, use the Done() channel for
// that.
Close() error
// Done returns a channel that will be closed when the processing goroutine
// has terminated, which will happen if Close() is called or an underlying
// stream is closed.
Done() <-chan struct{}
// Err returns an error if there was one from within the processing goroutine.
// If err returns non nil, the connection will be already closed or closing.
Err() error
}
type conn struct {
seq int64 // must only be accessed using atomic operations
writeMu sync.Mutex // protects writes to the stream
stream Stream
pendingMu sync.Mutex // protects the pending map
pending map[ID]chan *Response
done chan struct{}
err atomic.Value
}
// NewConn creates a new connection object around the supplied stream.
func NewConn(s Stream) Conn {
conn := &conn{
stream: s,
pending: make(map[ID]chan *Response),
done: make(chan struct{}),
}
return conn
}
func (c *conn) Notify(ctx context.Context, method string, params interface{}) (err error) {
notify, err := NewNotification(method, params)
if err != nil {
return fmt.Errorf("marshaling notify parameters: %v", err)
}
_, err = c.write(ctx, notify)
return err
}
func (c *conn) Call(ctx context.Context, method string, params, result interface{}) (_ ID, err error) {
// generate a new request identifier
id := ID{number: atomic.AddInt64(&c.seq, 1)}
call, err := NewCall(id, method, params)
if err != nil {
return id, fmt.Errorf("marshaling call parameters: %v", err)
}
// We have to add ourselves to the pending map before we send, otherwise we
// are racing the response. Also add a buffer to rchan, so that if we get a
// wire response between the time this call is cancelled and id is deleted
// from c.pending, the send to rchan will not block.
rchan := make(chan *Response, 1)
c.pendingMu.Lock()
c.pending[id] = rchan
c.pendingMu.Unlock()
defer func() {
c.pendingMu.Lock()
delete(c.pending, id)
c.pendingMu.Unlock()
}()
// now we are ready to send
_, err = c.write(ctx, call)
if err != nil {
// sending failed, we will never get a response, so don't leave it pending
return id, err
}
// now wait for the response
select {
case response := <-rchan:
// is it an error response?
if response.err != nil {
return id, response.err
}
if result == nil || len(response.result) == 0 {
return id, nil
}
if err := json.Unmarshal(response.result, result); err != nil {
return id, fmt.Errorf("unmarshaling result: %v", err)
}
return id, nil
case <-ctx.Done():
return id, ctx.Err()
}
}
func (c *conn) replier(req Request) Replier {
return func(ctx context.Context, result interface{}, err error) error {
call, ok := req.(*Call)
if !ok {
// request was a notify, no need to respond
return nil
}
response, err := NewResponse(call.id, result, err)
if err != nil {
return err
}
_, err = c.write(ctx, response)
if err != nil {
// TODO(iancottrell): if a stream write fails, we really need to shut down
// the whole stream
return err
}
return nil
}
}
func (c *conn) write(ctx context.Context, msg Message) (int64, error) {
c.writeMu.Lock()
defer c.writeMu.Unlock()
return c.stream.Write(ctx, msg)
}
func (c *conn) Go(ctx context.Context, handler Handler) {
go c.run(ctx, handler)
}
func (c *conn) run(ctx context.Context, handler Handler) {
defer close(c.done)
for {
// get the next message
msg, _, err := c.stream.Read(ctx)
if err != nil {
// The stream failed, we cannot continue.
c.fail(err)
return
}
switch msg := msg.(type) {
case Request:
if err := handler(ctx, c.replier(msg), msg); err != nil {
// delivery failed, not much we can do
log.Error().Err(err).Msg("jsonrpc2: message delivery failed")
}
case *Response:
// If method is not set, this should be a response, in which case we must
// have an id to send the response back to the caller.
c.pendingMu.Lock()
rchan, ok := c.pending[msg.id]
c.pendingMu.Unlock()
if ok {
rchan <- msg
}
}
}
}
func (c *conn) Close() error {
return c.stream.Close()
}
func (c *conn) Done() <-chan struct{} {
return c.done
}
func (c *conn) Err() error {
if err := c.err.Load(); err != nil {
return err.(error)
}
return nil
}
// fail sets a failure condition on the stream and closes it.
func (c *conn) fail(err error) {
c.err.Store(err)
c.stream.Close()
}