github.com/icyphox/x@v0.0.355-0.20220311094250-029bd783e8b8/sqlcon/dockertest/onexit.go (about)

     1  package dockertest
     2  
     3  import (
     4  	"os"
     5  	"os/signal"
     6  	"sync"
     7  	"syscall"
     8  )
     9  
    10  const interruptedExitCode = 130
    11  
    12  // OnExit helps with cleaning up docker test.
    13  type OnExit struct {
    14  	sync.Mutex
    15  	once     sync.Once
    16  	handlers []func()
    17  }
    18  
    19  // NewOnExit create a new OnExit instance.
    20  func NewOnExit() *OnExit {
    21  	return &OnExit{
    22  		handlers: make([]func(), 0),
    23  	}
    24  }
    25  
    26  // Add adds a task that is executed on SIGINT, SIGKILL, SIGTERM.
    27  func (at *OnExit) Add(f func()) {
    28  	at.Lock()
    29  	defer at.Unlock()
    30  	at.handlers = append(at.handlers, f)
    31  	at.once.Do(func() {
    32  		go func() {
    33  			c := make(chan os.Signal, 1)
    34  			signal.Notify(c, syscall.SIGINT, syscall.SIGKILL, syscall.SIGTERM)
    35  			<-c
    36  			at.Exit(interruptedExitCode)
    37  		}()
    38  	})
    39  }
    40  
    41  // Exit wraps os.Exit
    42  func (at *OnExit) Exit(status int) {
    43  	at.execute()
    44  	os.Exit(status)
    45  }
    46  
    47  func (at *OnExit) execute() {
    48  	at.Lock()
    49  	defer at.Unlock()
    50  	for _, f := range at.handlers {
    51  		f()
    52  	}
    53  	at.handlers = make([]func(), 0)
    54  }