fix: context key error in sql read write mode (#5000)

kevwan · web-flow · commit a71e56de52b6 · 2025-07-12T06:58:08.000+08:00
diff --git a/core/logx/fields.go b/core/logx/fields.go
@@ -7,12 +7,11 @@ import (
 )
 
 var (
-	fieldsContextKey contextKey
 	globalFields     atomic.Value
 	globalFieldsLock sync.Mutex
 )
 
-type contextKey struct{}
+type fieldsKey struct{}
 
 // AddGlobalFields adds global fields.
 func AddGlobalFields(fields ...LogField) {
@@ -29,16 +28,16 @@ func AddGlobalFields(fields ...LogField) {
 
 // ContextWithFields returns a new context with the given fields.
 func ContextWithFields(ctx context.Context, fields ...LogField) context.Context {
-	if val := ctx.Value(fieldsContextKey); val != nil {
+	if val := ctx.Value(fieldsKey{}); val != nil {
 		if arr, ok := val.([]LogField); ok {
 			allFields := make([]LogField, 0, len(arr)+len(fields))
 			allFields = append(allFields, arr...)
 			allFields = append(allFields, fields...)
-			return context.WithValue(ctx, fieldsContextKey, allFields)
+			return context.WithValue(ctx, fieldsKey{}, allFields)
 		}
 	}
 
-	return context.WithValue(ctx, fieldsContextKey, fields)
+	return context.WithValue(ctx, fieldsKey{}, fields)
 }
 
 // WithFields returns a new logger with the given fields.
diff --git a/core/logx/fields_test.go b/core/logx/fields_test.go
@@ -34,7 +34,7 @@ func TestAddGlobalFields(t *testing.T) {
 
 func TestContextWithFields(t *testing.T) {
 	ctx := ContextWithFields(context.Background(), Field("a", 1), Field("b", 2))
-	vals := ctx.Value(fieldsContextKey)
+	vals := ctx.Value(fieldsKey{})
 	assert.NotNil(t, vals)
 	fields, ok := vals.([]LogField)
 	assert.True(t, ok)
@@ -43,7 +43,7 @@ func TestContextWithFields(t *testing.T) {
 
 func TestWithFields(t *testing.T) {
 	ctx := WithFields(context.Background(), Field("a", 1), Field("b", 2))
-	vals := ctx.Value(fieldsContextKey)
+	vals := ctx.Value(fieldsKey{})
 	assert.NotNil(t, vals)
 	fields, ok := vals.([]LogField)
 	assert.True(t, ok)
@@ -55,7 +55,7 @@ func TestWithFieldsAppend(t *testing.T) {
 	ctx := context.WithValue(context.Background(), dummyKey, "dummy")
 	ctx = ContextWithFields(ctx, Field("a", 1), Field("b", 2))
 	ctx = ContextWithFields(ctx, Field("c", 3), Field("d", 4))
-	vals := ctx.Value(fieldsContextKey)
+	vals := ctx.Value(fieldsKey{})
 	assert.NotNil(t, vals)
 	fields, ok := vals.([]LogField)
 	assert.True(t, ok)
@@ -80,8 +80,8 @@ func TestWithFieldsAppendCopy(t *testing.T) {
 	ctxa := ContextWithFields(ctx, af)
 	ctxb := ContextWithFields(ctx, bf)
 
-	assert.EqualValues(t, af, ctxa.Value(fieldsContextKey).([]LogField)[count])
-	assert.EqualValues(t, bf, ctxb.Value(fieldsContextKey).([]LogField)[count])
+	assert.EqualValues(t, af, ctxa.Value(fieldsKey{}).([]LogField)[count])
+	assert.EqualValues(t, bf, ctxb.Value(fieldsKey{}).([]LogField)[count])
 }
 
 func BenchmarkAtomicValue(b *testing.B) {
diff --git a/core/logx/richlogger.go b/core/logx/richlogger.go
@@ -224,7 +224,7 @@ func (l *richLogger) buildFields(fields ...LogField) []LogField {
 		fields = append(fields, Field(spanKey, spanID))
 	}
 
-	val := l.ctx.Value(fieldsContextKey)
+	val := l.ctx.Value(fieldsKey{})
 	if val != nil {
 		if arr, ok := val.([]LogField); ok {
 			fields = append(fields, arr...)
diff --git a/core/stores/sqlx/rwstrategy.go b/core/stores/sqlx/rwstrategy.go
@@ -27,21 +27,21 @@ const (
 	notSpecifiedMode readWriteMode = ""
 )
 
-var readWriteModeKey struct{}
+type readWriteModeKey struct{}
 
 // WithReadPrimary sets the context to read-primary mode.
 func WithReadPrimary(ctx context.Context) context.Context {
-	return context.WithValue(ctx, readWriteModeKey, readPrimaryMode)
+	return context.WithValue(ctx, readWriteModeKey{}, readPrimaryMode)
 }
 
 // WithReadReplica sets the context to read-replica mode.
 func WithReadReplica(ctx context.Context) context.Context {
-	return context.WithValue(ctx, readWriteModeKey, readReplicaMode)
+	return context.WithValue(ctx, readWriteModeKey{}, readReplicaMode)
 }
 
 // WithWrite sets the context to write mode, indicating that the operation is a write operation.
 func WithWrite(ctx context.Context) context.Context {
-	return context.WithValue(ctx, readWriteModeKey, writeMode)
+	return context.WithValue(ctx, readWriteModeKey{}, writeMode)
 }
 
 type readWriteMode string
@@ -51,7 +51,7 @@ func (m readWriteMode) isValid() bool {
 }
 
 func getReadWriteMode(ctx context.Context) readWriteMode {
-	if mode := ctx.Value(readWriteModeKey); mode != nil {
+	if mode := ctx.Value(readWriteModeKey{}); mode != nil {
 		if v, ok := mode.(readWriteMode); ok && v.isValid() {
 			return v
 		}
diff --git a/core/stores/sqlx/rwstrategy_test.go b/core/stores/sqlx/rwstrategy_test.go
@@ -57,45 +57,45 @@ func TestWithReadMode(t *testing.T) {
 	ctx := context.Background()
 	readPrimaryCtx := WithReadPrimary(ctx)
 
-	val := readPrimaryCtx.Value(readWriteModeKey)
+	val := readPrimaryCtx.Value(readWriteModeKey{})
 	assert.Equal(t, readPrimaryMode, val)
 
 	readReplicaCtx := WithReadReplica(ctx)
-	val = readReplicaCtx.Value(readWriteModeKey)
+	val = readReplicaCtx.Value(readWriteModeKey{})
 	assert.Equal(t, readReplicaMode, val)
 }
 
 func TestWithWriteMode(t *testing.T) {
 	ctx := context.Background()
 	writeCtx := WithWrite(ctx)
 
-	val := writeCtx.Value(readWriteModeKey)
+	val := writeCtx.Value(readWriteModeKey{})
 	assert.Equal(t, writeMode, val)
 }
 
 func TestGetReadWriteMode(t *testing.T) {
 	t.Run("valid read-primary mode", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, readPrimaryMode)
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, readPrimaryMode)
 		assert.Equal(t, readPrimaryMode, getReadWriteMode(ctx))
 	})
 
 	t.Run("valid read-replica mode", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, readReplicaMode)
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, readReplicaMode)
 		assert.Equal(t, readReplicaMode, getReadWriteMode(ctx))
 	})
 
 	t.Run("valid write mode", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, writeMode)
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, writeMode)
 		assert.Equal(t, writeMode, getReadWriteMode(ctx))
 	})
 
 	t.Run("invalid mode value (wrong type)", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, "not-a-mode")
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, "not-a-mode")
 		assert.Equal(t, notSpecifiedMode, getReadWriteMode(ctx))
 	})
 
 	t.Run("invalid mode value (wrong value)", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, readWriteMode("delete"))
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, readWriteMode("delete"))
 		assert.Equal(t, notSpecifiedMode, getReadWriteMode(ctx))
 	})
 
@@ -107,22 +107,22 @@ func TestGetReadWriteMode(t *testing.T) {
 
 func TestUsePrimary(t *testing.T) {
 	t.Run("context with read-replica mode", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, readReplicaMode)
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, readReplicaMode)
 		assert.False(t, usePrimary(ctx))
 	})
 
 	t.Run("context with read-primary mode", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, readPrimaryMode)
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, readPrimaryMode)
 		assert.True(t, usePrimary(ctx))
 	})
 
 	t.Run("context with write mode", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, writeMode)
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, writeMode)
 		assert.True(t, usePrimary(ctx))
 	})
 
 	t.Run("context with invalid mode", func(t *testing.T) {
-		ctx := context.WithValue(context.Background(), readWriteModeKey, readWriteMode("invalid"))
+		ctx := context.WithValue(context.Background(), readWriteModeKey{}, readWriteMode("invalid"))
 		assert.True(t, usePrimary(ctx))
 	})
 
