protoc-gen-kit/pkg/xserver/xserver.go

package xserver

import (
	"context"
	"os"
	"os/signal"
	"sync"
	"time"
)

type Server interface {
	Start() error
	Shutdown(ctx context.Context) error
	Name() string
}

type baseServer struct {
	services  []Server
	wg        sync.WaitGroup
	errChan   chan error
	closeChan chan struct{}
}

func NewServer() *baseServer {
	return &baseServer{
		services:  make([]Server, 0),
		wg:        sync.WaitGroup{},
		errChan:   make(chan error),
		closeChan: make(chan struct{}),
	}
}

func (s *baseServer) Register(srv ...Server) {
	s.services = append(s.services, srv...)
}
func (s *baseServer) Start() error {
	var hasErr error
	for _, v := range s.services {
		if hasErr != nil {
			break
		}
		go func(v Server) {
			if err := v.Start(); err != nil {
				s.errChan <- err
				hasErr = err
			}
		}(v)
	}
	return hasErr
}

func (s *baseServer) Shutdown(ctx context.Context) error {
	for _, v := range s.services {
		go func(v Server) {
			<-s.closeChan
			if err := v.Shutdown(context.Background()); err != nil {
				s.errChan <- err
			}
		}(v)
	}
	return <-s.errChan
}
func (s *baseServer) NotifyClose() {
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	<-signalChan
	close(s.closeChan)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
	defer cancel()
	s.Shutdown(ctx)
	time.Sleep(time.Second)
}