github.com/keybase/client/go@v0.0.0-20241007131713-f10651d043c8/kbfs/libkbfs/coalescing_context.go (about) 1 package libkbfs 2 3 import ( 4 "reflect" 5 "sync" 6 "time" 7 8 "golang.org/x/net/context" 9 ) 10 11 // CoalescingContext allows many contexts to be treated as one. It waits on 12 // all its contexts' Context.Done() channels, and when all of them have 13 // returned, this CoalescingContext is canceled. At any point, a context can be 14 // added to the list, and will subsequently also be part of the wait condition. 15 // TODO: add timeout channel in case there is a goroutine leak. 16 type CoalescingContext struct { 17 context.Context 18 doneCh chan struct{} 19 mutateCh chan context.Context 20 selects []reflect.SelectCase 21 start sync.Once 22 } 23 24 const ( 25 mutateChanSelectIndex int = 0 26 closeChanSelectIndex int = 1 27 numExplicitlyHandledSelects int = 2 28 ) 29 30 func (ctx *CoalescingContext) loop() { 31 for { 32 chosen, val, _ := reflect.Select(ctx.selects) 33 switch chosen { 34 case mutateChanSelectIndex: 35 // request to mutate the select list 36 newCase := val.Interface().(context.Context) 37 if newCase != nil { 38 ctx.appendContext(newCase) 39 } 40 case closeChanSelectIndex: 41 // Done 42 close(ctx.doneCh) 43 return 44 default: 45 // The chosen channel has been closed. Remove it from our select list. 46 ctx.selects = append(ctx.selects[:chosen], ctx.selects[chosen+1:]...) 47 // If we have no more selects available, the request is done. 48 if len(ctx.selects) == numExplicitlyHandledSelects { 49 close(ctx.doneCh) 50 return 51 } 52 } 53 } 54 } 55 56 func (ctx *CoalescingContext) appendContext(other context.Context) { 57 ctx.selects = append(ctx.selects, reflect.SelectCase{ 58 Dir: reflect.SelectRecv, 59 Chan: reflect.ValueOf(other.Done()), 60 }) 61 } 62 63 // NewCoalescingContext creates a new CoalescingContext. The context _must_ be 64 // canceled to avoid a goroutine leak. 65 func NewCoalescingContext(parent context.Context) (*CoalescingContext, context.CancelFunc) { 66 ctx := &CoalescingContext{ 67 // Make the parent's `Value()` method available to consumers of this 68 // context. For example, this maintains the parent's log debug tags. 69 // TODO: Make _all_ parents' values available. 70 Context: parent, 71 doneCh: make(chan struct{}), 72 mutateCh: make(chan context.Context), 73 } 74 closeCh := make(chan struct{}) 75 ctx.selects = []reflect.SelectCase{ 76 { 77 Dir: reflect.SelectRecv, 78 Chan: reflect.ValueOf(ctx.mutateCh), 79 }, 80 { 81 Dir: reflect.SelectRecv, 82 Chan: reflect.ValueOf(closeCh), 83 }, 84 } 85 ctx.appendContext(parent) 86 cancelFunc := func() { 87 select { 88 case <-closeCh: 89 default: 90 close(closeCh) 91 } 92 } 93 return ctx, cancelFunc 94 } 95 96 func (ctx *CoalescingContext) startLoop() { 97 ctx.start.Do(func() { 98 go ctx.loop() 99 }) 100 } 101 102 // Deadline overrides the default parent's Deadline(). 103 func (ctx *CoalescingContext) Deadline() (time.Time, bool) { 104 return time.Time{}, false 105 } 106 107 // Done returns a channel that is closed when the CoalescingContext is 108 // canceled. 109 func (ctx *CoalescingContext) Done() <-chan struct{} { 110 ctx.startLoop() 111 return ctx.doneCh 112 } 113 114 // Err returns context.Canceled if the CoalescingContext has been canceled, and 115 // nil otherwise. 116 func (ctx *CoalescingContext) Err() error { 117 ctx.startLoop() 118 select { 119 case <-ctx.doneCh: 120 return context.Canceled 121 default: 122 } 123 return nil 124 } 125 126 // AddContext adds a context to the set of contexts that we're waiting on. 127 func (ctx *CoalescingContext) AddContext(other context.Context) error { 128 ctx.startLoop() 129 select { 130 case ctx.mutateCh <- other: 131 return nil 132 case <-ctx.doneCh: 133 return context.Canceled 134 } 135 }