diff --git a/cancelablecontext.go b/cancelablecontext.go new file mode 100644 index 0000000..3efd344 --- /dev/null +++ b/cancelablecontext.go @@ -0,0 +1,55 @@ +package xcontext + +import ( + "context" +) + +// CancelableContext is a custom implementation of context.Context with the Cancel() method to cancel context.Context. +type CancelableContext interface { + context.Context + Cancel() +} + +type cancelableContext struct { + context.Context + context.CancelFunc +} + +// Cancel calls the cancel function of the underlying context. +func (c *cancelableContext) Cancel() { + c.CancelFunc() +} + +// NewCancelableContext returns the underlying context as CancelableContext. +// The code should call CancelableContext.Cancel method or the cancel function to release resources associated with it. +func NewCancelableContext(ctx context.Context, cancel context.CancelFunc) (CancelableContext, context.CancelFunc) { + result := new(cancelableContext) + result.Context, result.CancelFunc = ctx, cancel + return result, cancel +} + +// WithCancelable creates a new cancel context as CancelableContext. +// The code should call CancelableContext.Cancel method or the cancel function to release resources associated with it. +func WithCancelable(parent context.Context) (CancelableContext, context.CancelFunc) { + return NewCancelableContext(context.WithCancel(parent)) +} + +// WithCancelable2 is similar with WithCancelable. +// It returns only a new context inherited from parent. +func WithCancelable2(parent context.Context) CancelableContext { + ctx, _ := WithCancelable(parent) + return ctx +} + +// WithCancelableAutoCancel is similar with WithCancelable. +// But it cancels the context when it was done through parent. +func WithCancelableAutoCancel(parent context.Context) (CancelableContext, context.CancelFunc) { + return NewCancelableContext(AutoCancel(context.WithCancel(parent))) +} + +// WithCancelableAutoCancel2 is similar with WithCancelableAutoCancel. +// It returns only a new context inherited from parent. +func WithCancelableAutoCancel2(parent context.Context) CancelableContext { + ctx, _ := WithCancelableAutoCancel(parent) + return ctx +} diff --git a/terminatecontext.go b/terminatecontext.go deleted file mode 100644 index b5df935..0000000 --- a/terminatecontext.go +++ /dev/null @@ -1,55 +0,0 @@ -package xcontext - -import ( - "context" -) - -// TerminateContext is a custom implementation of context.Context with Terminate() method to cancel context.Context. -type TerminateContext interface { - context.Context - Terminate() -} - -type terminateContext struct { - context.Context - context.CancelFunc -} - -// Terminate calls cancel function of the underlying context. -func (c *terminateContext) Terminate() { - c.CancelFunc() -} - -// NewTerminateContext returns the underlying context as TerminateContext. -// The code should call TerminateContext.Terminate method or cancel function to release resources associated with it. -func NewTerminateContext(ctx context.Context, cancel context.CancelFunc) (TerminateContext, context.CancelFunc) { - result := new(terminateContext) - result.Context, result.CancelFunc = ctx, cancel - return result, cancel -} - -// WithTerminate creates a new cancel context as TerminateContext. -// The code should call TerminateContext.Terminate method or cancel function to release resources associated with it. -func WithTerminate(parent context.Context) (TerminateContext, context.CancelFunc) { - return NewTerminateContext(context.WithCancel(parent)) -} - -// WithTerminate2 is similar with WithTerminate. -// It returns only a new context inherited from parent. -func WithTerminate2(parent context.Context) TerminateContext { - ctx, _ := WithTerminate(parent) - return ctx -} - -// WithTerminateAutoCancel is similar with WithTerminate. -// But it cancels the context when it was done through parent. -func WithTerminateAutoCancel(parent context.Context) (TerminateContext, context.CancelFunc) { - return NewTerminateContext(AutoCancel(context.WithCancel(parent))) -} - -// WithTerminateAutoCancel2 is similar with WithTerminateAutoCancel. -// It returns only a new context inherited from parent. -func WithTerminateAutoCancel2(parent context.Context) TerminateContext { - ctx, _ := WithTerminateAutoCancel(parent) - return ctx -} diff --git a/xcontext.go b/xcontext.go index 964eb31..0fb78ff 100644 --- a/xcontext.go +++ b/xcontext.go @@ -13,10 +13,11 @@ import ( // Calling the cancel function is not necessary. func DelayContext(ctx context.Context, delay time.Duration) (context.Context, context.CancelFunc) { newCtx, newCtxCancel := context.WithCancel(context.Background()) + delayTimer := time.NewTimer(delay) go func() { parentCtx := ctx parentCtxOk := context.Background() - delayCh := time.After(delay) + delayCh := delayTimer.C delayOkCh := make(<-chan time.Time) for done := false; !done; { select { @@ -35,6 +36,7 @@ func DelayContext(ctx context.Context, delay time.Duration) (context.Context, co } } newCtxCancel() + delayTimer.Stop() }() return newCtx, newCtxCancel } @@ -51,16 +53,19 @@ func DelayContext2(ctx context.Context, delay time.Duration) context.Context { // Calling the cancel function is not necessary. func DelayAfterContext(ctx context.Context, delay time.Duration) (context.Context, context.CancelFunc) { newCtx, newCtxCancel := context.WithCancel(context.Background()) + delayTimer := time.NewTimer(0) + if !delayTimer.Stop() { + <-delayTimer.C + } go func() { parentCtx := ctx parentCtxOk := context.Background() - delayCh := make(<-chan time.Time) for done := false; !done; { select { case <-parentCtx.Done(): parentCtx = parentCtxOk - delayCh = time.After(delay) - case <-delayCh: + delayTimer.Reset(delay) + case <-delayTimer.C: if parentCtx == parentCtxOk { done = true } @@ -69,6 +74,7 @@ func DelayAfterContext(ctx context.Context, delay time.Duration) (context.Contex } } newCtxCancel() + delayTimer.Stop() }() return newCtx, newCtxCancel }