github.com/angenalZZZ/gofunc@v0.0.0-20210507121333-48ff1be3917b/f/context.go (about)

     1  package f
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"os/signal"
     7  	"sync"
     8  	"sync/atomic"
     9  	"syscall"
    10  	"time"
    11  )
    12  
    13  // ContextOnInterrupt creates a new context that cancels on SIGINT or SIGTERM.
    14  func ContextOnInterrupt() (context.Context, func()) {
    15  	return ContextWrapSignal(context.Background(), syscall.SIGINT, syscall.SIGTERM)
    16  }
    17  
    18  // ContextOnSignal creates a new context that cancels on the given signals.
    19  func ContextOnSignal(signals ...os.Signal) (context.Context, func()) {
    20  	return ContextWrapSignal(context.Background(), signals...)
    21  }
    22  
    23  // ContextWrapSignal creates a new context that cancels on the given signals. It wraps the provided context.
    24  func ContextWrapSignal(parent context.Context, signals ...os.Signal) (ctx context.Context, closer func()) {
    25  	ctx, closer = context.WithCancel(parent)
    26  
    27  	c := make(chan os.Signal, 1)
    28  	signal.Notify(c, signals...)
    29  
    30  	go func() {
    31  		select {
    32  		case <-c:
    33  			closer()
    34  		case <-ctx.Done():
    35  		}
    36  	}()
    37  
    38  	return ctx, closer
    39  }
    40  
    41  // contextWaitChan is an unexported type for channel used as the blocking statement in
    42  // the waitFunc returned by wait.WithWait.
    43  type contextWaitChan chan struct{}
    44  
    45  // contextWaitKey is the key for wait.contextWaitChan values in Contexts. It is unexported;
    46  // clients use wait.WithWait and wait.Done instead of using this key directly.
    47  const contextWaitKey = 32 << (uint64(^uint(0)) >> 63)
    48  
    49  // contextWaitMutex is the mutual exclusion lock used to prevent a race condition in wait.Done.
    50  var contextWaitMutex sync.Mutex
    51  
    52  // ContextWithWait creates a new context that wait on the given ContextDone().
    53  func ContextWithWait(parent context.Context) (ctx context.Context, waitFunc func()) {
    54  	ctx = parent
    55  
    56  	var ok bool
    57  	var wait contextWaitChan
    58  	if wait, ok = ctx.Value(contextWaitKey).(contextWaitChan); !ok {
    59  		wait = make(contextWaitChan)
    60  		ctx = context.WithValue(parent, contextWaitKey, wait)
    61  	}
    62  
    63  	waitFunc = func() {
    64  		select {
    65  		case <-wait:
    66  		}
    67  	}
    68  	return
    69  }
    70  
    71  // DoneContext unblocks the waitFunc returned by wait.WithWait for the provided ctx, if available.
    72  // If waitFunc has already been unblocked, then nothing happens.
    73  func DoneContext(ctx context.Context) {
    74  	if wait, ok := ctx.Value(contextWaitKey).(contextWaitChan); ok {
    75  		contextWaitMutex.Lock()
    76  		defer contextWaitMutex.Unlock()
    77  
    78  		select {
    79  		case <-wait:
    80  		default:
    81  			close(wait)
    82  		}
    83  	}
    84  }
    85  
    86  // NumIncrWait to wait succeeded-number equals published-number.
    87  func NumIncrWait(publishedNumber, succeededNumber *int64) {
    88  	for atomic.LoadInt64(succeededNumber) < atomic.LoadInt64(publishedNumber) {
    89  		time.Sleep(time.Nanosecond)
    90  	}
    91  }