server-go/web/fasthttp/server.go

package fasthttp

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

	"git.loafle.net/commons/logging-go"
	"git.loafle.net/commons/server-go"
	"git.loafle.net/commons/server-go/web"
	"github.com/valyala/fasthttp"
)

type Server struct {
	ServerHandler ServerHandler

	ctx      server.ServerCtx
	stopChan chan struct{}
	stopWg   sync.WaitGroup

	hs *fasthttp.Server
}

func (s *Server) ListenAndServe() error {
	var (
		err      error
		listener net.Listener
	)
	if nil == s.ServerHandler {
		return fmt.Errorf("%s server handler must be specified", s.logHeader())
	}
	s.ServerHandler.Validate()

	if s.stopChan != nil {
		return fmt.Errorf("%s already running. Stop it before starting it again", s.logHeader())
	}

	s.ctx = s.ServerHandler.ServerCtx()
	if nil == s.ctx {
		return fmt.Errorf("%s ServerCtx is nil", s.logHeader())
	}

	s.hs = &fasthttp.Server{
		Handler:         s.httpHandler,
		Name:            s.ServerHandler.GetName(),
		Concurrency:     s.ServerHandler.GetConcurrency(),
		ReadBufferSize:  s.ServerHandler.GetReadBufferSize(),
		WriteBufferSize: s.ServerHandler.GetWriteBufferSize(),
		ReadTimeout:     s.ServerHandler.GetReadTimeout(),
		WriteTimeout:    s.ServerHandler.GetWriteTimeout(),
	}

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

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

	s.stopChan = make(chan struct{})

	s.stopWg.Add(1)
	return s.handleServer(listener)
}

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

	s.ServerHandler.Destroy(s.ctx)

	s.stopChan = nil

	return nil
}

func (s *Server) logHeader() string {
	return fmt.Sprintf("Server[%s]:", s.ServerHandler.GetName())
}

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

	defer func() {
		if nil != listener {
			listener.Close()
		}
		s.ServerHandler.OnStop(s.ctx)

		logging.Logger().Infof("%s Stopped", s.logHeader())

		s.stopWg.Done()
	}()

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

	hsCloseChan := make(chan error)
	go func() {
		if err := s.hs.Serve(listener); nil != err {
			if nil == stopping.Load() {
				hsCloseChan <- err
				return
			}
		}
		hsCloseChan <- nil
	}()

	logging.Logger().Infof("%s Started", s.logHeader())

	select {
	case err, _ := <-hsCloseChan:
		if nil != err {
			return err
		}
	case <-s.stopChan:
		stopping.Store(true)
		listener.Close()
		<-hsCloseChan
		listener = nil
	}

	return nil
}

func (s *Server) httpHandler(ctx *fasthttp.RequestCtx) {
	var (
		servlet Servlet
	)

	if s.ServerHandler.CheckOrigin(ctx) {
		return
	}

	if servlet = s.ServerHandler.Servlet(s.ctx, ctx); nil == servlet {
		s.onError(ctx, web.NewError(fasthttp.StatusNotFound, fmt.Errorf("Not Found")))
		return
	}

	servletCtx := servlet.ServletCtx(s.ctx)

	if err := servlet.Handle(servletCtx, ctx); nil != err {
		s.onError(ctx, err)
	}
}

func (s *Server) onError(ctx *fasthttp.RequestCtx, err *web.Error) {
	s.ServerHandler.OnError(s.ctx, ctx, err)
}