增加上下文时限管理能力

util/gopool/pool.go

@@ -16,6 +16,7 @@ package gopool
 
 import (
 	"context"
+	"github.com/bytedance/gopkg/util/hcontext"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -30,6 +31,8 @@ type Pool interface {
 	Go(f func())
 	// CtxGo executes f and accepts the context.
 	CtxGo(ctx context.Context, f func())
+	// CtxGoNoDeadlineOrCancel executes f and will be no timeouts and cancellations
+	CtxGoNoDeadlineOrCancel(ctx context.Context,f func())
 	// SetPanicHandler sets the panic handler.
 	SetPanicHandler(f func(context.Context, interface{}))
 }
@@ -138,6 +141,12 @@ func (p *pool) CtxGo(ctx context.Context, f func()) {
 	}
 }
 
+func (p *pool) CtxGoNoDeadlineOrCancel(ctx context.Context,f func())  {
+	ctx = hcontext.WithNoCancel(ctx)
+	ctx = hcontext.WithNoDeadline(ctx)
+	CtxGo(ctx,f)
+}
+
 // SetPanicHandler the func here will be called after the panic has been recovered.
 func (p *pool) SetPanicHandler(f func(context.Context, interface{})) {
 	p.panicHandler = f

util/hcontext/README.md

@@ -0,0 +1,15 @@
+# hcontext
+
+## 功能
+
+对 context.Context 的包装函数
+
+## usage
+
+```golang
+//移除超时信息，保留value信息
+ctx = hcontext.WithNoDeadline(ctx)
+
+//移除取消信息，保留超时和value信息
+ctx = hcontext.WithNoCancel(ctx)
+```

util/hcontext/hcontext.go

@@ -0,0 +1,30 @@
+package hcontext
+
+import (
+	"context"
+	"time"
+)
+
+type valueOnlyContext struct{ context.Context }
+
+func (valueOnlyContext) Deadline() (deadline time.Time, ok bool) { return }
+func (valueOnlyContext) Done() <-chan struct{}                   { return nil }
+func (valueOnlyContext) Err() error                              { return nil }
+
+//WithNoDeadline 移除超时控制，保留value信息
+func WithNoDeadline(ctx context.Context) context.Context {
+	return valueOnlyContext{ctx}
+}
+
+//WithNoCancel 不受上层context cancel的影响，但是保留上层ctx的Deadline信息
+func WithNoCancel(ctx context.Context) context.Context {
+	deadline, ok := ctx.Deadline()
+	if ok {
+		var cancel context.CancelFunc
+		ctx = WithNoDeadline(ctx)
+		ctx, cancel = context.WithDeadline(ctx, deadline)
+		_ = cancel
+		return ctx
+	}
+	return WithNoDeadline(ctx)
+}

util/hcontext/hcontext_test.go

@@ -0,0 +1,152 @@
+package hcontext
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"gopkg.in/go-playground/assert.v1"
+)
+
+type ctxtype string
+
+var k ctxtype = "k"
+
+func TestWithNoDeadline(t *testing.T) {
+	type args struct {
+		ctx context.Context
+	}
+	tests := []struct {
+		name string
+		args args
+		stub func(*testing.T, context.Context)
+	}{
+		{
+			name: "value正常传递",
+			args: args{
+				context.WithValue(context.TODO(), k, "world"),
+			},
+			stub: func(t *testing.T, ctx context.Context) {
+				assert.Equal(t, ctx.Value(k).(string), "world")
+			},
+		},
+		{
+			name: "nil Deadline()",
+			args: args{
+				ctx: func() context.Context {
+					ctx, cancel := context.WithTimeout(context.TODO(), time.Second)
+					_ = cancel
+					return ctx
+				}(),
+			},
+			stub: func(t *testing.T, ctx context.Context) {
+				deadline, ok := ctx.Deadline()
+				assert.Equal(t, deadline.IsZero(), true)
+				assert.Equal(t, ok, false)
+
+			},
+		},
+		{
+			name: "nil Done()",
+			args: args{
+				ctx: func() context.Context {
+					ctx, cancel := context.WithTimeout(context.TODO(), time.Microsecond)
+					_ = cancel
+					return ctx
+				}(),
+			},
+			stub: func(t *testing.T, ctx context.Context) {
+				assert.Equal(t, ctx.Done(), nil)
+			},
+		},
+		{
+			name: "nil Err()",
+			args: args{
+				ctx: func() context.Context {
+					ctx, cancel := context.WithTimeout(context.TODO(), time.Second)
+					cancel()
+					return ctx
+				}(),
+			},
+			stub: func(t *testing.T, ctx context.Context) {
+				assert.Equal(t, ctx.Err(), nil)
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := WithNoDeadline(tt.args.ctx)
+			tt.stub(t, got)
+		})
+	}
+}
+
+func TestWithNoCancel(t *testing.T) {
+	type args struct {
+		ctx context.Context
+	}
+	tests := []struct {
+		name string
+		args args
+		stub func(*testing.T, context.Context)
+	}{
+		{
+			name: "value正常传递",
+			args: args{
+				context.WithValue(context.TODO(), k, "world"),
+			},
+			stub: func(t *testing.T, ctx context.Context) {
+				assert.Equal(t, ctx.Value(k).(string), "world")
+			},
+		},
+
+		{
+			name: "new Deadline()",
+			args: args{
+				ctx: func() context.Context {
+					ctx, cancel := context.WithDeadline(context.TODO(), time.Date(2022, 1, 1, 1, 1, 1, 1, time.Local))
+					_ = cancel
+					return ctx
+				}(),
+			},
+			stub: func(t *testing.T, ctx context.Context) {
+				deadline, ok := ctx.Deadline()
+				assert.Equal(t, ok, true)
+				assert.Equal(t, deadline, time.Date(2022, 1, 1, 1, 1, 1, 1, time.Local))
+			},
+		},
+
+		{
+			name: "cancel withtimeout context",
+			args: args{
+				ctx: func() context.Context {
+					now := time.Now()
+					ctx := context.WithValue(context.Background(), "t", now)
+					ctx, cancel := context.WithDeadline(ctx, now.Add(1000*time.Microsecond))
+					go func() {
+						time.Sleep(100 * time.Microsecond)
+						cancel()
+					}()
+					return ctx
+				}(),
+			},
+
+			stub: func(t *testing.T, ctx context.Context) {
+				<-ctx.Done()
+				duration := time.Since(ctx.Value("t").(time.Time))
+				if ctx.Err() == context.Canceled {
+					t.Error("context canceled")
+				}
+				if duration < 200*time.Microsecond {
+					t.Error("timeout cancel ")
+				}
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := WithNoCancel(tt.args.ctx)
+			tt.stub(t, got)
+		})
+	}
+}