probe/probe/container-servlet.go

package probe

import (
	"fmt"

	logging "git.loafle.net/commons/logging-go"
	crp "git.loafle.net/commons/rpc-go/protocol"
	crr "git.loafle.net/commons/rpc-go/registry"
	"git.loafle.net/commons/server-go"
	css "git.loafle.net/commons/server-go/socket"
	cssw "git.loafle.net/commons/server-go/socket/web"
	cuc "git.loafle.net/commons/util-go/context"
	occc "git.loafle.net/overflow/commons-go/config/container"
	occp "git.loafle.net/overflow/commons-go/config/probe"
	"git.loafle.net/overflow/probe/container"
	"github.com/valyala/fasthttp"
)

var ContainerTypeKey = cuc.ContextKey("ContainerType")

type ContainerServlet interface {
	cssw.Servlet
}

type ContainerServlets struct {
	cssw.Servlets

	RPCInvoker     crr.RPCInvoker
	RPCServerCodec crp.ServerCodec

	OnConnectChan    chan<- *container.OnConnectInfo
	OnDisconnectChan chan<- occp.ContainerType
}

func (s *ContainerServlets) Init(serverCtx server.ServerCtx) error {
	if err := s.Servlets.Init(serverCtx); nil != err {
		return err
	}

	return nil
}

func (s *ContainerServlets) OnStart(serverCtx server.ServerCtx) error {
	if err := s.Servlets.OnStart(serverCtx); nil != err {
		return err
	}

	return nil
}

func (s *ContainerServlets) OnStop(serverCtx server.ServerCtx) {

	s.Servlets.OnStop(serverCtx)
}

func (s *ContainerServlets) Destroy(serverCtx server.ServerCtx) {

	s.Servlets.Destroy(serverCtx)
}

func (s *ContainerServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) {
	bMethod := ctx.Request.Header.Peek(occc.HTTPRequestHeaderKey_Container_Method)
	if nil == bMethod {
		return nil, fmt.Errorf("Unexpected probe method: %v", bMethod)
	}

	method := string(bMethod)

	switch method {
	case occc.HTTPRequestHeaderValue_Container_Method_Connect:
	default:
		return nil, fmt.Errorf("Unexpected probe method: %s", method)
	}

	bContainerType := ctx.Request.Header.Peek(occc.HTTPRequestHeaderKey_Container_Type)
	if nil == bContainerType {
		return nil, fmt.Errorf("Unexpected Container type : %v", bContainerType)
	}

	containerType := string(bContainerType)

	servletCtx.SetAttribute(ContainerTypeKey, containerType)

	return nil, nil
}

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

}

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

}

func (s *ContainerServlets) Handle(servletCtx server.ServletCtx,
	stopChan <-chan struct{}, doneChan chan<- struct{},
	readChan <-chan css.SocketMessage, writeChan chan<- css.SocketMessage) {

	_containerType := servletCtx.GetAttribute(ContainerTypeKey)

	defer func() {
		if nil != _containerType {
			containerType := occp.ToContainerType(_containerType.(string))
			s.OnDisconnectChan <- containerType
		}

		doneChan <- struct{}{}
	}()

	var (
		src            crp.ServerRequestCodec
		messageType    int
		message        []byte
		result         interface{}
		resMessageType int
		resMessage     []byte
		err            error
	)

	if nil != _containerType {
		containerType := occp.ToContainerType(_containerType.(string))
		s.OnConnectChan <- &container.OnConnectInfo{
			ContainerType: containerType,
			WriteChan:     writeChan,
		}
	}

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

			if !s.RPCInvoker.HasMethod(src.Method()) {
				logging.Logger().Error(err)
				s.writeError(src, writeChan, crp.E_NO_METHOD, "", err)
				break
			}

			result, err = s.RPCInvoker.Invoke(src)
			if nil != err {
				logging.Logger().Error(err)
			}

			if !src.HasResponse() {
				break
			}
			resMessageType, resMessage, err = src.NewResponse(result, err)
			if nil != err {
				logging.Logger().Error(err)
				s.writeError(src, writeChan, crp.E_INTERNAL, "", err)
				break
			}

			writeChan <- css.MakeSocketMessage(resMessageType, resMessage)
		case <-stopChan:
			return
		}
	}
}

func (s *ContainerServlets) writeError(src crp.ServerRequestCodec, writeChan chan<- css.SocketMessage, code crp.ErrorCode, message string, data interface{}) {
	if !src.HasResponse() {
		return
	}

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

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