src.elv.sh@v0.21.0-dev.0.20240515223629-06979efb9a2a/pkg/rpc/client.go (about) 1 // Copyright 2009 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 package rpc 6 7 import ( 8 "bufio" 9 "encoding/gob" 10 "errors" 11 "io" 12 "log" 13 "net" 14 "sync" 15 ) 16 17 // ServerError represents an error that has been returned from 18 // the remote side of the RPC connection. 19 type ServerError string 20 21 func (e ServerError) Error() string { 22 return string(e) 23 } 24 25 var ErrShutdown = errors.New("connection is shut down") 26 27 // Call represents an active RPC. 28 type Call struct { 29 ServiceMethod string // The name of the service and method to call. 30 Args any // The argument to the function (*struct). 31 Reply any // The reply from the function (*struct). 32 Error error // After completion, the error status. 33 Done chan *Call // Receives *Call when Go is complete. 34 } 35 36 // Client represents an RPC Client. 37 // There may be multiple outstanding Calls associated 38 // with a single Client, and a Client may be used by 39 // multiple goroutines simultaneously. 40 type Client struct { 41 codec ClientCodec 42 43 reqMutex sync.Mutex // protects following 44 request Request 45 46 mutex sync.Mutex // protects following 47 seq uint64 48 pending map[uint64]*Call 49 closing bool // user has called Close 50 shutdown bool // server has told us to stop 51 } 52 53 // A ClientCodec implements writing of RPC requests and 54 // reading of RPC responses for the client side of an RPC session. 55 // The client calls WriteRequest to write a request to the connection 56 // and calls ReadResponseHeader and ReadResponseBody in pairs 57 // to read responses. The client calls Close when finished with the 58 // connection. ReadResponseBody may be called with a nil 59 // argument to force the body of the response to be read and then 60 // discarded. 61 // See NewClient's comment for information about concurrent access. 62 type ClientCodec interface { 63 WriteRequest(*Request, any) error 64 ReadResponseHeader(*Response) error 65 ReadResponseBody(any) error 66 67 Close() error 68 } 69 70 func (client *Client) send(call *Call) { 71 client.reqMutex.Lock() 72 defer client.reqMutex.Unlock() 73 74 // Register this call. 75 client.mutex.Lock() 76 if client.shutdown || client.closing { 77 client.mutex.Unlock() 78 call.Error = ErrShutdown 79 call.done() 80 return 81 } 82 seq := client.seq 83 client.seq++ 84 client.pending[seq] = call 85 client.mutex.Unlock() 86 87 // Encode and send the request. 88 client.request.Seq = seq 89 client.request.ServiceMethod = call.ServiceMethod 90 err := client.codec.WriteRequest(&client.request, call.Args) 91 if err != nil { 92 client.mutex.Lock() 93 call = client.pending[seq] 94 delete(client.pending, seq) 95 client.mutex.Unlock() 96 if call != nil { 97 call.Error = err 98 call.done() 99 } 100 } 101 } 102 103 func (client *Client) input() { 104 var err error 105 var response Response 106 for err == nil { 107 response = Response{} 108 err = client.codec.ReadResponseHeader(&response) 109 if err != nil { 110 break 111 } 112 seq := response.Seq 113 client.mutex.Lock() 114 call := client.pending[seq] 115 delete(client.pending, seq) 116 client.mutex.Unlock() 117 118 switch { 119 case call == nil: 120 // We've got no pending call. That usually means that 121 // WriteRequest partially failed, and call was already 122 // removed; response is a server telling us about an 123 // error reading request body. We should still attempt 124 // to read error body, but there's no one to give it to. 125 err = client.codec.ReadResponseBody(nil) 126 if err != nil { 127 err = errors.New("reading error body: " + err.Error()) 128 } 129 case response.Error != "": 130 // We've got an error response. Give this to the request; 131 // any subsequent requests will get the ReadResponseBody 132 // error if there is one. 133 call.Error = ServerError(response.Error) 134 err = client.codec.ReadResponseBody(nil) 135 if err != nil { 136 err = errors.New("reading error body: " + err.Error()) 137 } 138 call.done() 139 default: 140 err = client.codec.ReadResponseBody(call.Reply) 141 if err != nil { 142 call.Error = errors.New("reading body " + err.Error()) 143 } 144 call.done() 145 } 146 } 147 // Terminate pending calls. 148 client.reqMutex.Lock() 149 client.mutex.Lock() 150 client.shutdown = true 151 closing := client.closing 152 if err == io.EOF { 153 if closing { 154 err = ErrShutdown 155 } else { 156 err = io.ErrUnexpectedEOF 157 } 158 } 159 for _, call := range client.pending { 160 call.Error = err 161 call.done() 162 } 163 client.mutex.Unlock() 164 client.reqMutex.Unlock() 165 if debugLog && err != io.EOF && !closing { 166 log.Println("rpc: client protocol error:", err) 167 } 168 } 169 170 func (call *Call) done() { 171 select { 172 case call.Done <- call: 173 // ok 174 default: 175 // We don't want to block here. It is the caller's responsibility to make 176 // sure the channel has enough buffer space. See comment in Go(). 177 if debugLog { 178 log.Println("rpc: discarding Call reply due to insufficient Done chan capacity") 179 } 180 } 181 } 182 183 // NewClient returns a new Client to handle requests to the 184 // set of services at the other end of the connection. 185 // It adds a buffer to the write side of the connection so 186 // the header and payload are sent as a unit. 187 // 188 // The read and write halves of the connection are serialized independently, 189 // so no interlocking is required. However each half may be accessed 190 // concurrently so the implementation of conn should protect against 191 // concurrent reads or concurrent writes. 192 func NewClient(conn io.ReadWriteCloser) *Client { 193 encBuf := bufio.NewWriter(conn) 194 client := &gobClientCodec{conn, gob.NewDecoder(conn), gob.NewEncoder(encBuf), encBuf} 195 return NewClientWithCodec(client) 196 } 197 198 // NewClientWithCodec is like NewClient but uses the specified 199 // codec to encode requests and decode responses. 200 func NewClientWithCodec(codec ClientCodec) *Client { 201 client := &Client{ 202 codec: codec, 203 pending: make(map[uint64]*Call), 204 } 205 go client.input() 206 return client 207 } 208 209 type gobClientCodec struct { 210 rwc io.ReadWriteCloser 211 dec *gob.Decoder 212 enc *gob.Encoder 213 encBuf *bufio.Writer 214 } 215 216 func (c *gobClientCodec) WriteRequest(r *Request, body any) (err error) { 217 if err = c.enc.Encode(r); err != nil { 218 return 219 } 220 if err = c.enc.Encode(body); err != nil { 221 return 222 } 223 return c.encBuf.Flush() 224 } 225 226 func (c *gobClientCodec) ReadResponseHeader(r *Response) error { 227 return c.dec.Decode(r) 228 } 229 230 func (c *gobClientCodec) ReadResponseBody(body any) error { 231 return c.dec.Decode(body) 232 } 233 234 func (c *gobClientCodec) Close() error { 235 return c.rwc.Close() 236 } 237 238 // Dial connects to an RPC server at the specified network address. 239 func Dial(network, address string) (*Client, error) { 240 conn, err := net.Dial(network, address) 241 if err != nil { 242 return nil, err 243 } 244 return NewClient(conn), nil 245 } 246 247 // Close calls the underlying codec's Close method. If the connection is already 248 // shutting down, ErrShutdown is returned. 249 func (client *Client) Close() error { 250 client.mutex.Lock() 251 if client.closing { 252 client.mutex.Unlock() 253 return ErrShutdown 254 } 255 client.closing = true 256 client.mutex.Unlock() 257 return client.codec.Close() 258 } 259 260 // Go invokes the function asynchronously. It returns the Call structure representing 261 // the invocation. The done channel will signal when the call is complete by returning 262 // the same Call object. If done is nil, Go will allocate a new channel. 263 // If non-nil, done must be buffered or Go will deliberately crash. 264 func (client *Client) Go(serviceMethod string, args any, reply any, done chan *Call) *Call { 265 call := new(Call) 266 call.ServiceMethod = serviceMethod 267 call.Args = args 268 call.Reply = reply 269 if done == nil { 270 done = make(chan *Call, 10) // buffered. 271 } else { 272 // If caller passes done != nil, it must arrange that 273 // done has enough buffer for the number of simultaneous 274 // RPCs that will be using that channel. If the channel 275 // is totally unbuffered, it's best not to run at all. 276 if cap(done) == 0 { 277 log.Panic("rpc: done channel is unbuffered") 278 } 279 } 280 call.Done = done 281 client.send(call) 282 return call 283 } 284 285 // Call invokes the named function, waits for it to complete, and returns its error status. 286 func (client *Client) Call(serviceMethod string, args any, reply any) error { 287 call := <-client.Go(serviceMethod, args, reply, make(chan *Call, 1)).Done 288 return call.Error 289 }