35 lines
848 B
Go
35 lines
848 B
Go
package uniws
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
)
|
|
|
|
// customCancelContext wraps context and cancels as soon as channel closed.
|
|
type customCancelContext struct {
|
|
context.Context
|
|
ch <-chan struct{}
|
|
}
|
|
|
|
// Deadline not used.
|
|
func (c customCancelContext) Deadline() (time.Time, bool) { return time.Time{}, false }
|
|
|
|
// Done returns channel that will be closed as soon as connection closed.
|
|
func (c customCancelContext) Done() <-chan struct{} { return c.ch }
|
|
|
|
// Err returns context error.
|
|
func (c customCancelContext) Err() error {
|
|
select {
|
|
case <-c.ch:
|
|
return context.Canceled
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// NewCancelContext returns a wrapper context around original context that will
|
|
// be canceled on channel close.
|
|
func NewCancelContext(ctx context.Context, ch <-chan struct{}) context.Context {
|
|
return customCancelContext{Context: ctx, ch: ch}
|
|
}
|