From 12ba111b101f98396eed83cad5c95c537741bb62 Mon Sep 17 00:00:00 2001
From: Googol Lee <i@googol.im>
Date: Thu, 11 Apr 2024 16:35:29 +0200
Subject: [PATCH] Optimize repo.InjectTo()

---
 module/context.go | 23 +++++++++++++++++++----
 module/repo.go    | 39 ++++++++++++++++++++++++++++++---------
 test.sh           |  2 +-
 3 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/module/context.go b/module/context.go
index b4e7314..8312488 100644
--- a/module/context.go
+++ b/module/context.go
@@ -14,13 +14,13 @@ type createPanic struct {
 	err error
 }
 
-type moduleContext struct {
+type buildContext struct {
 	context.Context
-	providers map[moduleKey]Provider
+	providers map[moduleKey]providerWithLine
 	instances map[moduleKey]any
 }
 
-func (c *moduleContext) Value(key any) any {
+func (c *buildContext) Value(key any) any {
 	moduleKey, ok := key.(moduleKey)
 	if !ok {
 		return c.Context.Value(key)
@@ -35,10 +35,25 @@ func (c *moduleContext) Value(key any) any {
 		panic(createPanic{key: moduleKey, err: ErrNoPrivoder})
 	}
 
-	instance, err := provider.value(c)
+	instance, err := provider.provider.value(c)
 	if err != nil {
 		panic(createPanic{key: moduleKey, err: err})
 	}
 	c.instances[moduleKey] = instance
 	return instance
 }
+
+type moduleContext struct {
+	context.Context
+	instances map[moduleKey]any
+}
+
+func (c *moduleContext) Value(key any) any {
+	if moduleKey, ok := key.(moduleKey); ok {
+		if instance, ok := c.instances[moduleKey]; ok {
+			return instance
+		}
+	}
+
+	return c.Context.Value(key)
+}
diff --git a/module/repo.go b/module/repo.go
index ad8a3fa..20d68b5 100644
--- a/module/repo.go
+++ b/module/repo.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"runtime"
+	"sync"
 )
 
 type providerWithLine struct {
@@ -15,6 +16,8 @@ type providerWithLine struct {
 // Repo is a repository of modules, and to inject instances creating by modules into a context.
 type Repo struct {
 	providers map[moduleKey]providerWithLine
+
+	locker    sync.RWMutex // protects filed `instances` below
 	instances map[moduleKey]any
 }
 
@@ -43,24 +46,42 @@ func (r *Repo) Add(provider Provider) {
 
 // InjectTo injects instances created by modules into a context `ctx`.
 // It returns a new context with all injections. If any module creates an instance with an error, `InjectTo` returns that error with the module name.
-func (r *Repo) InjectTo(ctx context.Context) (ret context.Context, err error) {
+// Injecting instances only create once if necessary. Calling `InjectTo` mutlple times share instances between returning contexts.
+// InjectTo ignores all new providers adding to the Repo after the first run. So adding all providers before calling `InjectTo`.
+func (r *Repo) InjectTo(ctx context.Context) (context.Context, error) {
+	r.locker.RLock()
+	needCreating := len(r.instances) == 0
+	r.locker.RUnlock()
+
+	if needCreating {
+		r.locker.Lock()
+		err := r.buildValues(ctx)
+		r.locker.Unlock()
+
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return &moduleContext{
+		Context:   ctx,
+		instances: r.instances,
+	}, nil
+}
+
+func (r *Repo) buildValues(ctx context.Context) (err error) {
 	defer func() {
 		err = r.catchError(recover())
 	}()
 
-	providers := make(map[moduleKey]Provider)
-	for k, p := range r.providers {
-		providers[k] = p.provider
-	}
-
-	ret = &moduleContext{
+	builder := &buildContext{
 		Context:   ctx,
-		providers: providers,
+		providers: r.providers,
 		instances: r.instances,
 	}
 
 	for key := range r.providers {
-		_ = ret.Value(key)
+		_ = builder.Value(key)
 	}
 
 	return
diff --git a/test.sh b/test.sh
index 5b7013e..11f8d24 100755
--- a/test.sh
+++ b/test.sh
@@ -1,3 +1,3 @@
 #!/bin/sh
 
-GODEBUG=httpmuxgo121=0 go test ./... -v
+GODEBUG=httpmuxgo121=0 go test ./... -v -race