diff --git a/constants.go b/constants.go new file mode 100644 index 0000000..2735701 --- /dev/null +++ b/constants.go @@ -0,0 +1,7 @@ +package overflow_probe_container + +import cuc "git.loafle.net/commons_go/util/context" + +var ( + RPCServletKey = cuc.ContextKey("RPCServlet") +) diff --git a/server/socket_handlers.go b/server/socket_handlers.go index 0b134f9..23be9fd 100644 --- a/server/socket_handlers.go +++ b/server/socket_handlers.go @@ -5,10 +5,11 @@ import ( "sync" cRPC "git.loafle.net/commons_go/rpc" - crsrs "git.loafle.net/commons_go/rpc/server/rwc/socket" - "git.loafle.net/commons_go/rpc/protocol/json" + crsrs "git.loafle.net/commons_go/rpc/server/rwc/socket" "git.loafle.net/commons_go/server" + + oopc "git.loafle.net/overflow/overflow_probe_container" ) func newSocketHandler(rpcSH RPCServletHandler) SocketHandler { @@ -45,12 +46,12 @@ func (sh *SocketHandlers) OnConnect(soc server.Socket) { sh.SocketHandlers.OnConnect(soc) soc.Context().SetAttribute(cRPC.ContentTypeKey, json.Name) - + soc.Context().SetAttribute(oopc.RPCServletKey, retainRPCServlet(sh.rpcSH, sh.rpcRWCSH)) } func (sh *SocketHandlers) Handle(soc server.Socket, stopChan <-chan struct{}, doneChan chan<- error) { var err error - rpcServlet := retainRPCServlet(sh.rpcSH, sh.rpcRWCSH) + rpcServlet := soc.Context().GetAttribute(oopc.RPCServletKey).(cRPC.Servlet) defer func() { releaseRPCServlet(rpcServlet) @@ -72,6 +73,7 @@ func (sh *SocketHandlers) Handle(soc server.Socket, stopChan <-chan struct{}, do } func (sh *SocketHandlers) OnDisconnect(soc server.Socket) { + soc.Context().RemoveAttribute(oopc.RPCServletKey) sh.SocketHandlers.OnDisconnect(soc) }