zquit.go

package zquit

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

type ZQuit struct {
	isQuit bool
	wg     *zwg.Zwg
}

func New() *ZQuit {
	return &ZQuit{
		wg: zwg.New(),
	}
}

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

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

func (q *ZQuit) Wait() {
	q.Shutdown()

	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()) {
	if pre != nil {
		pre()
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)

	<-c

	q.Wait()

	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"))
		return
	}

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

	c.Next()
}

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