// Copyright 2009 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 geerpc import ( "bufio" "context" "encoding/json" "errors" "fmt" "geerpc/codec" "io" "log" "net" "net/http" "strings" "sync" "time" ) // Call represents an active RPC. type Call struct { Seq uint64 ServiceMethod string // format "." Args interface{} // arguments to the function Reply interface{} // reply from the function Error error // if error occurs, it will be set Done chan *Call // Strobes when call is complete. } func (call *Call) done() { call.Done <- call } // Client represents an RPC Client. // There may be multiple outstanding Calls associated // with a single Client, and a Client may be used by // multiple goroutines simultaneously. type Client struct { cc codec.Codec opt *Option sending sync.Mutex // protect following header codec.Header mu sync.Mutex // protect following seq uint64 pending map[uint64]*Call closing bool // user has called Close shutdown bool // server has told us to stop } var _ io.Closer = (*Client)(nil) var ErrShutdown = errors.New("connection is shut down") // Close the connection func (client *Client) Close() error { client.mu.Lock() defer client.mu.Unlock() if client.closing { return ErrShutdown } client.closing = true return client.cc.Close() } // IsAvailable return true if the client does work func (client *Client) IsAvailable() bool { client.mu.Lock() defer client.mu.Unlock() return !client.shutdown && !client.closing } func (client *Client) registerCall(call *Call) (uint64, error) { client.mu.Lock() defer client.mu.Unlock() if client.closing || client.shutdown { return 0, ErrShutdown } call.Seq = client.seq client.pending[call.Seq] = call client.seq++ return call.Seq, nil } func (client *Client) removeCall(seq uint64) *Call { client.mu.Lock() defer client.mu.Unlock() call := client.pending[seq] delete(client.pending, seq) return call } func (client *Client) terminateCalls(err error) { client.sending.Lock() defer client.sending.Unlock() client.mu.Lock() defer client.mu.Unlock() client.shutdown = true for _, call := range client.pending { call.Error = err call.done() } } func (client *Client) send(call *Call) { // make sure that the client will send a complete request client.sending.Lock() defer client.sending.Unlock() // register this call. seq, err := client.registerCall(call) if err != nil { call.Error = err call.done() return } // prepare request header client.header.ServiceMethod = call.ServiceMethod client.header.Seq = seq client.header.Error = "" // encode and send the request if err := client.cc.Write(&client.header, call.Args); err != nil { call := client.removeCall(seq) // call may be nil, it usually means that Write partially failed, // client has received the response and handled if call != nil { call.Error = err call.done() } } } func (client *Client) receive() { var err error for err == nil { var h codec.Header if err = client.cc.ReadHeader(&h); err != nil { break } call := client.removeCall(h.Seq) switch { case call == nil: // it usually means that Write partially failed // and call was already removed. err = client.cc.ReadBody(nil) case h.Error != "": call.Error = fmt.Errorf(h.Error) err = client.cc.ReadBody(nil) call.done() default: err = client.cc.ReadBody(call.Reply) if err != nil { call.Error = errors.New("reading body " + err.Error()) } call.done() } } // error occurs, so terminateCalls pending calls client.terminateCalls(err) } // Go invokes the function asynchronously. // It returns the Call structure representing the invocation. func (client *Client) Go(serviceMethod string, args, reply interface{}, done chan *Call) *Call { if done == nil { done = make(chan *Call, 10) } else if cap(done) == 0 { log.Panic("rpc client: done channel is unbuffered") } call := &Call{ ServiceMethod: serviceMethod, Args: args, Reply: reply, Done: done, } client.send(call) return call } // Call invokes the named function, waits for it to complete, // and returns its error status. func (client *Client) Call(ctx context.Context, serviceMethod string, args, reply interface{}) error { call := client.Go(serviceMethod, args, reply, make(chan *Call, 1)) select { case <-ctx.Done(): client.removeCall(call.Seq) return errors.New("rpc client: call failed: " + ctx.Err().Error()) case call := <-call.Done: return call.Error } } func parseOptions(opts ...*Option) (*Option, error) { // if opts is nil or pass nil as parameter if len(opts) == 0 || opts[0] == nil { return DefaultOption, nil } if len(opts) != 1 { return nil, errors.New("number of options is more than 1") } opt := opts[0] opt.MagicNumber = DefaultOption.MagicNumber if opt.CodecType == "" { opt.CodecType = DefaultOption.CodecType } return opt, nil } func NewClient(conn net.Conn, opt *Option) (*Client, error) { f := codec.NewCodecFuncMap[opt.CodecType] if f == nil { err := fmt.Errorf("invalid codec type %s", opt.CodecType) log.Println("rpc client: codec error:", err) return nil, err } // send options with server if err := json.NewEncoder(conn).Encode(opt); err != nil { log.Println("rpc client: options error: ", err) _ = conn.Close() return nil, err } return newClientCodec(f(conn), opt), nil } func newClientCodec(cc codec.Codec, opt *Option) *Client { client := &Client{ seq: 1, // seq starts with 1, 0 means invalid call cc: cc, opt: opt, pending: make(map[uint64]*Call), } go client.receive() return client } type clientResult struct { client *Client err error } type newClientFunc func(conn net.Conn, opt *Option) (client *Client, err error) func dialTimeout(f newClientFunc, network, address string, opts ...*Option) (client *Client, err error) { opt, err := parseOptions(opts...) if err != nil { return nil, err } conn, err := net.DialTimeout(network, address, opt.ConnectTimeout) if err != nil { return nil, err } // close the connection if client is nil defer func() { if err != nil { _ = conn.Close() } }() ch := make(chan clientResult) go func() { client, err := f(conn, opt) ch <- clientResult{client: client, err: err} }() if opt.ConnectTimeout == 0 { result := <-ch return result.client, result.err } select { case <-time.After(opt.ConnectTimeout): return nil, fmt.Errorf("rpc client: connect timeout: expect within %s", opt.ConnectTimeout) case result := <-ch: return result.client, result.err } } // Dial connects to an RPC server at the specified network address func Dial(network, address string, opts ...*Option) (*Client, error) { return dialTimeout(NewClient, network, address, opts...) } // NewHTTPClient new a Client instance via HTTP as transport protocol func NewHTTPClient(conn net.Conn, opt *Option) (*Client, error) { _, _ = io.WriteString(conn, fmt.Sprintf("CONNECT %s HTTP/1.0\n\n", defaultRPCPath)) // Require successful HTTP response // before switching to RPC protocol. resp, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: "CONNECT"}) if err == nil && resp.Status == connected { return NewClient(conn, opt) } if err == nil { err = errors.New("unexpected HTTP response: " + resp.Status) } return nil, err } // DialHTTP connects to an HTTP RPC server at the specified network address // listening on the default HTTP RPC path. func DialHTTP(network, address string, opts ...*Option) (*Client, error) { return dialTimeout(NewHTTPClient, network, address, opts...) } // XDial calls different functions to connect to a RPC server // according the first parameter rpcAddr. // rpcAddr is a general format (protocol@addr) to represent a rpc server // eg, http@10.0.0.1:7001, tcp@10.0.0.1:9999, unix@/tmp/geerpc.sock func XDial(rpcAddr string, opts ...*Option) (*Client, error) { parts := strings.Split(rpcAddr, "@") if len(parts) != 2 { return nil, fmt.Errorf("rpc client err: wrong format '%s', expect protocol@addr", rpcAddr) } protocol, addr := parts[0], parts[1] switch protocol { case "http": return DialHTTP("tcp", addr, opts...) default: // tcp, unix or other transport protocol return Dial(protocol, addr, opts...) } }