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