zquit.go

package zquit

import (
	"fmt"
	"git.barsukov.pro/barsukov/zgo/zwg"
	"github.com/gin-gonic/gin"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

type ZQuit struct {
	Name        string
	PostWaitDur time.Duration

	isQuit bool
	wg     *zwg.Zwg
}

func New(name string, postWaitDur time.Duration) *ZQuit {
	return &ZQuit{
		Name:        name,
		PostWaitDur: postWaitDur,
		wg:          zwg.New(),
	}
}

func Default(name ...string) *ZQuit {
	if len(name) == 0 {
		name = []string{"MAIN"}
	}

	return New(name[0], time.Second)
}

func (q *ZQuit) Inc() {
	q.wg.Inc()
}

func (q *ZQuit) Done() {
	q.wg.Done()
}

func (q *ZQuit) Wait() {
	q.wg.Wait()
}

func (q *ZQuit) IncDefer() func() {
	return q.wg.IncDefer()
}

func (q *ZQuit) IsQuit() bool {
	return q.isQuit
}

func (q *ZQuit) Val() int {
	return q.wg.Val()
}

func (q *ZQuit) WaitInterruptPrePost(pre func(), post func()) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	<-c

	if pre != nil {
		pre()
	}

	q.Shutdown()
	q.Wait()
	time.Sleep(q.PostWaitDur)

	if post != nil {
		post()
	}
}

func (q *ZQuit) WaitInterrupt() {
	q.WaitInterruptPrePost(nil, nil)
}

func (q *ZQuit) Middleware(c *gin.Context) {
	if q.isQuit {
		c.Data(http.StatusServiceUnavailable, "text/plain", []byte("server shutdown"))
		c.Abort()

		return
	}

	cancel := q.IncDefer()
	defer cancel()

	c.Next()
}

func (q *ZQuit) GetStat() string {
	return fmt.Sprintf("%s: %v", q.Name, q.Val())
}
func (q *ZQuit) PrintStat(sec int) {
	go func() {
		for {
			time.Sleep(time.Second * time.Duration(sec))
			log.Printf("ZQUIT: %s", q.GetStat())
		}
	}()
}

func (q *ZQuit) Shutdown() {
	q.isQuit = true
}