@@ -137,6 +137,6 @@ func TestWithModeTwice(t *testing.T) {
 	ctx = WithReadPrimary(ctx)
 	writeCtx := WithWrite(ctx)
 
-	val := writeCtx.Value(readWriteModeKey)
+	val := writeCtx.Value(readWriteModeKey{})
 	assert.Equal(t, writeMode, val)
 }

Original file line number	Diff line number	Diff line change
`@@ -7,12 +7,11 @@ import (`
`7`	`7`	`)`
`8`	`8`
`9`	`9`	`var (`
`10`		`- fieldsContextKey contextKey`
`11`	`10`	`globalFields atomic.Value`
`12`	`11`	`globalFieldsLock sync.Mutex`
`13`	`12`	`)`
`14`	`13`
`15`		`-type contextKey struct{}`
	`14`	`+type fieldsKey struct{}`
`16`	`15`
`17`	`16`	`// AddGlobalFields adds global fields.`
`18`	`17`	`func AddGlobalFields(fields ...LogField) {`
`@@ -29,16 +28,16 @@ func AddGlobalFields(fields ...LogField) {`
`29`	`28`
`30`	`29`	`// ContextWithFields returns a new context with the given fields.`
`31`	`30`	`func ContextWithFields(ctx context.Context, fields ...LogField) context.Context {`
`32`		`- if val := ctx.Value(fieldsContextKey); val != nil {`
	`31`	`+ if val := ctx.Value(fieldsKey{}); val != nil {`
