fix(dependency-injection): honor custom scopes in toOverriding, bind(value).to(Class), and close()

dependency-injection/src/main/java/build/codemodel/injection/CustomScopedClassBinding.java

@@ -21,7 +21,12 @@
  */
 
 import java.lang.annotation.Annotation;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Stream;
 
 /**
  * A {@link ClassBinding} whose instance lifecycle is managed by a custom {@link Scope} (i.e. any
@@ -59,6 +64,12 @@ class CustomScopedClassBinding<T>
      */
     private final ValueBinding<T> delegate;
 
+    /**
+     * All distinct instances produced by this binding, tracked by identity for {@code @PreDestroy}.
+     */
+    private final Set<T> instances = Collections.synchronizedSet(
+        Collections.newSetFromMap(new IdentityHashMap<>()));
+
     /**
      * Constructs a {@link CustomScopedClassBinding}.
      *
@@ -103,6 +114,18 @@ Class<? extends Annotation> scopeAnnotation() {
 
     @Override
     public T value() {
-        return this.delegate.value();
+        final var v = this.delegate.value();
+        this.instances.add(v);
+        return v;
+    }
+
+    boolean hasInstantiatedValues() {
+        return !this.instances.isEmpty();
+    }
+
+    Stream<T> instantiatedValues() {
+        synchronized (this.instances) {
+            return new ArrayList<>(this.instances).stream();
+        }
     }
 }

dependency-injection/src/main/java/build/codemodel/injection/InjectionContext.java

@@ -199,6 +199,14 @@ public Binding<T> toOverriding(final Class<? extends T> concreteClass) {
                 final var typeDescriptor = codeModel.getJDKTypeDescriptor(concreteClass)
                     .orElseThrow(() -> new IllegalArgumentException(
                         "Could not resolve a TypeDescriptor for " + concreteClass));
+                final var scopeEntry = this.injectionFramework.findScopeEntry(typeDescriptor);
+                if (scopeEntry.isPresent()) {
+                    final ValueBinding<T> factory = new SupplierBinding<>(
+                        dependency, () -> (T) InjectionContext.this.createUnscoped(concreteClass));
+                    final Binding<?> scoped = scopeEntry.get().getValue().scope(factory);
+                    return replaceBinding(dependency,
+                        new CustomScopedClassBinding<>(dependency, concreteClass, scopeEntry.get().getKey(), scoped));
+                }
                 return replaceBinding(dependency, this.injectionFramework.isSingleton(typeDescriptor)
                     ? new LazySingletonClassBinding<>(dependency, concreteClass)
                     : new NonSingletonClassBinding<T>(dependency, concreteClass));
@@ -303,6 +311,14 @@ public Binding<T> to(final Class<? extends T> concreteClass) {
                 final var typeDescriptor = codeModel.getJDKTypeDescriptor(concreteClass)
                     .orElseThrow(() -> new IllegalArgumentException(
                         "Could not resolve a TypeDescriptor for " + concreteClass));
+                final var scopeEntry = this.injectionFramework.findScopeEntry(typeDescriptor);
+                if (scopeEntry.isPresent()) {
+                    final ValueBinding<T> factory = new SupplierBinding<>(
+                        dependency, () -> (T) InjectionContext.this.createUnscoped(concreteClass));
+                    final Binding<?> scoped = scopeEntry.get().getValue().scope(factory);
+                    return addBinding(dependency,
+                        new CustomScopedClassBinding<>(dependency, concreteClass, scopeEntry.get().getKey(), scoped));
+                }
                 return addBinding(dependency, this.injectionFramework.isSingleton(typeDescriptor)
                     ? new LazySingletonClassBinding<>(dependency, concreteClass)
                     : new NonSingletonClassBinding<>(dependency, concreteClass));
@@ -624,30 +640,40 @@ public Context initializeEagerSingletons() {
     @Override
     @SuppressWarnings("unchecked")
     public void close() {
-        // Collect all instantiated singleton bindings
+        // Collect all instantiated singleton and custom-scoped bindings
         final var instantiatedSingletons = this.bindingsByDependency.values().stream()
             .filter(LazySingletonClassBinding.class::isInstance)
             .map(b -> (LazySingletonClassBinding<Object>) b)
             .filter(b -> b.value().optional().isPresent())
             .toList();
 
-        if (instantiatedSingletons.isEmpty()) {
+        final var instantiatedCustomScoped = this.bindingsByDependency.values().stream()
+            .filter(CustomScopedClassBinding.class::isInstance)
+            .map(b -> (CustomScopedClassBinding<Object>) b)
+            .filter(CustomScopedClassBinding::hasInstantiatedValues)
+            .toList();
+
+        if (instantiatedSingletons.isEmpty() && instantiatedCustomScoped.isEmpty()) {
             return;
         }
 
-        // Build a dependency graph over instantiated singletons only
-        final var singletonDeps = instantiatedSingletons.stream()
-            .map(Binding::dependency)
+        // Build a dependency graph over all instantiated bindings
+        final var instantiatedDeps = Stream.concat(
+                instantiatedSingletons.stream().map(Binding::dependency),
+                instantiatedCustomScoped.stream().map(Binding::dependency))
             .collect(Collectors.toSet());
 
         final var graphBuilder = Graph.<Dependency>directed();
-        instantiatedSingletons.forEach(b -> {
+        Stream.concat(
+            instantiatedSingletons.stream().map(b -> (ClassBinding<Object>) b),
+            instantiatedCustomScoped.stream().map(b -> (ClassBinding<Object>) b)
+        ).forEach(b -> {
             final var bindingDep = b.dependency();
             graphBuilder.addVertex(bindingDep);
             this.injectionFramework.getInjectableDescriptor(b.concreteClass())
                 .injectionPoints()
                 .flatMap(InjectionPoint::dependencies)
-                .filter(singletonDeps::contains)
+                .filter(instantiatedDeps::contains)
                 .forEach(dep -> graphBuilder.addEdge(bindingDep, dep));
         });
 
@@ -658,33 +684,34 @@ public void close() {
         Collections.reverse(destroyOrder);
 
         // Index bindings by dependency for lookup during destruction
-        final var depToBinding = instantiatedSingletons.stream()
+        final var depToBinding = Stream.concat(
+                instantiatedSingletons.stream().map(b -> (ClassBinding<Object>) b),
+                instantiatedCustomScoped.stream().map(b -> (ClassBinding<Object>) b))
             .collect(Collectors.toMap(Binding::dependency, b -> b));
 
-        // Invoke @PreDestroy methods on each singleton in destruction order
+        // Invoke @PreDestroy methods in destruction order
         destroyOrder.forEach(dep -> {
             final var binding = depToBinding.get(dep);
-            if (binding == null) {
-                return;
-            }
-            final var instance = binding.value().optional().orElse(null);
-            if (instance == null) {
-                return;
-            }
-            this.injectionFramework.getInjectableDescriptor(binding.concreteClass())
-                .preDestroyMethods()
-                .map(md -> md.getTrait(MethodType.class).orElse(null))
-                .filter(Objects::nonNull)
-                .map(MethodType::method)
-                .forEach(method -> {
-                    try {
-                        method.setAccessible(true);
-                        method.invoke(instance);
-                    } catch (final IllegalAccessException | InvocationTargetException e) {
-                        throw new InjectionException(
-                            "Invoking @PreDestroy method " + method + " on " + instance.getClass(), e);
-                    }
-                });
+            final Stream<Object> instances = switch (binding) {
+                case LazySingletonClassBinding<Object> lsb -> lsb.value().optional().stream();
+                case CustomScopedClassBinding<Object> csb -> csb.instantiatedValues();
+                default -> throw new IllegalStateException("Unexpected binding type: " + binding.getClass());
+            };
+            instances.forEach(instance ->
+                this.injectionFramework.getInjectableDescriptor(binding.concreteClass())
+                    .preDestroyMethods()
+                    .filter(md -> md.hasTrait(MethodType.class))
+                    .map(md -> md.trait(MethodType.class))
+                    .map(MethodType::method)
+                    .forEach(method -> {
+                        try {
+                            method.setAccessible(true);
+                            method.invoke(instance);
+                        } catch (final IllegalAccessException | InvocationTargetException e) {
+                            throw new InjectionException(
+                                "Invoking @PreDestroy method " + method + " on " + instance.getClass(), e);
+                        }
+                    }));
         });
     }
 

dependency-injection/src/test/java/build/codemodel/injection/LifecycleTests.java

@@ -85,7 +85,7 @@ static class NoDestroyService {
     static class ContextScopedService {
     }
 
-    // ---- custom scope fixture ----
+    // ---- custom scope fixtures ----
 
     @Retention(RetentionPolicy.RUNTIME)
     @Target(ElementType.TYPE)
@@ -97,6 +97,14 @@ static class ContextScopedService {
     static class CustomScopedService {
     }
 
+    @AlwaysSame
+    static class CustomScopedWithPreDestroy {
+        @PreDestroy
+        void destroy() {
+            events.add("custom-scoped");
+        }
+    }
+
     // ---- @PreDestroy: reverse topological order ----
 
     /**
@@ -228,4 +236,57 @@ void shouldSupportCustomScopeViaBindScope() {
         assertThat(context.create(CustomScopedService.class)).isSameAs(fixedInstance);
         assertThat(context.create(CustomScopedService.class)).isSameAs(fixedInstance);
     }
+
+    /**
+     * Ensures {@link BindingBuilder#toOverriding(Class)} honours a custom scope annotation on the
+     * concrete class instead of silently falling back to prototype.
+     */
+    @Test
+    void shouldRespectCustomScopeInToOverriding() {
+        final var framework = createInjectionFramework();
+        final var fixedInstance = new CustomScopedService();
+        framework.bindScope(AlwaysSame.class, binding -> ValueBinding.of(binding.dependency(), fixedInstance));
+
+        final var context = framework.newContext();
+        context.bind(CustomScopedService.class).to(CustomScopedService.class);
+        context.bind(CustomScopedService.class).toOverriding(CustomScopedService.class);
+
+        assertThat(context.create(CustomScopedService.class)).isSameAs(fixedInstance);
+    }
+
+    /**
+     * Ensures {@link Context#bind(Object)} followed by {@link BindingBuilder#to(Class)} honours
+     * a custom scope annotation on the concrete class instead of silently falling back to prototype.
+     */
+    @Test
+    void shouldRespectCustomScopeInBindValueToClass() {
+        final var framework = createInjectionFramework();
+        final var fixedInstance = new CustomScopedService();
+        framework.bindScope(AlwaysSame.class, binding -> ValueBinding.of(binding.dependency(), fixedInstance));
+
+        final var context = framework.newContext();
+        final var dummy = new CustomScopedService();
+        context.bind(dummy).to(CustomScopedService.class);
+
+        assertThat(context.create(CustomScopedService.class)).isSameAs(fixedInstance);
+    }
+
+    /**
+     * Ensures {@link Context#close()} invokes {@link PreDestroy} on instantiated custom-scoped instances.
+     */
+    @Test
+    void shouldInvokePreDestroyOnCustomScopedInstances() {
+        events.clear();
+        final var framework = createInjectionFramework();
+        final var fixedInstance = new CustomScopedWithPreDestroy();
+        framework.bindScope(AlwaysSame.class, binding -> ValueBinding.of(binding.dependency(), fixedInstance));
+
+        final var context = framework.newContext();
+        context.bind(CustomScopedWithPreDestroy.class).to(CustomScopedWithPreDestroy.class);
+        context.create(CustomScopedWithPreDestroy.class);
+
+        context.close();
+
+        assertThat(events).containsExactly("custom-scoped");
+    }
 }