1 /*
2 * Copyright 2006-2007 the original author or authors.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 package org.springframework.batch.core.scope.context;
17
18 import java.util.HashMap;
19 import java.util.Map;
20 import java.util.Stack;
21 import java.util.concurrent.atomic.AtomicInteger;
22
23 import org.springframework.batch.core.Step;
24 import org.springframework.batch.core.StepExecution;
25
26 /**
27 * Central convenience class for framework use in managing the step scope
28 * context. Generally only to be used by implementations of {@link Step}. N.B.
29 * it is the responsibility of every {@link Step} implementation to ensure that
30 * a {@link StepContext} is available on every thread that might be involved in
31 * a step execution, including worker threads from a pool.
32 *
33 * @author Dave Syer
34 *
35 */
36 public class StepSynchronizationManager {
37
38 /*
39 * We have to deal with single and multi-threaded execution, with a single
40 * and with multiple step execution instances. That's 2x2 = 4 scenarios.
41 */
42
43 /**
44 * Storage for the current step execution; has to be ThreadLocal because it
45 * is needed to locate a StepContext in components that are not part of a
46 * Step (like when re-hydrating a scoped proxy). Doesn't use
47 * InheritableThreadLocal because there are side effects if a step is trying
48 * to run multiple child steps (e.g. with partitioning). The Stack is used
49 * to cover the single threaded case, so that the API is the same as
50 * multi-threaded.
51 */
52 private static final ThreadLocal<Stack<StepExecution>> executionHolder = new ThreadLocal<Stack<StepExecution>>();
53
54 /**
55 * Reference counter for each step execution: how many threads are using the
56 * same one?
57 */
58 private static final Map<StepExecution, AtomicInteger> counts = new HashMap<StepExecution, AtomicInteger>();
59
60 /**
61 * Simple map from a running step execution to the associated context.
62 */
63 private static final Map<StepExecution, StepContext> contexts = new HashMap<StepExecution, StepContext>();
64
65 /**
66 * Getter for the current context if there is one, otherwise returns null.
67 *
68 * @return the current {@link StepContext} or null if there is none (if one
69 * has not been registered for this thread).
70 */
71 public static StepContext getContext() {
72 if (getCurrent().isEmpty()) {
73 return null;
74 }
75 synchronized (contexts) {
76 return contexts.get(getCurrent().peek());
77 }
78 }
79
80 /**
81 * Register a context with the current thread - always put a matching
82 * {@link #close()} call in a finally block to ensure that the correct
83 * context is available in the enclosing block.
84 *
85 * @param stepExecution the step context to register
86 * @return a new {@link StepContext} or the current one if it has the same
87 * {@link StepExecution}
88 */
89 public static StepContext register(StepExecution stepExecution) {
90 if (stepExecution == null) {
91 return null;
92 }
93 getCurrent().push(stepExecution);
94 StepContext context;
95 synchronized (contexts) {
96 context = contexts.get(stepExecution);
97 if (context == null) {
98 context = new StepContext(stepExecution);
99 contexts.put(stepExecution, context);
100 }
101 }
102 increment();
103 return context;
104 }
105
106 /**
107 * Method for de-registering the current context - should always and only be
108 * used by in conjunction with a matching {@link #register(StepExecution)}
109 * to ensure that {@link #getContext()} always returns the correct value.
110 * Does not call {@link StepContext#close()} - that is left up to the caller
111 * because he has a reference to the context (having registered it) and only
112 * he has knowledge of when the step actually ended.
113 */
114 public static void close() {
115 StepContext oldSession = getContext();
116 if (oldSession == null) {
117 return;
118 }
119 decrement();
120 }
121
122 private static void decrement() {
123 StepExecution current = getCurrent().pop();
124 if (current != null) {
125 int remaining = counts.get(current).decrementAndGet();
126 if (remaining <= 0) {
127 synchronized (contexts) {
128 contexts.remove(current);
129 counts.remove(current);
130 }
131 }
132 }
133 }
134
135 private static void increment() {
136 StepExecution current = getCurrent().peek();
137 if (current != null) {
138 AtomicInteger count;
139 synchronized (counts) {
140 count = counts.get(current);
141 if (count == null) {
142 count = new AtomicInteger();
143 counts.put(current, count);
144 }
145 }
146 count.incrementAndGet();
147 }
148 }
149
150 private static Stack<StepExecution> getCurrent() {
151 if (executionHolder.get() == null) {
152 executionHolder.set(new Stack<StepExecution>());
153 }
154 return executionHolder.get();
155 }
156
157 /**
158 * A convenient "deep" close operation. Call this instead of
159 * {@link #close()} if the step execution for the current context is ending.
160 * Delegates to {@link StepContext#close()} and then ensures that
161 * {@link #close()} is also called in a finally block.
162 */
163 public static void release() {
164 StepContext context = getContext();
165 try {
166 if (context != null) {
167 context.close();
168 }
169 }
170 finally {
171 close();
172 }
173 }
174
175 }