server.go

package rpch

import (
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net"
	"reflect"
	"sync"
	"time"
)

type Server struct {
	ReadTimeOut  time.Duration
	WriteTimeOut time.Duration
	services     sync.Map
}

var DefaultServer = NewServer()

func NewServer() *Server {
	return &Server{
		ReadTimeOut:  10 * time.Second,
		WriteTimeOut: 10 * time.Second,
	}
}

func (svr *Server) ListenAndServe(addr string) error {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	defer l.Close()
	return svr.Serve(l)
}

func (svr *Server) Serve(l net.Listener) error {
	var tempDelay time.Duration
	for {
		rwc, err := l.Accept()
		if err != nil {
			if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				time.Sleep(tempDelay)
				continue
			}
			return err
		}
		tempDelay = 0
		c := newConn(svr, rwc)
		go func() {
			err := svr.handleConn(c)
			if err != nil && err != io.EOF {
				log.Println(err)
			}
		}()
	}
}

func (svr *Server) handleConn(conn *conn) error {
	var err error
	var req *request
	defer func() {
		if e := recover(); e != nil {
			log.Printf("err recovered, err=%v\n", e)
		}
		conn.close()
	}()
	if err = conn.setReadDeadline(); err != nil {
		return err
	}
	_magic, err := conn.readMagic()
	if err != nil {
		return err
	}
	if magic != _magic {
		return errInvalidMagic
	}
	for {
		req, err = conn.readRequest()
		if err != nil {
			return err
		}
		if err = svr.handleRequest(req); err != nil {
			return err
		}
		if err = req.finishRequest(); err != nil {
			return err
		}
	}
}

func (svr *Server) handleRequest(req *request) error {
	iservice, ok := svr.services.Load(req.service)
	if !ok {
		return errBadRequestService
	}
	service := iservice.(*Service)
	methodDesc, ok := service.Methods[req.method]
	if !ok {
		return errBadRequestMethod
	}
	//NumIn还包括receiver这个参数，所以+1
	if methodDesc.MethodType.NumIn() != int(req.argCnt)+1 {
		return errBadRequestArgCnt
	}
	values, err := req.parseArgs()
	if err != nil {
		return err
	}
	rtns := methodDesc.Method.Call(values)
	if req.streamingArg != nil {
		//consume the rest data in istream if user doesn't do that in handler
		//otherwise it will affect the parse of the next request
		if r := req.streamingArg.streamReader; r != nil {
			io.Copy(ioutil.Discard, r)
		}
		// if stream is a ostream, we need to make w(chunkWriter) send an EOF signal to client after
		// handler, which indicates that there are no more data to be written to ostream.
		// Only by this, can client know it's time to accept Return Value of registered methods
		if w := req.streamingArg.streamWriter; w != nil {
			//it will send 0\r\n\r\n
			w.Write(nil)
		}
	}
	return req.conn.sendResponse(rtns, methodDesc, req.seq)
}

// A valid method should have at least one and at most three return values.
// The last return value must be an error. when error is not nil, then only the
// error will be sent to the client.

// when there is a stream type in return values, then the number of return
// values of this method must be three which is stream, a callback funciton
// and an error respectively.
// This callback will be called when stream finishs being read or writen
// For instance, you open a file in your handler and return this *os.File
// You shouldn't close this file in handler, or client will not be able to read
// or write this stream. So to avoid resource leak, you can put (*os.File).Close
// in this callback function.
func checkServiceValidation(service *Service) error {
	for _, methodDesc := range service.Methods {
		f := methodDesc.MethodType
		out := f.NumOut()
		if out == 0 {
			return errors.New("Registered method should have at least one return value")
		}
		lastType := f.Out(out - 1)
		ierror := reflect.TypeOf((*error)(nil)).Elem()
		if !lastType.Implements(ierror) {
			return errors.New("Last return value should implement interface error")
		}
		if out > 3 {
			return errors.New("Registered method should have at most 2 return values besides error")
		}
		if out == 3 {
			secondType := f.Out(1)
			if secondType.AssignableTo(reflect.TypeOf(func() {})) {
				continue
			}
			return errors.New("The type of second return value must be func() type")
		}
	}
	return nil
}

func (svr *Server) Register(service *Service) {
	if service == nil {
		panic("register a nil service")
	}
	if err := checkServiceValidation(service); err != nil {
		//use goroutine, protect panic from getting recovered
		//this happens when developers edit and modify the code automatically generated by hgen
		go func() {
			panic(err)
		}()
	}
	svr.services.Store(service.Name, service)
}

func (svr *Server) UnRegister(serviceName string) {
	svr.services.Delete(serviceName)
}

func UnRegister(serviceName string) {
	DefaultServer.UnRegister(serviceName)
}

func Register(service *Service) {
	DefaultServer.Register(service)
}

func Serve(l net.Listener) error {
	return DefaultServer.Serve(l)
}

func ListenAndServe(addr string) error {
	return DefaultServer.ListenAndServe(addr)
}

func Go(f func()) {
	go func() {
		defer func() {
			if err := recover(); err != nil {
				log.Printf("panic recovered err: %v\n", err)
			}
		}()
		f()
	}()
}