container-go/servlet/rpc-servlet.go

package servlet

import (
	"net"

	"git.loafle.net/commons/logging-go"
	crp "git.loafle.net/commons/rpc-go/protocol"
	crpj "git.loafle.net/commons/rpc-go/protocol/json"
	crr "git.loafle.net/commons/rpc-go/registry"
	"git.loafle.net/commons/server-go"
	css "git.loafle.net/commons/server-go/socket"
	cssn "git.loafle.net/commons/server-go/socket/net"
)

type RPCServlet interface {
	cssn.Servlet
}

type RPCServlets struct {
	cssn.Servlets

	RPCInvoker   crr.RPCInvoker
	RPCWriteChan <-chan []byte
}

func (s *RPCServlets) Handshake(servletCtx server.ServletCtx, conn net.Conn) error {
	return nil
}

func (s *RPCServlets) OnConnect(servletCtx server.ServletCtx, conn css.Conn) {
	s.Servlets.OnConnect(servletCtx, conn)

}

func (s *RPCServlets) OnDisconnect(servletCtx server.ServletCtx) {
	s.Servlets.OnDisconnect(servletCtx)
}

func (s *RPCServlets) Handle(servletCtx server.ServletCtx,
	stopChan <-chan struct{}, doneChan chan<- struct{},
	readChan <-chan []byte, writeChan chan<- []byte) {
	defer func() {
		doneChan <- struct{}{}
	}()

	var (
		src       crp.ServerRequestCodec
		reply     interface{}
		replyBuff []byte
		err       error
	)
	sc := crpj.NewServerCodec()
	go s.handleRPCWrite(stopChan, writeChan)

	for {
		select {
		case msg, ok := <-readChan:
			if !ok {
				return
			}
			// grpc exec method call
			src, err = sc.NewRequest(msg)
			if nil != err {
				logging.Logger().Error(err)
				break
			}

			reply, err = s.RPCInvoker.Invoke(src)
			replyBuff, err = src.NewResponse(reply, err)
			if nil != err {
				logging.Logger().Error(err)
				s.writeError(src, writeChan, crp.E_INTERNAL, "", err)
				break
			}

			writeChan <- replyBuff
		case <-stopChan:
			return
		}
	}
}

func (s *RPCServlets) handleRPCWrite(stopChan <-chan struct{}, writeChan chan<- []byte) {
	for {
		select {
		case buf, ok := <-s.RPCWriteChan:
			if !ok {
				break
			}

			writeChan <- buf
		case <-stopChan:
			return
		}
	}
}

func (s *RPCServlets) writeError(src crp.ServerRequestCodec, writeChan chan<- []byte, code crp.ErrorCode, message string, data interface{}) {
	if !src.HasResponse() {
		return
	}

	pErr := &crp.Error{
		Code:    code,
		Message: message,
		Data:    data,
	}

	buf, err := src.NewResponse(nil, pErr)
	if nil != err {
		logging.Logger().Error(err)
		return
	}
	writeChan <- buf
}