gohttp/server/manager.go

package server

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"strings"

	"bufio"
	"io"

	"git.pyer.club/kingecg/gologger"
	logger "git.pyer.club/kingecg/gologger"
	"github.com/soheilhy/cmux"
)

var ServerManager map[string]*ServerListener = make(map[string]*ServerListener)

func makeMatcher(name string, s *ServerListener) cmux.Matcher {
	return func(r io.Reader) bool {
		if s.ServerCount() == 1 {
			return true
		}
		req, err := http.ReadRequest(bufio.NewReader(r))
		if err != nil {
			return false
		}
		return strings.HasPrefix(req.Host, name)
	}
}

type ServerListener struct {
	port     int
	listener cmux.CMux
	servers  map[string]*http.Server
}

func (s *ServerListener) AddServer(name string, server *http.Server) {
	s.servers[name] = server
}

func (s *ServerListener) ServerCount() int {
	return len(s.servers)
}

func (s *ServerListener) StartServer(name string) error {
	server, ok := s.servers[name]
	if !ok {
		panic("No named server")
	}

	l := s.listener.Match(makeMatcher(name, s))
	return server.Serve(l)
}

func (s *ServerListener) Start() {
	for name, _ := range s.servers {
		go func(name string) {
			err := s.StartServer(name)
			if err != nil {
				l := logger.GetLogger("ServerListener")
				l.Error("Start server error:", name, err)
			}
		}(name)
	}
}

func (s *ServerListener) ShutDown() {
	// may need add wait group
	for _, server := range s.servers {
		server.Shutdown(context.Background())
	}
	s.listener.Close()
}

func (s *ServerListener) Serve() {
	l := gologger.GetLogger("Listener")
	l.Debug("listen on :", s.port)
	go s.listener.Serve()
}

func NewServerListener(port int) *ServerListener {
	listenStr := fmt.Sprintf(":%d", port)
	l, err := net.Listen("tcp", listenStr)

	if err != nil {
		l := logger.GetLogger("ServerListener")
		l.Error("Listen error:", err)
	}
	muxer := cmux.New(l)
	s := &ServerListener{port: port, listener: muxer, servers: make(map[string]*http.Server)}
	return s
}

func AddServer(name string, server *http.Server, port int) {
	listenStr := fmt.Sprintf(":%d", port)
	listener, ok := ServerManager[listenStr]
	if !ok {
		listener = NewServerListener(port)
		ServerManager[listenStr] = listener
	}
	listener.AddServer(name, server)
}