github.com/qioalice/ekago/v3@v3.3.2-0.20221202205325-5c262d586ee4/ekaserv/observe.go (about)

     1  // Copyright © 2022. All rights reserved.
     2  // Author: Ilya Stroy.
     3  // Contacts: iyuryevich@pm.me, https://github.com/qioalice
     4  // License: https://opensource.org/licenses/MIT
     5  
     6  package ekaserv
     7  
     8  import (
     9  	"context"
    10  	"sync"
    11  	"sync/atomic"
    12  	"unsafe"
    13  
    14  	"github.com/qioalice/ekago/v3/ekadeath"
    15  )
    16  
    17  // TODO: Comment
    18  
    19  type (
    20  	ObserveController struct {
    21  		ctx  context.Context
    22  		wg   sync.WaitGroup
    23  		sema sync.Mutex
    24  	}
    25  
    26  	// ObserveCallback is just a function alias for callback
    27  	// that you will use in Observe() function.
    28  	ObserveCallback func(oc *ObserveController)
    29  )
    30  
    31  const (
    32  	mutexLocked = 1 << iota
    33  )
    34  
    35  var (
    36  	// observeCalled is a special package-level variable,
    37  	// that checks whether Observe() is called only once using atomic CAS operation.
    38  	observeCalled uint32
    39  )
    40  
    41  func (oc *ObserveController) Join() bool {
    42  	ptr := unsafe.Pointer(&oc.sema)
    43  	if !atomic.CompareAndSwapInt32((*int32)(ptr), 0, mutexLocked) {
    44  		return false
    45  	}
    46  
    47  	oc.wg.Add(1)
    48  	oc.sema.Unlock()
    49  	return true
    50  }
    51  
    52  func (oc *ObserveController) Leave() {
    53  	oc.wg.Done()
    54  }
    55  
    56  func (oc *ObserveController) Context() context.Context {
    57  	return oc.ctx
    58  }
    59  
    60  func (oc *ObserveController) wait() {
    61  	oc.sema.Lock()
    62  	oc.wg.Wait()
    63  }
    64  
    65  func newObserveController(ctx context.Context) *ObserveController {
    66  	return &ObserveController{ctx: ctx}
    67  }
    68  
    69  // Observe is your service initializer.
    70  // It provides a mechanism of graceful shutdown, supporting OS signals (like SIGINT)
    71  // and requested shutdowns from ekadeath or ekalog.
    72  //
    73  // First of all. You need a special "main" context.
    74  // If you pass nil as context.Context, the context.Background() will be used.
    75  // The difference is:
    76  //  - Overwritten background context can be canceled ONLY by OS signals;
    77  //  - Your context can be cancelled by your hands manually AND by OS signals.
    78  // This context will be used to be able to say your jobs that it's time to stop
    79  // their work and they need to be completed completely.
    80  //
    81  // WARNING!
    82  // You can call Observe() ONLY ONE per whole lifetime of your application.
    83  // The next call will lead to panic.
    84  func Observe(ctx context.Context, cb ObserveCallback) {
    85  
    86  	if !atomic.CompareAndSwapUint32(&observeCalled, 0, 1) {
    87  		panic("Observe must be called only once at all!")
    88  	}
    89  
    90  	if cb == nil {
    91  		panic("Observe initialization callback must not be nil")
    92  	}
    93  
    94  	if ctx == nil {
    95  		ctx = context.Background()
    96  	}
    97  
    98  	ctx, cancelFunc := context.WithCancel(ctx)
    99  	oc := newObserveController(ctx)
   100  
   101  	ekadeath.Reg(func() {
   102  		cancelFunc()
   103  		oc.wait()
   104  	})
   105  
   106  	cb(oc)
   107  }