1package libkb 2 3import ( 4 "sync" 5 6 "golang.org/x/net/context" 7) 8 9type RPCCancelerReason uint 10 11const ( 12 RPCCancelerReasonLogout RPCCancelerReason = 1 << iota 13 RPCCancelerReasonBackground 14 RPCCancelerReasonMax 15) 16 17const RPCCancelerReasonAll = RPCCancelerReasonMax - 1 18 19type RPCCancelerKey string 20 21func NewRPCCancelerKey() RPCCancelerKey { 22 return RPCCancelerKey(RandStringB64(3)) 23} 24 25type liveContext struct { 26 ctx context.Context 27 cancelFn context.CancelFunc 28 reason RPCCancelerReason 29} 30 31type RPCCanceler struct { 32 sync.Mutex 33 liveCtxs map[RPCCancelerKey]liveContext 34} 35 36func NewRPCCanceler() *RPCCanceler { 37 return &RPCCanceler{ 38 liveCtxs: make(map[RPCCancelerKey]liveContext), 39 } 40} 41 42func (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 53func (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 62func (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} 72