From f1766bdbdc2d8ec41654a3ed150ed01d7aead258 Mon Sep 17 00:00:00 2001
From: John Guo <john@johng.cn>
Date: Tue, 12 Jul 2022 19:27:42 +0800
Subject: [PATCH] add init ctx feature (#1995)

---
 os/gctx/gctx.go             | 17 +++++++++++++++--
 os/gctx/gctx_z_unit_test.go |  9 +++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/os/gctx/gctx.go b/os/gctx/gctx.go
index e04066ae9..02a8d6083 100644
--- a/os/gctx/gctx.go
+++ b/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
+}
diff --git a/os/gctx/gctx_z_unit_test.go b/os/gctx/gctx_z_unit_test.go
index 2c62a44da..aa09b92ff 100644
--- a/os/gctx/gctx_z_unit_test.go
+++ b/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)
+	})
+}