1package cancelctx 2 3import ( 4 "context" 5 "time" 6) 7 8// customCancelContext wraps context and cancels as soon as channel closed. 9type customCancelContext struct { 10 context.Context 11 ch <-chan struct{} 12} 13 14// Deadline not used. 15func (c customCancelContext) Deadline() (time.Time, bool) { return time.Time{}, false } 16 17// Done returns channel that will be closed as soon as connection closed. 18func (c customCancelContext) Done() <-chan struct{} { return c.ch } 19 20// Err returns context error. 21func (c customCancelContext) Err() error { 22 select { 23 case <-c.ch: 24 return context.Canceled 25 default: 26 return nil 27 } 28} 29 30// New returns a wrapper context around original context that will 31// be canceled on channel close. 32func New(ctx context.Context, ch <-chan struct{}) context.Context { 33 return customCancelContext{Context: ctx, ch: ch} 34} 35