server-go/server-readwriter.go

package server

import (
	"sync"

	logging "git.loafle.net/commons/logging-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 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 []byte)
	writeChan := make(chan []byte)

	readerDoneChan := make(chan error)
	writerDoneChan := make(chan error)

	go servlet.Handle(servletCtx, stopChan, servletDoneChan, readChan, writeChan)
	go connReadHandler(srw.ReadwriteHandler, conn, stopChan, readerDoneChan, readChan)
	go connWriteHandler(srw.ReadwriteHandler, conn, stopChan, writerDoneChan, 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
	}
}