ing

server.go

@@ -151,6 +151,8 @@ func handleConnection(s *server, soc *Socket, socketHandler SocketHandler) {
 	defer s.stopWg.Done()
 
 	logging.Logger().Debug(fmt.Sprintf("Server: Client[%s] is connected.", soc.RemoteAddr()))
+	socketHandler.OnConnect(soc)
+
 	clientStopChan := make(chan struct{})
 	handleDoneChan := make(chan struct{})
 
@@ -159,11 +161,11 @@ func handleConnection(s *server, soc *Socket, socketHandler SocketHandler) {
 	select {
 	case <-s.stopChan:
 		close(clientStopChan)
-		soc.Close()
 		<-handleDoneChan
 	case <-handleDoneChan:
 		close(clientStopChan)
 		logging.Logger().Debug(fmt.Sprintf("Server: Client[%s] is disconnected.", soc.RemoteAddr()))
+		socketHandler.OnDisconnect(soc)
 		soc.Close()
 	}
 }

socket_handler.go

@@ -10,7 +10,9 @@ type SocketHandler interface {
 	// Handshake do handshake client and server
 	// id is identity of client socket. if id is nil, disallow connection
 	Handshake(ctx *fasthttp.RequestCtx) (id interface{}, extensionsHeader *fasthttp.ResponseHeader)
+	OnConnect(soc *Socket)
 	Handle(soc *Socket, stopChan <-chan struct{}, doneChan chan<- struct{})
+	OnDisconnect(soc *Socket)
 
 	GetMaxMessageSize() int64
 	GetWriteTimeout() time.Duration

socket_handlers.go

@@ -30,11 +30,18 @@ type SocketHandlers struct {
 func (sh *SocketHandlers) Handshake(ctx *fasthttp.RequestCtx) (id interface{}, extensionsHeader *fasthttp.ResponseHeader) {
 	return nil, nil
 }
+func (sh *SocketHandlers) OnConnect(soc *Socket) {
+	// no op
+}
 
 func (sh *SocketHandlers) Handle(soc *Socket, stopChan <-chan struct{}, doneChan chan<- struct{}) {
 	// no op
 }
 
+func (sh *SocketHandlers) OnDisconnect(soc *Socket) {
+	// no op
+}
+
 func (sh *SocketHandlers) GetMaxMessageSize() int64 {
 	return sh.MaxMessageSize
 }