server.go

package rpc

import (
	"bufio"
	"context"
	"encoding/gob"
	"errors"
	"io"
	"log"
	"net"
	"reflect"
	"strings"
	"sync"
)

// A value sent as a placeholder for the server's response value when the server
// receives an invalid request. It is never decoded by the client since the Response
// contains an error when it is used.
var invalidRequest = struct{}{}

// Server represents an RPC Server.
type Server struct {
	serviceMap  sync.Map
	reqLock     sync.Mutex // protects freeReq
	freeReq     *Request
	respLock    sync.Mutex // protects freeResp
	freeResp    *Response
	middlewares []middleware
}

// NewServer returns a new Server.
func NewServer() *Server {
	return &Server{}
}

// Register publishes in the server the set of methods of the
// receiver value that satisfy the following conditions:
//	- exported method of exported type
//	- three arguments
//  - the first argument is a context.Context
//  - the second argument is an exported type
//	- the third argument is a pointer and an exported type
//	- one return value, of type error
// It returns an error if the receiver is not an exported type or has
// no suitable methods. It also logs the error using package log.
// The client accesses each method using a string of the form "Type.Method",
// where Type is the receiver's concrete type.
func (server *Server) Register(rcvr interface{}) error {
	return server.register(rcvr, "", false)
}

// RegisterName is like Register but uses the provided name for the type
// instead of the receiver's concrete type.
func (server *Server) RegisterName(name string, rcvr interface{}) error {
	return server.register(rcvr, name, true)
}

// Accept accepts connections on the listener and serves requests
// for each incoming connection. Accept blocks until the listener
// returns a non-nil error. The caller typically invokes Accept in a
// go statement.
func (server *Server) Accept(ctx context.Context, lis net.Listener) {
	for {
		conn, err := lis.Accept()
		if err != nil {
			if debugLog && !errors.Is(err, net.ErrClosed) {
				log.Print("rpc.Serve: accept:", err.Error())
			}
			return
		}

		// attempt some optimizations for tcp
		if tcp, ok := conn.(*net.TCPConn); ok {
			_ = tcp.SetKeepAlive(true)
			_ = tcp.SetNoDelay(true)
		}

		go server.ServeConn(ctx, conn)
	}
}

// ServeConn runs the server on a single connection.
// ServeConn blocks, serving the connection until the client hangs up.
// The caller typically invokes ServeConn in a go statement.
// ServeConn uses the gob wire format (see package gob) on the
// connection. To use an alternate codec, use ServeCodec.
// See NewClient's comment for information about concurrent access.
func (server *Server) ServeConn(ctx context.Context, conn io.ReadWriteCloser) {
	buf := bufio.NewWriter(conn)
	srv := &gobServerCodec{
		rwc:    conn,
		dec:    gob.NewDecoder(conn),
		enc:    gob.NewEncoder(buf),
		encBuf: buf,
	}
	server.ServeCodec(ctx, srv)
}

// ServeCodec is like ServeConn but uses the specified codec to
// decode requests and encode responses.
func (server *Server) ServeCodec(ctx context.Context, codec ServerCodec) {
	sending := new(sync.Mutex)
	wg := new(sync.WaitGroup)

	for {
		service, mtype, req, argv, replyv, keepReading, err := server.readRequest(codec)

		if err != nil {
			if debugLog && err != io.EOF {
				log.Println("rpc:", err)
			}
			if !keepReading {
				break
			}
			// send a response if we actually managed to read a header.
			if req != nil {
				resp := server.getResponse(req)
				resp.callErr = err
				server.sendResponse(sending, resp, invalidRequest, codec)
				server.freeRequest(req)
			}
			continue
		}
		wg.Add(1)
		go service.call(ctx, server, sending, wg, mtype, req, argv, replyv, codec)
	}

	// We've seen that there are no more requests.
	// Wait for responses to be sent before closing codec.
	wg.Wait()
	codec.Close()
}

// ServeRequest is like ServeCodec but synchronously serves a single request.
// It does not close the codec upon completion.
func (server *Server) ServeRequest(ctx context.Context, codec ServerCodec) error {
	sending := new(sync.Mutex)
	service, mtype, req, argv, replyv, keepReading, err := server.readRequest(codec)

	if err != nil {
		if !keepReading {
			return err
		}

		// send a response if we actually managed to read a header.
		if req != nil {
			resp := server.getResponse(req)
			resp.callErr = err
			server.sendResponse(sending, resp, invalidRequest, codec)
			server.freeRequest(req)
		}
		return err
	}

	service.call(ctx, server, sending, nil, mtype, req, argv, replyv, codec)
	return nil
}

