From 81fad9ba1af0a2b1b3cffedadc035ef1402c96b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=98=BF=E7=94=B7?= <l.weinan@gmail.com>
Date: Fri, 15 Dec 2023 00:24:23 +0800
Subject: [PATCH] JBERET-561 Change visibility of scopes implementations

---
 .../jberet/creation/JobScopedContextImpl.java | 32 +++++++++++++++----
 .../creation/PartitionScopedContextImpl.java  |  8 +++++
 .../creation/StepScopedContextImpl.java       |  8 +++++
 3 files changed, 42 insertions(+), 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
index 3309606c8..62234e5bb 100644
--- 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
index 91a32680c..cfcdf0deb 100644
--- 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
index 3d7ccb128..3a84a1c18 100644
--- 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();