server-go/socket/server-readwriter.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()
		}
		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")

		servlet.OnDisconnect(servletCtx)
		olog.Logger().Info("Client has been disconnected", zap.String("Address", addr.String()))

		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
	}
}