Modules: Concurrency issue with module context fix

exchange/exchange_test.go

@@ -5996,7 +5996,7 @@ func (e mockUpdateBidRequestHook) HandleBidderRequestHook(_ context.Context, mct
 		}, hookstage.MutationUpdate, "bidRequest", "site.domain",
 	)
 
-	mctx.ModuleContext = map[string]interface{}{"some-ctx": "some-ctx"}
+	mctx.ModuleContext.Set("some-ctx", "some-ctx")
 
 	return hookstage.HookResult[hookstage.BidderRequestPayload]{ChangeSet: c, ModuleContext: mctx.ModuleContext}, nil
 }

hooks/hookexecution/context.go

@@ -43,24 +43,27 @@ func (ctx executionContext) getModuleContext(moduleName string) hookstage.Module
 // moduleContexts preserves data the module wants to pass to itself from earlier stages to later stages.
 type moduleContexts struct {
 	sync.RWMutex
-	ctxs map[string]hookstage.ModuleContext // format: {"module_name": hookstage.ModuleContext}
+	ctxs map[string]*hookstage.ModuleContext // format: {"module_name": hookstage.ModuleContext}
 }
 
-func (mc *moduleContexts) put(moduleName string, mCtx hookstage.ModuleContext) {
+func (mc *moduleContexts) put(moduleName string, mCtx *hookstage.ModuleContext) {
+	if mc == nil {
+		return
+	}
 	mc.Lock()
 	defer mc.Unlock()
 
-	newCtx := mCtx
-	if existingCtx, ok := mc.ctxs[moduleName]; ok && existingCtx != nil {
-		for k, v := range mCtx {
-			existingCtx[k] = v
-		}
-		newCtx = existingCtx
+	existingCtx, ok := mc.ctxs[moduleName]
+	if !ok {
+		mc.ctxs[moduleName] = mCtx
+		return
 	}
-	mc.ctxs[moduleName] = newCtx
+
+	// Add new data to existing context
+	existingCtx.SetAll(mCtx.GetAll())
 }
 
-func (mc *moduleContexts) get(moduleName string) (hookstage.ModuleContext, bool) {
+func (mc *moduleContexts) get(moduleName string) (*hookstage.ModuleContext, bool) {
 	mc.RLock()
 	defer mc.RUnlock()
 	mCtx, ok := mc.ctxs[moduleName]
@@ -72,4 +75,4 @@ type stageModuleContext struct {
 	groupCtx []groupModuleContext
 }
 
-type groupModuleContext map[string]hookstage.ModuleContext
+type groupModuleContext map[string]*hookstage.ModuleContext

hooks/hookexecution/executor.go

@@ -68,7 +68,7 @@ func NewHookExecutor(builder hooks.ExecutionPlanBuilder, endpoint string, me met
 		endpoint:       endpoint,
 		planBuilder:    builder,
 		stageOutcomes:  []StageOutcome{},
-		moduleContexts: &moduleContexts{ctxs: make(map[string]hookstage.ModuleContext)},
+		moduleContexts: &moduleContexts{ctxs: make(map[string]*hookstage.ModuleContext)},
 		metricEngine:   me,
 	}
 }

hooks/hookexecution/mocks_test.go

@@ -201,42 +201,50 @@ type mockModuleContextHook struct {
 }
 
 func (e mockModuleContextHook) HandleEntrypointHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.EntrypointPayload) (hookstage.HookResult[hookstage.EntrypointPayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.EntrypointPayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 
 func (e mockModuleContextHook) HandleRawAuctionHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.RawAuctionRequestPayload) (hookstage.HookResult[hookstage.RawAuctionRequestPayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.RawAuctionRequestPayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 
 func (e mockModuleContextHook) HandleProcessedAuctionHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.ProcessedAuctionRequestPayload) (hookstage.HookResult[hookstage.ProcessedAuctionRequestPayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.ProcessedAuctionRequestPayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 
 func (e mockModuleContextHook) HandleBidderRequestHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.BidderRequestPayload) (hookstage.HookResult[hookstage.BidderRequestPayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.BidderRequestPayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 
 func (e mockModuleContextHook) HandleRawBidderResponseHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.RawBidderResponsePayload) (hookstage.HookResult[hookstage.RawBidderResponsePayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.RawBidderResponsePayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 
 func (e mockModuleContextHook) HandleAllProcessedBidResponsesHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.AllProcessedBidResponsesPayload) (hookstage.HookResult[hookstage.AllProcessedBidResponsesPayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.AllProcessedBidResponsesPayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 
 func (e mockModuleContextHook) HandleAuctionResponseHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.AuctionResponsePayload) (hookstage.HookResult[hookstage.AuctionResponsePayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.AuctionResponsePayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 
 func (e mockModuleContextHook) HandleExitpointHook(_ context.Context, miCtx hookstage.ModuleInvocationContext, _ hookstage.ExitpointPayload) (hookstage.HookResult[hookstage.ExitpointPayload], error) {
-	miCtx.ModuleContext = map[string]interface{}{e.key: e.val}
+	miCtx.ModuleContext = hookstage.NewModuleContext()
+	miCtx.ModuleContext.Set(e.key, e.val)
 	return hookstage.HookResult[hookstage.ExitpointPayload]{ModuleContext: miCtx.ModuleContext}, nil
 }
 

hooks/hookstage/invocation.go

@@ -2,6 +2,8 @@ package hookstage
 
 import (
 	"encoding/json"
+	"maps"
+	"sync"
 
 	"github.com/prebid/prebid-server/v4/hooks/hookanalytics"
 )
@@ -16,7 +18,7 @@ type HookResult[T any] struct {
 	Warnings      []string
 	DebugMessages []string
 	AnalyticsTags hookanalytics.Analytics
-	ModuleContext ModuleContext // holds values that the module wants to pass to itself at later stages
+	ModuleContext *ModuleContext // holds values that the module wants to pass to itself at later stages
 }
 
 // ModuleInvocationContext holds data passed to the module hook during invocation.
@@ -28,11 +30,77 @@ type ModuleInvocationContext struct {
 	// Endpoint represents the path of the current endpoint.
 	Endpoint string
 	// ModuleContext holds values that the module passes to itself from the previous stages.
-	ModuleContext ModuleContext
+	ModuleContext *ModuleContext
 	// HookImplCode is the hook_impl_code for a module instance to differentiate between multiple hooks
 	HookImplCode string
 }
 
 // ModuleContext holds arbitrary data passed between module hooks at different stages.
 // We use interface as we do not know exactly how the modules will use their inner context.
-type ModuleContext map[string]interface{}
+type ModuleContext struct {
+	sync.RWMutex
+	data map[string]any
+}
+
+// NewModuleContext creates a new module context
+func NewModuleContext() *ModuleContext {
+	moduleContext := ModuleContext{
+		data: make(map[string]any),
+	}
+	return &moduleContext
+}
+
+// Get retrieves a value from the module context with read lock
+func (mc *ModuleContext) Get(key string) (any, bool) {
+	if mc == nil {
+		return nil, false
+	}
+	mc.RLock()
+	defer mc.RUnlock()
+	if mc.data == nil {
+		return nil, false
+	}
+	val, ok := mc.data[key]
+	return val, ok
+}
+
+// Set stores a value in the module context with write lock
+func (mc *ModuleContext) Set(key string, value any) {
+	if mc == nil {
+		return
+	}
+	mc.Lock()
+	defer mc.Unlock()
+	if mc.data == nil {
+		mc.data = make(map[string]any)
+	}
+	mc.data[key] = value
+}
+
+// GetAll returns a copy of all data in the context
+func (mc *ModuleContext) GetAll() map[string]any {
+	if mc == nil {
+		return nil
+	}
+	mc.RLock()
+	defer mc.RUnlock()
+	if mc.data == nil {
+		return nil
+	}
+	result := make(map[string]any, len(mc.data))
+	maps.Copy(result, mc.data)
+	return result
+}
+
+// SetAll merges the provided data into the module context
+func (mc *ModuleContext) SetAll(data map[string]any) {
+	if mc == nil {
+		return
+	}
+	mc.Lock()
+	defer mc.Unlock()
+	if mc.data == nil {
+		mc.data = make(map[string]any, len(data))
+	}
+	maps.Copy(mc.data, data)
+}

modules/fiftyonedegrees/devicedetection/evidence_extractor.go

@@ -57,16 +57,18 @@ func merge(val1, val2 []stringEvidence) []stringEvidence {
 	return evidence
 }
 
-func (x *defaultEvidenceExtractor) extract(ctx hookstage.ModuleContext) ([]onpremise.Evidence, string, error) {
+func (x *defaultEvidenceExtractor) extract(ctx *hookstage.ModuleContext) ([]onpremise.Evidence, string, error) {
 	if ctx == nil {
 		return nil, "", errors.New("context is nil")
 	}
 
-	suaStrings, err := x.getEvidenceStrings(ctx[evidenceFromSuaCtxKey])
+	evidenceFromSuaCtx, _ := ctx.Get(evidenceFromSuaCtxKey)
+	suaStrings, err := x.getEvidenceStrings(evidenceFromSuaCtx)
 	if err != nil {
 		return nil, "", fmt.Errorf("error extracting sua evidence: %w", err)
 	}
-	headerString, err := x.getEvidenceStrings(ctx[evidenceFromHeadersCtxKey])
+	evidenceFromHeadersCtx, _ := ctx.Get(evidenceFromHeadersCtxKey)
+	headerString, err := x.getEvidenceStrings(evidenceFromHeadersCtx)
 	if err != nil {
 		return nil, "", fmt.Errorf("error extracting header evidence: %w", err)
 	}