github.com/keybase/client/go@v0.0.0-20241007131713-f10651d043c8/libkb/rpccanceler.go (about)

     1  package libkb
     2  
     3  import (
     4  	"sync"
     5  
     6  	"golang.org/x/net/context"
     7  )
     8  
     9  type RPCCancelerReason uint
    10  
    11  const (
    12  	RPCCancelerReasonLogout RPCCancelerReason = 1 << iota
    13  	RPCCancelerReasonBackground
    14  	RPCCancelerReasonMax
    15  )
    16  
    17  const RPCCancelerReasonAll = RPCCancelerReasonMax - 1
    18  
    19  type RPCCancelerKey string
    20  
    21  func NewRPCCancelerKey() RPCCancelerKey {
    22  	return RPCCancelerKey(RandStringB64(3))
    23  }
    24  
    25  type liveContext struct {
    26  	ctx      context.Context
    27  	cancelFn context.CancelFunc
    28  	reason   RPCCancelerReason
    29  }
    30  
    31  type RPCCanceler struct {
    32  	sync.Mutex
    33  	liveCtxs map[RPCCancelerKey]liveContext
    34  }
    35  
    36  func NewRPCCanceler() *RPCCanceler {
    37  	return &RPCCanceler{
    38  		liveCtxs: make(map[RPCCancelerKey]liveContext),
    39  	}
    40  }
    41  
    42  func (r *RPCCanceler) RegisterContext(ctx context.Context, reason RPCCancelerReason) (context.Context, RPCCancelerKey) {
    43  	r.Lock()
    44  	defer r.Unlock()
    45  	var lc liveContext
    46  	lc.ctx, lc.cancelFn = context.WithCancel(ctx)
    47  	id := NewRPCCancelerKey()
    48  	lc.reason = reason
    49  	r.liveCtxs[id] = lc
    50  	return lc.ctx, id
    51  }
    52  
    53  func (r *RPCCanceler) UnregisterContext(id RPCCancelerKey) {
    54  	r.Lock()
    55  	defer r.Unlock()
    56  	if lc, ok := r.liveCtxs[id]; ok {
    57  		lc.cancelFn()
    58  		delete(r.liveCtxs, id)
    59  	}
    60  }
    61  
    62  func (r *RPCCanceler) CancelLiveContexts(reason RPCCancelerReason) {
    63  	r.Lock()
    64  	defer r.Unlock()
    65  	for id, liveCtx := range r.liveCtxs {
    66  		if liveCtx.reason&reason != 0 {
    67  			liveCtx.cancelFn()
    68  			delete(r.liveCtxs, id)
    69  		}
    70  	}
    71  }