func (server *Server) readRequestHeader(codec ServerCodec) (svc *service, mtype *methodType, req *Request, keepReading bool, err error) {
	// Grab the request header.
	req = server.getRequest()
	err = codec.ReadRequestHeader(req)
	if err != nil {
		req = nil
		if err == io.EOF || err == io.ErrUnexpectedEOF {
			return
		}
		err = errors.New("rpc: server cannot decode request: " + err.Error())
		return
	}

	// We read the header successfully. If we see an error now,
	// we can still recover and move on to the next request.
	keepReading = true

	dot := strings.LastIndex(req.ServiceMethod, ".")
	if dot < 0 {
		err = errors.New("rpc: service/method request ill-formed: " + req.ServiceMethod)
		return
	}
	serviceName := req.ServiceMethod[:dot]
	methodName := req.ServiceMethod[dot+1:]

	// Look up the request.
	svci, ok := server.serviceMap.Load(serviceName)
	if !ok {
		err = errors.New("rpc: can't find service " + req.ServiceMethod)
		return
	}
	svc = svci.(*service)
	mtype = svc.method[methodName]
	if mtype == nil {
		err = errors.New("rpc: can't find method " + req.ServiceMethod)
	}
	return
}

func (server *Server) readRequest(codec ServerCodec) (service *service, mtype *methodType, req *Request, argv, replyv reflect.Value, keepReading bool, err error) {
	service, mtype, req, keepReading, err = server.readRequestHeader(codec)
	if err != nil {
		if !keepReading {
			return
		}

		// discard body
		_ = codec.ReadRequestBody(nil)
		return
	}

	// Decode the argument value.
	argIsValue := false // if true, need to indirect before calling.
	if mtype.ArgType.Kind() == reflect.Ptr {
		argv = reflect.New(mtype.ArgType.Elem())
	} else {
		argv = reflect.New(mtype.ArgType)
		argIsValue = true
	}

	// argv guaranteed to be a pointer now.
	if err = codec.ReadRequestBody(argv.Interface()); err != nil {
		return
	}
	if argIsValue {
		argv = argv.Elem()
	}

	replyv = reflect.New(mtype.ReplyType.Elem())

	switch mtype.ReplyType.Elem().Kind() {
	case reflect.Map:
		replyv.Elem().Set(reflect.MakeMap(mtype.ReplyType.Elem()))
	case reflect.Slice:
		replyv.Elem().Set(reflect.MakeSlice(mtype.ReplyType.Elem(), 0, 0))
	}
	return
}

func (server *Server) sendResponse(sending *sync.Mutex, resp *Response, reply interface{}, codec ServerCodec) {
	// write the Error string if we have one
	if resp.callErr != nil {
		resp.Error = resp.callErr.Error()
	}

	sending.Lock()
	err := codec.WriteResponse(resp, reply)
	if debugLog && err != nil {
		log.Println("rpc: writing response:", err)
	}
	sending.Unlock()
	server.freeResponse(resp)
}

func (server *Server) getResponse(req *Request) *Response {
	server.respLock.Lock()
	resp := server.freeResp
	if resp == nil {
		resp = new(Response)
	} else {
		server.freeResp = resp.next
		*resp = Response{}
	}

	// Encode the response header
	resp.ServiceMethod = req.ServiceMethod
	resp.ID = req.ID
	resp.Header = Header{}

	server.respLock.Unlock()
	return resp
}

func (server *Server) freeResponse(resp *Response) {
	server.respLock.Lock()
	resp.next = server.freeResp
	server.freeResp = resp
	server.respLock.Unlock()
}

func (server *Server) getRequest() *Request {
	server.reqLock.Lock()
	req := server.freeReq
	if req == nil {
		req = new(Request)
	} else {
		server.freeReq = req.next
		*req = Request{}
	}
	server.reqLock.Unlock()
	return req
}

func (server *Server) freeRequest(req *Request) {
	server.reqLock.Lock()
	req.next = server.freeReq
	server.freeReq = req
	server.reqLock.Unlock()
}

func (s *service) call(ctx context.Context, server *Server, sending *sync.Mutex, wg *sync.WaitGroup, mtype *methodType, req *Request, argv, replyv reflect.Value, codec ServerCodec) {
	ctx = context.WithValue(ctx, contextID, req.ID)
	ctx = context.WithValue(ctx, contextServiceMethod, req.ServiceMethod)
	ctx = context.WithValue(ctx, contextHeader, req.Header)

	if wg != nil {
		defer wg.Done()
	}

	mtype.Lock()
	mtype.numCalls++
	mtype.Unlock()
	function := mtype.method.Func

	resp := server.getResponse(req)

	sentResponse := false

	// create our root handler
	handler := func(ctx context.Context, rw ResponseWriter, req *Request) {
		// Invoke the method, providing a new value for the reply.
		returnValues := function.Call([]reflect.Value{s.rcvr, reflect.ValueOf(ctx), argv, replyv})

		// The return value for the method is an error.
		errInter := returnValues[0].Interface()
		if errInter != nil {
			resp.callErr = errInter.(error)
		}

		server.sendResponse(sending, resp, replyv.Interface(), codec)
		sentResponse = true
	}

	// do middleware stuff
	for i := len(server.middlewares) - 1; i >= 0; i-- {
		handler = server.middlewares[i].Middleware(handler)
	}

	handler(ctx, responseWriter{Response: resp}, req)

	// we still want to send a response, regardless of middleware firing
	if !sentResponse {
		server.sendResponse(sending, resp, replyv.Interface(), codec)
	}

	// cleanup
	server.freeRequest(req)
}