github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/synckit/chain_cancel.go (about)

     1  // Copyright 2020 Insolar Network Ltd.
     2  // All rights reserved.
     3  // This material is licensed under the Insolar License version 1.0,
     4  // available at https://github.com/insolar/assured-ledger/blob/master/LICENSE.md.
     5  
     6  package synckit
     7  
     8  import (
     9  	"context"
    10  	"sync/atomic"
    11  )
    12  
    13  func NewChainedCancel() *ChainedCancel {
    14  	return &ChainedCancel{}
    15  }
    16  
    17  type ChainedCancel struct {
    18  	state uint32 // atomic
    19  	chain atomic.Value
    20  }
    21  
    22  const (
    23  	stateCancelled = 1 << iota
    24  	stateChainHandlerBeingSet
    25  	stateChainHandlerSet
    26  )
    27  
    28  func (p *ChainedCancel) Cancel() {
    29  	if p == nil {
    30  		return
    31  	}
    32  	for {
    33  		lastState := atomic.LoadUint32(&p.state)
    34  		switch {
    35  		case lastState&stateCancelled != 0:
    36  			return
    37  		case !atomic.CompareAndSwapUint32(&p.state, lastState, lastState|stateCancelled):
    38  			continue
    39  		case lastState == stateChainHandlerSet:
    40  			p.runChain()
    41  		}
    42  		return
    43  	}
    44  }
    45  
    46  func (p *ChainedCancel) runChain() {
    47  	// here is a potential problem, because Go spec doesn't provide ANY ordering on atomic operations
    48  	// but Go compiler does provide some guarantees, so lets hope for the best
    49  
    50  	fn := (p.chain.Load()).(context.CancelFunc)
    51  	if fn == nil {
    52  		// this can only happen when atomic ordering is broken
    53  		panic("unexpected atomic ordering")
    54  	}
    55  
    56  	// prevent repeated calls as well as retention of references & possible memory leaks
    57  	// type cast is mandatory due to specifics of atomic.Value
    58  	p.chain.Store(context.CancelFunc(func() {}))
    59  	fn()
    60  }
    61  
    62  func (p *ChainedCancel) IsCancelled() bool {
    63  	return p != nil && atomic.LoadUint32(&p.state)&stateCancelled != 0
    64  }
    65  
    66  /*
    67  	SetChain sets a chained function once.
    68  	The chained function can only be set once to a non-null value, further calls will panic.
    69      But if the chained function was not set, the SetChain(nil) can be called multiple times.
    70  
    71  	The chained function is guaranteed to be called only once, it will also be called on set when IsCancelled is already true.
    72  */
    73  func (p *ChainedCancel) SetChain(chain context.CancelFunc) {
    74  	if chain == nil {
    75  		if p.chain.Load() == nil {
    76  			return
    77  		}
    78  		panic("illegal state")
    79  	}
    80  	for {
    81  		lastState := atomic.LoadUint32(&p.state)
    82  		switch {
    83  		case lastState&^stateCancelled != 0: // chain is set or being set
    84  			panic("illegal state")
    85  		case !atomic.CompareAndSwapUint32(&p.state, lastState, lastState|stateChainHandlerBeingSet): //
    86  			continue
    87  		}
    88  		break
    89  	}
    90  
    91  	p.chain.Store(chain)
    92  
    93  	for {
    94  		lastState := atomic.LoadUint32(&p.state)
    95  		switch {
    96  		case lastState&^stateCancelled != stateChainHandlerBeingSet:
    97  			// this can only happen when atomic ordering is broken
    98  			panic("unexpected atomic ordering")
    99  		case !atomic.CompareAndSwapUint32(&p.state, lastState, (lastState&stateCancelled)|stateChainHandlerSet):
   100  			continue
   101  		case lastState&stateCancelled != 0:
   102  			// if cancel was set then call the chained cancel fn here
   103  			// otherwise, the cancelling process will be responsible to call the chained cancel
   104  			p.runChain()
   105  		}
   106  		return
   107  	}
   108  }