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