zquit.go

package zquit

import (
	"github.com/gin-gonic/gin"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
)

type ZQuit struct {
	isQuit bool
	delta  atomic.Int32
	wg     sync.WaitGroup
}

func Default() *ZQuit {
	return &ZQuit{}
}

func (q *ZQuit) Add(delta int) {
	q.delta.Add(int32(delta))
	q.wg.Add(delta)
}

func (q *ZQuit) Inc() {
	q.Add(1)
}

func (q *ZQuit) Done() {
	q.delta.Add(-1)
	q.wg.Done()
}

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

func (q *ZQuit) IncDefer() func() {
	q.Inc()

	return func() {
		q.Done()
	}
}

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

func (q *ZQuit) Delta() int32 {
	return q.delta.Load()
}

func (q *ZQuit) WaitInterruptPrePost(pre func(), post func()) {
	pre()

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

	<-c

	q.Wait()

	post()
}

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

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

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

	c.Next()
}