rpc/registry.go

package rpc

import (
	"fmt"
	"io"
	"reflect"
	"strings"

	"git.loafle.net/commons_go/rpc/protocol"
)

/**

Network connection
	Handshake(1..n)
		Send request
			HTTP


*/

type WriteHookFunc func(io.Writer)

// NewRPCRegistry returns a new RPC registry.
func NewRegistry() Registry {
	return &rpcRegistry{
		codecs:   make(map[string]protocol.Codec),
		services: new(serviceMap),
	}
}

type Registry interface {
	RegisterCodec(codec protocol.Codec, contentType string)
	RegisterService(receiver interface{}, name string) error
	HasMethod(method string) bool
	Invoke(contentType string, reader io.Reader, writer io.Writer, beforeWrite WriteHookFunc, afterWrite WriteHookFunc) error
}

// RPCRegistry serves registered RPC services using registered codecs.
type rpcRegistry struct {
	codecs   map[string]protocol.Codec
	services *serviceMap
}

// RegisterCodec adds a new codec to the server.
//
// Codecs are defined to process a given serialization scheme, e.g., JSON or
// XML. A codec is chosen based on the "Content-Type" header from the request,
// excluding the charset definition.
func (rr *rpcRegistry) RegisterCodec(codec protocol.Codec, contentType string) {
	rr.codecs[strings.ToLower(contentType)] = codec
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has two arguments: *args, *reply.
//    - All two arguments are pointers.
//    - The first and second arguments are exported or local.
//    - The method has return type error.
//
// All other methods are ignored.
func (rr *rpcRegistry) RegisterService(receiver interface{}, name string) error {
	return rr.services.register(receiver, name)
}

// HasMethod returns true if the given method is registered.
//
// The method uses a dotted notation as in "Service.Method".
func (rr *rpcRegistry) HasMethod(method string) bool {
	if _, _, err := rr.services.get(method); err == nil {
		return true
	}
	return false
}

// Invoke execute a method.
//
// Codecs are defined to process a given serialization scheme, e.g., JSON or
// XML. A codec is chosen based on the "Content-Type" header from the request,
// excluding the charset definition.
func (rr *rpcRegistry) Invoke(contentType string, r io.Reader, w io.Writer, beforeWrite WriteHookFunc, afterWrite WriteHookFunc) error {
	var codec protocol.Codec
	if contentType == "" && len(rr.codecs) == 1 {
		// If Content-Type is not set and only one codec has been registered,
		// then default to that codec.
		for _, c := range rr.codecs {
			codec = c
		}
	} else if codec = rr.codecs[strings.ToLower(contentType)]; codec == nil {
		return fmt.Errorf("Unrecognized Content-Type: %s", contentType)
	}

	// Create a new codec request.
	codecReq, errNew := codec.NewRequest(r)
	if nil != errNew {
		return errNew
	}
	// Get service method to be called.
	method, errMethod := codecReq.Method()
	if errMethod != nil {
		return write(codecReq, w, beforeWrite, afterWrite, nil, errMethod)
	}
	serviceSpec, methodSpec, errGet := rr.services.get(method)
	if errGet != nil {
		return write(codecReq, w, beforeWrite, afterWrite, nil, errGet)
	}
	// Decode the args.
	args := reflect.New(methodSpec.argsType)
	if errRead := codecReq.ReadRequest(args.Interface()); errRead != nil {
		return write(codecReq, w, beforeWrite, afterWrite, nil, errRead)
	}
	// Call the service method.
	reply := reflect.New(methodSpec.replyType)
	errValue := methodSpec.method.Func.Call([]reflect.Value{
		serviceSpec.rcvr,
		args,
		reply,
	})

	// Cast the result to error if needed.
	var errResult error
	errInter := errValue[0].Interface()
	if errInter != nil {
		errResult = errInter.(error)
	}

	if errResult != nil {
		return write(codecReq, w, beforeWrite, afterWrite, nil, errResult)
	}

	return write(codecReq, w, beforeWrite, afterWrite, reply.Interface(), nil)
}

func write(codecReq protocol.CodecRequest, w io.Writer, beforeWrite WriteHookFunc, afterWrite WriteHookFunc, result interface{}, err error) error {
	if nil != beforeWrite {
		beforeWrite(w)
	}

	var wErr error

	if err == nil {
		wErr = codecReq.WriteResponse(w, result)
	} else {
		wErr = codecReq.WriteError(w, 400, err)
	}

	if nil != wErr {
		return wErr
	}

	if nil != afterWrite {
		afterWrite(w)
	}

	return nil
}