`33`	`32`	`if arr, ok := val.([]LogField); ok {`
`34`	`33`	`allFields := make([]LogField, 0, len(arr)+len(fields))`
`35`	`34`	`allFields = append(allFields, arr...)`
`36`	`35`	`allFields = append(allFields, fields...)`
`37`		`- return context.WithValue(ctx, fieldsContextKey, allFields)`
	`36`	`+ return context.WithValue(ctx, fieldsKey{}, allFields)`
`38`	`37`	`}`
`39`	`38`	`}`
`40`	`39`
`41`		`- return context.WithValue(ctx, fieldsContextKey, fields)`
	`40`	`+ return context.WithValue(ctx, fieldsKey{}, fields)`
`42`	`41`	`}`
`43`	`42`
`44`	`43`	`// WithFields returns a new logger with the given fields.`
Original file line number	Diff line number	Diff line change
`@@ -224,7 +224,7 @@ func (l *richLogger) buildFields(fields ...LogField) []LogField {`
`224`	`224`	`fields = append(fields, Field(spanKey, spanID))`
`225`	`225`	`}`
`226`	`226`
`227`		`- val := l.ctx.Value(fieldsContextKey)`
	`227`	`+ val := l.ctx.Value(fieldsKey{})`
`228`	`228`	`if val != nil {`
`229`	`229`	`if arr, ok := val.([]LogField); ok {`
`230`	`230`	`fields = append(fields, arr...)`
Original file line number	Diff line number	Diff line change
`@@ -27,21 +27,21 @@ const (`
`27`	`27`	`notSpecifiedMode readWriteMode = ""`
`28`	`28`	`)`
`29`	`29`
`30`		`-var readWriteModeKey struct{}`
	`30`	`+type readWriteModeKey struct{}`
`31`	`31`
`32`	`32`	`// WithReadPrimary sets the context to read-primary mode.`
`33`	`33`	`func WithReadPrimary(ctx context.Context) context.Context {`
`34`		`- return context.WithValue(ctx, readWriteModeKey, readPrimaryMode)`
	`34`	`+ return context.WithValue(ctx, readWriteModeKey{}, readPrimaryMode)`
`35`	`35`	`}`
`36`	`36`
`37`	`37`	`// WithReadReplica sets the context to read-replica mode.`
`38`	`38`	`func WithReadReplica(ctx context.Context) context.Context {`
`39`		`- return context.WithValue(ctx, readWriteModeKey, readReplicaMode)`
	`39`	`+ return context.WithValue(ctx, readWriteModeKey{}, readReplicaMode)`
`40`	`40`	`}`
`41`	`41`
`42`	`42`	`// WithWrite sets the context to write mode, indicating that the operation is a write operation.`
`43`	`43`	`func WithWrite(ctx context.Context) context.Context {`
`44`		`- return context.WithValue(ctx, readWriteModeKey, writeMode)`
	`44`	`+ return context.WithValue(ctx, readWriteModeKey{}, writeMode)`
`45`	`45`	`}`
`46`	`46`
`47`	`47`	`type readWriteMode string`
`@@ -51,7 +51,7 @@ func (m readWriteMode) isValid() bool {`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`func getReadWriteMode(ctx context.Context) readWriteMode {`
`54`		`- if mode := ctx.Value(readWriteModeKey); mode != nil {`
	`54`	`+ if mode := ctx.Value(readWriteModeKey{}); mode != nil {`
`55`	`55`	`if v, ok := mode.(readWriteMode); ok && v.isValid() {`
`56`	`56`	`return v`
`57`	`57`	`}`