View Javadoc

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 }