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  }