GO Shutdown

shutdown.go

package shutdown

import (
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/go-playground/validator/v10"
	"github.com/pkg/errors"
)

// Shutdown Shutdown handles explicit shutdown or shutdown from signals
type Shutdown struct {
	done      bool
	mu        *sync.RWMutex
	shutdownC chan interface{}
	signalsC  chan os.Signal
	timeout   time.Duration
	timeoutC  chan interface{}
	wg        *sync.WaitGroup
}

// Options to New
type Options struct {
	Timeout time.Duration
}

// New create Shutdown
func New(ops Options) *Shutdown {
	v := validator.New()
	err := v.Struct(ops)
	if err != nil {
		panic(errors.Wrap(err, "can not create shutdown"))
	}
	return &Shutdown{
		done:      false,
		mu:        &sync.RWMutex{},
		shutdownC: make(chan interface{}, 1),
		signalsC:  make(chan os.Signal, 1),
		timeout:   ops.Timeout,
		timeoutC:  make(chan interface{}, 1),
		wg:        &sync.WaitGroup{},
	}
}

// Setup start trapping signals
func (s *Shutdown) Setup(signals ...os.Signal) {
	s.wg.Add(1)
	if len(signals) < 1 {
		signals = []os.Signal{os.Interrupt, syscall.SIGTERM}
	}
	signal.Notify(s.signalsC, signals...)
	if s.timeout > 0 {
		go func() {
			s.timeoutC <- time.After(s.timeout)
		}()
	}
	go func() {
		select {
		case <-s.signalsC:
			s.finish()
		case <-s.shutdownC:
			s.finish()
		case <-s.timeoutC:
			s.finish()
		}
	}()
}

func (s *Shutdown) finish() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.done = true
	s.wg.Done()
}

// Wait blocks until wait group is completely released
func (s *Shutdown) Wait() {
	s.wg.Wait()
}

// IsDone checks that the program is done
func (s *Shutdown) IsDone() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.done
}

// Shutdown shutdown immediately
func (s *Shutdown) Shutdown() {
	s.shutdownC <- true
}

// WGAdd add to wait group
func (s *Shutdown) WGAdd() {
	s.wg.Add(1)
}

// WGRelease release wait group
func (s *Shutdown) WGRelease() {
	s.wg.Done()
}