overflow_probe_container/server/socket_handlers.go

package server

import (
	"net"
	"sync"

	cr "git.loafle.net/commons_go/rpc"
	"git.loafle.net/commons_go/rpc/protocol/json"
	crsrs "git.loafle.net/commons_go/rpc/server/rwc/socket"
	"git.loafle.net/commons_go/server"

	oopcs "git.loafle.net/overflow/overflow_probe_container/service"
)

func NewSocketHandler(rpcSH RPCServletHandler, probeService *oopcs.ProbeService) SocketHandler {
	rpcRWCSH := crsrs.New()

	sh := &SocketHandlers{
		rpcSH:        rpcSH,
		rpcRWCSH:     rpcRWCSH,
		probeService: probeService,
	}

	return sh
}

type SocketHandlers struct {
	server.SocketHandlers

	rpcRWCSH     cr.ServletReadWriteCloseHandler
	rpcSH        RPCServletHandler
	probeService *oopcs.ProbeService
}

func (sh *SocketHandlers) Init(serverCTX server.ServerContext) error {
	if err := sh.SocketHandlers.Init(serverCTX); nil != err {
		return err
	}

	return nil
}

func (sh *SocketHandlers) Handshake(socketCTX server.SocketContext, conn net.Conn) (id string) {
	return "server"
}

func (sh *SocketHandlers) OnConnect(soc server.Socket) {
	sh.SocketHandlers.OnConnect(soc)

	soc.Context().SetAttribute(cr.ContentTypeKey, json.Name)
}

func (sh *SocketHandlers) Handle(soc server.Socket, stopChan <-chan struct{}, doneChan chan<- error) {
	var err error
	rpcServlet := retainRPCServlet(sh.rpcSH, sh.rpcRWCSH)
	sh.probeService.RPCServlet = rpcServlet

	defer func() {
		rpcServlet.Stop()
		sh.probeService.RPCServlet = nil
		releaseRPCServlet(rpcServlet)
		doneChan <- err
	}()

	rpcDoneChan := make(chan error, 1)

	if err = rpcServlet.Start(soc.Context(), soc, rpcDoneChan); nil != err {
		return
	}

	select {
	case err = <-rpcDoneChan:
	case <-stopChan:
	}
}

func (sh *SocketHandlers) OnDisconnect(soc server.Socket) {

	sh.SocketHandlers.OnDisconnect(soc)
}

func (sh *SocketHandlers) Destroy() {

	sh.SocketHandlers.Destroy()
}

func (sh *SocketHandlers) Validate() {

}

var rpcServletPool sync.Pool

func retainRPCServlet(sh RPCServletHandler, rpcRWCSH cr.ServletReadWriteCloseHandler) cr.Servlet {
	v := rpcServletPool.Get()
	if v == nil {
		return cr.NewServlet(sh, rpcRWCSH)
	}
	return v.(cr.Servlet)
}

func releaseRPCServlet(s cr.Servlet) {

	rpcServletPool.Put(s)
}