Merge pull request #22 from goinsane/develop

v1.7.0

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
+}

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
 }