85 lines
1.9 KiB
Go
85 lines
1.9 KiB
Go
|
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
|
||
|
|
}
|
||
|
|
}
|