chore: update scope.go to support context-based disposal

Author: cng217

scope.go

@@ -1,6 +1,7 @@
 package submodule
 
 import (
+	"context"
 	"reflect"
 	"sync"
 )
@@ -34,7 +35,8 @@ type Scope interface {
 	InitError(g Retrievable, e error)
 
 	Dispose() error
-	AppendMiddleware(Middleware)
+	DisposeWithContext(ctx context.Context) error
+	AppendMiddleware(...Middleware)
 	Apply(Submodule[Middleware])
 }
 
@@ -140,44 +142,94 @@ func (s *scope) Apply(submodule Submodule[Middleware]) {
 	s.AppendMiddleware(m)
 }
 
-// Dispose scope to free up all resolved values and trigger scope end middlewares
-func (s *scope) Dispose() error {
+// remove all values in the scope
+func (s *scope) release() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for k := range s.values {
+		delete(s.values, k)
+	}
+}
+
+type middlewareCaller func(Middleware) error
+
+func (s *scope) dispose(cond middlewareCaller) error {
 	for _, m := range s.middleware {
 		if m.hasOnScopeEnd {
-			e := m.onScopeEnd()
-			if e != nil {
-				return e
+			if err := cond(m); err != nil {
+				return err
 			}
 		}
 	}
+	return nil
+}
 
-	s.mu.Lock()
-	defer s.mu.Unlock()
+var disposeCond = func(m Middleware) error {
+	if m.onScopeEnd != nil {
+		return m.onScopeEnd()
+	}
+	return nil
+}
 
-	for k := range s.values {
-		delete(s.values, k)
+var disposeWithContextCond = func(ctx context.Context) func(m Middleware) error {
+	return func(m Middleware) error {
+		if m.onScopeEndWithContext != nil {
+			return m.onScopeEndWithContext(ctx)
+		}
+		return nil
 	}
+}
 
+// DisposeWithContext dispose scope with context
+func (s *scope) DisposeWithContext(ctx context.Context) error {
+	if err := s.dispose(disposeWithContextCond(ctx)); err != nil {
+		return err
+	}
+	if err := s.dispose(disposeCond); err != nil {
+		return err
+	}
+	s.release()
+	return nil
+}
+
+// Dispose scope to free up all resolved values and trigger scope end middlewares
+func (s *scope) Dispose() error {
+	if err := s.dispose(disposeWithContextCond(context.TODO())); err != nil {
+		return err
+	}
+	if err := s.dispose(disposeCond); err != nil {
+		return err
+	}
+	s.release()
 	return nil
 }
 
 // Append middleware to the scope
-func (s *scope) AppendMiddleware(m Middleware) {
-	s.middleware = append(s.middleware, m)
+func (s *scope) AppendMiddleware(m ...Middleware) {
+	if len(m) == 0 {
+		return
+	}
+	s.middleware = append(s.middleware, m...)
 }
 
 // Append global middleware to the global scope
 func AppendGlobalMiddleware(ms ...Middleware) {
-	for _, m := range ms {
-		globalScope.AppendMiddleware(m)
+	if len(ms) == 0 {
+		return
 	}
+	globalScope.AppendMiddleware(ms...)
 }
 
 // Dispose global scope to free up all resolved values and trigger scope end middlewares
 func DisposeGlobalScope() error {
 	return globalScope.Dispose()
 }
 
+func DisposeGlobalScopeWithContext(ctx context.Context) error {
+	return globalScope.DisposeWithContext(ctx)
+}
+
 // Apply middleware to the global scope
 func Apply(s Middleware) {
 	globalScope.AppendMiddleware(s)
@@ -194,7 +246,8 @@ type Middleware struct {
 	onScopeResolveType reflect.Type
 	onScopeResolve     reflect.Value
 
-	onScopeEnd func() error
+	onScopeEnd            func() error
+	onScopeEndWithContext func(context.Context) error
 }
 
 type MiddlewareFn func(Middleware) Middleware
@@ -214,6 +267,13 @@ func WithScopeEnd(fn func() error) Middleware {
 	}
 }
 
+func WithContextScopeEnd(fn func(context.Context) error) Middleware {
+	return Middleware{
+		hasOnScopeEnd:         true,
+		onScopeEndWithContext: fn,
+	}
+}
+
 type ScopeOpts struct {
 	inherit     bool
 	parent      Scope

scope_test.go

@@ -1,31 +1,31 @@
 package submodule_test
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/submodule-org/submodule.go/v2"
 )
 
-type k struct {
-	v int
-}
-
 func TestScope(t *testing.T) {
 	var seed int
+	var isDisposeGlobal bool
+	var isDisposeLocal bool
 	intValue := submodule.Make[int](func(self submodule.Self) int {
-		self.Scope.AppendMiddleware(submodule.WithScopeEnd(func() error {
-			fmt.Println("add 2 at the end of the day")
-			seed = seed + 2
-			return nil
-		}))
-
-		self.Scope.AppendMiddleware(submodule.WithScopeResolve(func(i any) any {
-			fmt.Println("catch everything")
-			return i
-		}))
-
+		self.Scope.AppendMiddleware(
+			submodule.WithScopeEnd(func() error {
+				fmt.Println("add 2 at the end of the day")
+				seed = seed + 2
+				isDisposeLocal = true
+				return nil
+			}),
+			submodule.WithScopeResolve(func(i any) any {
+				fmt.Println("catch everything")
+				return i
+			}),
+		)
 		return seed
 	})
 
@@ -36,6 +36,7 @@ func TestScope(t *testing.T) {
 
 	onEnd := submodule.WithScopeEnd(func() error {
 		fmt.Println("scope is ended")
+		isDisposeGlobal = true
 		return nil
 	})
 
@@ -50,4 +51,29 @@ func TestScope(t *testing.T) {
 	assert.Nil(t, e)
 
 	assert.Equal(t, 2, seed)
+	assert.True(t, isDisposeGlobal)
+	assert.True(t, isDisposeLocal)
+}
+
+func Test_Scope_With_Context(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	var isDispose bool
+	sub := submodule.Make[int](func(self submodule.Self) int {
+		self.Scope.AppendMiddleware(
+			submodule.WithContextScopeEnd(func(ctx context.Context) error {
+				<-ctx.Done()
+				isDispose = true
+				return nil
+			}),
+		)
+		return 0
+	})
+	scope := submodule.CreateScope()
+	_ = sub.ResolveWith(scope)
+	go func() {
+		cancel()
+	}()
+	err := scope.DisposeWithContext(ctx)
+	assert.Nil(t, err)
+	assert.True(t, isDispose)
 }