code.cestus.io/tools/fabricator@v0.4.3/pkg/helpers/context.go (about)

     1  package helpers
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"os/signal"
     8  	"sync"
     9  
    10  	"code.cestus.io/tools/fabricator/pkg/fabricator"
    11  )
    12  
    13  // WithCancelOnSignal returns a context that will get cancelled whenever one of
    14  // the specified signals is caught.
    15  func WithCancelOnSignal(ctx context.Context, io fabricator.IOStreams, signals ...os.Signal) (context.Context, func()) {
    16  	var once sync.Once
    17  	ctx, cancel := context.WithCancel(ctx)
    18  
    19  	ch := make(chan os.Signal)
    20  
    21  	signal.Notify(ch, signals...)
    22  
    23  	go func() {
    24  		if sig := <-ch; sig != nil {
    25  			fmt.Fprintf(io.Out, "Context received signal: %s\n", sig)
    26  		}
    27  
    28  		cancel()
    29  	}()
    30  
    31  	return ctx, func() {
    32  		once.Do(func() {
    33  			signal.Reset(signals...)
    34  			close(ch)
    35  		})
    36  		<-ctx.Done()
    37  	}
    38  }