package socket import ( "sync" "git.loafle.net/commons/logging-go" "git.loafle.net/commons/server-go" ) type ServerReadWriter struct { connections sync.Map ReadwriteHandler ReadWriteHandler ServerStopChan <-chan struct{} ServerStopWg *sync.WaitGroup } func (srw *ServerReadWriter) ConnectionSize() int { var sz int srw.connections.Range(func(k, v interface{}) bool { sz++ return true }) return sz } func (srw *ServerReadWriter) HandleConnection(servlet Servlet, servletCtx server.ServletCtx, conn Conn) { addr := conn.RemoteAddr() defer func() { if nil != conn { conn.Close() } servlet.OnDisconnect(servletCtx) logging.Logger().Infof("Client[%s] has been disconnected", addr) srw.ServerStopWg.Done() }() logging.Logger().Infof("Client[%s] has been connected", addr) srw.connections.Store(conn, true) defer srw.connections.Delete(conn) servlet.OnConnect(servletCtx, conn) conn.SetCloseHandler(func(code int, text string) error { logging.Logger().Debugf("close") return nil }) stopChan := make(chan struct{}) servletDoneChan := make(chan struct{}) readChan := make(chan SocketMessage) writeChan := make(chan SocketMessage) readerDoneChan := make(chan error) writerDoneChan := make(chan error) go connReadHandler(srw.ReadwriteHandler, conn, stopChan, readerDoneChan, readChan) go connWriteHandler(srw.ReadwriteHandler, conn, stopChan, writerDoneChan, writeChan) go servlet.Handle(servletCtx, stopChan, servletDoneChan, readChan, writeChan) select { case <-readerDoneChan: close(stopChan) <-writerDoneChan <-servletDoneChan case <-writerDoneChan: close(stopChan) <-readerDoneChan <-servletDoneChan case <-servletDoneChan: close(stopChan) <-readerDoneChan <-writerDoneChan case <-srw.ServerStopChan: close(stopChan) <-readerDoneChan <-writerDoneChan <-servletDoneChan } }