package socket import ( "sync" "go.uber.org/zap" olog "git.loafle.net/overflow/log-go" "git.loafle.net/overflow/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) olog.Logger().Info("Client has been disconnected", zap.String("Address", addr.String())) srw.ServerStopWg.Done() }() olog.Logger().Info("Client has been connected", zap.String("Address", addr.String())) srw.connections.Store(conn, true) defer srw.connections.Delete(conn) servlet.OnConnect(servletCtx, conn) conn.SetCloseHandler(func(code int, text string) error { olog.Logger().Debug("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 } }