shutdown.go

package shutdown

import (
	"context"
	"sync"
)

type Shutdown struct {
	cancel context.CancelFunc
	wg     sync.WaitGroup
	once   sync.Once
	err    error
}

func WithContext(ctx context.Context) (*Shutdown, context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	return &Shutdown{cancel: cancel}, ctx
}

func (s *Shutdown) Wait() {
	s.wg.Wait()
}

func (s *Shutdown) Go(f func()) {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		f()
		s.once.Do(func() {
			s.cancel()
		})
	}()
}