JBERET-561 Change visibility of scopes implementations

jberet · Dec 16, 2023 · 81fad9b · 81fad9b
1 parent c6d30f8
commit 81fad9b
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 6 deletions.
diff --git a/jberet-core/src/main/java/org/jberet/creation/JobScopedContextImpl.java b/jberet-core/src/main/java/org/jberet/creation/JobScopedContextImpl.java
@@ -26,6 +26,10 @@ public class JobScopedContextImpl implements Context {
     private JobScopedContextImpl() {
     }
 
+    public static JobScopedContextImpl getInstance() {
+        return INSTANCE;
+    }
+
     @Override
     public Class<? extends Annotation> getScope() {
         return JobScoped.class;
@@ -61,6 +65,10 @@ public boolean isActive() {
         return ArtifactCreationContext.getCurrentArtifactCreationContext().jobContext != null;
     }
 
+    public void destroy(Contextual<?> contextual) {
+        JobScopedContextImpl.ScopedInstance.destroy(getJobScopedBeans(), contextual);
+    }
+
     private ConcurrentMap<Contextual<?>, ScopedInstance<?>> getJobScopedBeans() {
         final JobContextImpl jobContext = ArtifactCreationContext.getCurrentArtifactCreationContext().jobContext;
         return jobContext.getScopedBeans();
@@ -75,15 +83,27 @@ public ScopedInstance(final T instance, final CreationalContext<T> creationalCon
             this.creationalContext = creationalContext;
         }
 
-        @SuppressWarnings("unchecked")
         public static void destroy(final ConcurrentMap<Contextual<?>, JobScopedContextImpl.ScopedInstance<?>> scopedBeans) {
-            if (scopedBeans.size() > 0) {
-                for (final Map.Entry<Contextual<?>, JobScopedContextImpl.ScopedInstance<?>> e : scopedBeans.entrySet()) {
-                    final Contextual<?> contextual = e.getKey();
-                    final ScopedInstance<?> value = e.getValue();
+            destroy(scopedBeans, null);
+        }
+
+        @SuppressWarnings({"rawtypes", "unchecked"})
+        public static void destroy(final ConcurrentMap<Contextual<?>, JobScopedContextImpl.ScopedInstance<?>> scopedBeans,
+                                   final Contextual<?> contextual) {
+            if (contextual == null) {
+                if (scopedBeans.size() > 0) {
+                    for (final Map.Entry<Contextual<?>, JobScopedContextImpl.ScopedInstance<?>> e : scopedBeans.entrySet()) {
+                        final Contextual<?> key = e.getKey();
+                        final ScopedInstance<?> value = e.getValue();
+                        ((Contextual) key).destroy(value.instance, value.creationalContext);
+                    }
+                    scopedBeans.clear();
+                }
+            } else {
+                final ScopedInstance<?> value = scopedBeans.remove(contextual);
+                if (value != null) {
                     ((Contextual) contextual).destroy(value.instance, value.creationalContext);
                 }
-                scopedBeans.clear();
             }
         }
     }

diff --git a/jberet-core/src/main/java/org/jberet/creation/PartitionScopedContextImpl.java b/jberet-core/src/main/java/org/jberet/creation/PartitionScopedContextImpl.java
@@ -25,6 +25,10 @@ public class PartitionScopedContextImpl implements Context {
     private PartitionScopedContextImpl() {
     }
 
+    public static PartitionScopedContextImpl getInstance() {
+        return INSTANCE;
+    }
+
     @Override
     public Class<? extends Annotation> getScope() {
         return PartitionScoped.class;
@@ -61,6 +65,10 @@ public boolean isActive() {
         return stepContext != null && stepContext.getPartitionScopedBeans() != null;
     }
 
+    public void destroy(Contextual<?> contextual) {
+        JobScopedContextImpl.ScopedInstance.destroy(getPartitionScopedBeans(), contextual);
+    }
+
     private ConcurrentMap<Contextual<?>, JobScopedContextImpl.ScopedInstance<?>> getPartitionScopedBeans() {
         final StepContextImpl stepContext = ArtifactCreationContext.getCurrentArtifactCreationContext().stepContext;
         return stepContext.getPartitionScopedBeans();

diff --git a/jberet-core/src/main/java/org/jberet/creation/StepScopedContextImpl.java b/jberet-core/src/main/java/org/jberet/creation/StepScopedContextImpl.java
@@ -25,6 +25,10 @@ public class StepScopedContextImpl implements Context {
     private StepScopedContextImpl() {
     }
 
+    public static StepScopedContextImpl getInstance() {
+        return INSTANCE;
+    }
+
     @Override
     public Class<? extends Annotation> getScope() {
         return StepScoped.class;
@@ -60,6 +64,10 @@ public boolean isActive() {
         return ArtifactCreationContext.getCurrentArtifactCreationContext().stepContext != null;
     }
 
+    public void destroy(Contextual<?> contextual) {
+        JobScopedContextImpl.ScopedInstance.destroy(getStepScopedBeans(), contextual);
+    }
+
     private ConcurrentMap<Contextual<?>, JobScopedContextImpl.ScopedInstance<?>> getStepScopedBeans() {
         final StepContextImpl stepContext = ArtifactCreationContext.getCurrentArtifactCreationContext().stepContext;
         return stepContext.getScopedBeans();