add init ctx feature (#1995)

os/gctx/gctx.go

@@ -23,8 +23,8 @@ type (
 )
 
 var (
-	// processCtx is the context initialized from process environment.
-	processCtx context.Context
+	processCtx context.Context // processCtx is the context initialized from process environment.
+	initCtx    context.Context // initCtx is the context for init function of packages.
 )
 
 func init() {
@@ -40,6 +40,8 @@ func init() {
 		context.Background(),
 		propagation.MapCarrier(m),
 	)
+	// Initialize initialization context.
+	initCtx = New()
 }
 
 // New creates and returns a context which contains context id.
@@ -64,3 +66,14 @@ func WithCtx(ctx context.Context) context.Context {
 func CtxId(ctx context.Context) string {
 	return gtrace.GetTraceID(ctx)
 }
+
+// SetInitCtx sets custom initialization context.
+// Note that this function cannot be called in multiple goroutines.
+func SetInitCtx(ctx context.Context) {
+	initCtx = ctx
+}
+
+// GetInitCtx returns the initialization context.
+func GetInitCtx() context.Context {
+	return initCtx
+}

os/gctx/gctx_z_unit_test.go

@@ -30,3 +30,12 @@ func Test_WithCtx(t *testing.T) {
 		t.Assert(ctx.Value("TEST"), 1)
 	})
 }
+
+func Test_SetInitCtx(t *testing.T) {
+	gtest.C(t, func(t *gtest.T) {
+		ctx := context.WithValue(context.TODO(), "TEST", 1)
+		gctx.SetInitCtx(ctx)
+		t.AssertNE(gctx.GetInitCtx(), "")
+		t.Assert(gctx.GetInitCtx().Value("TEST"), 1)
+	})
+}