server-go/server.go

package server

import (
	"context"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"git.loafle.net/commons/logging-go"
)

type Server struct {
	ServerHandler ServerHandler

	ctx      ServerContext
	servlets sync.Map
	stopChan chan struct{}
	stopWg   sync.WaitGroup
}

func (s *Server) ListenAndServe() error {
	if s.stopChan != nil {
		return fmt.Errorf("Server: server is already running. Stop it before starting it again")
	}

	var (
		err      error
		listener net.Listener
	)
	if nil == s.ServerHandler {
		panic("Server: server handler must be specified.")
	}
	s.ServerHandler.Validate()

	s.ctx = s.ServerHandler.ServerContext()
	if nil == s.ctx {
		return fmt.Errorf("Server: ServerContext is nil")
	}

	if err = s.ServerHandler.Init(s.ctx); nil != err {
		return fmt.Errorf("Server: Initialization of server has been failed %v", err)
	}

	if listener, err = s.ServerHandler.Listen(s.ctx); nil != err {
		return err
	}

	s.stopChan = make(chan struct{})
	s.stopWg.Add(1)
	return s.handleLoop(listener)
}

func (s *Server) Shutdown(ctx context.Context) error {
	if s.stopChan == nil {
		return fmt.Errorf("Server: server must be started before stopping it")
	}
	close(s.stopChan)
	s.stopWg.Wait()

	s.ServerHandler.Destroy(s.ctx)

	s.stopChan = nil

	return nil
}

func (s *Server) ConnectionSize() int {
	var sz int
	s.servlets.Range(func(k, v interface{}) bool {
		sz++
		return true
	})
	return sz
}

func (s *Server) handleLoop(listener net.Listener) error {
	var (
		stopping atomic.Value
		conn     net.Conn
		err      error
	)

	defer func() {
		s.ServerHandler.OnStop(s.ctx)

		s.stopWg.Done()
	}()

	if err = s.ServerHandler.OnStart(s.ctx); nil != err {
		return err
	}

	for {
		acceptChan := make(chan struct{})

		go func() {
			if conn, err = listener.Accept(); err != nil {
				if nil == stopping.Load() {
					logging.Logger().Errorf("Server: Cannot accept new connection: [%s]", err)
				}
			}
			close(acceptChan)
		}()

		select {
		case <-s.stopChan:
			stopping.Store(true)
			listener.Close()
			<-acceptChan
			return nil
		case <-acceptChan:
		}

		if nil != err {
			select {
			case <-s.stopChan:
				return nil
			case <-time.After(time.Second):
			}
			continue
		}

		if 0 < s.ServerHandler.GetMaxConnections() {
			sz := s.ConnectionSize()
			if sz >= s.ServerHandler.GetMaxConnections() {
				logging.Logger().Warnf("max connections size %d, refuse\n", sz)
				conn.Close()
				continue
			}
		}

		s.stopWg.Add(1)
		go s.handleConnection(conn)
	}
}

func (s *Server) handleConnection(conn net.Conn) {
	servlet := s.ServerHandler.Servlet()

	defer func() {
		s.servlets.Delete(servlet)
		s.stopWg.Done()
	}()

	if nil == servlet {
		logging.Logger().Errorf("Server: Servlet is nil")
	}
	s.servlets.Store(servlet, true)

	servletStopChan := make(chan struct{})
	doneChan := make(chan struct{})

	go servlet.Handle(s.ctx, conn, doneChan, servletStopChan)

	select {
	case <-doneChan:
		close(servletStopChan)
		conn.Close()
	case <-s.stopChan:
		close(servletStopChan)
		conn.Close()
		<-doneChan
	}
}