|
38 | 38 | import java.util.Queue; |
39 | 39 | import java.util.concurrent.CancellationException; |
40 | 40 | import java.util.concurrent.ConcurrentHashMap; |
| 41 | +import java.util.concurrent.CountDownLatch; |
41 | 42 | import java.util.concurrent.ExecutionException; |
42 | 43 | import java.util.concurrent.ExecutorService; |
43 | 44 | import java.util.concurrent.Future; |
@@ -140,14 +141,19 @@ private void execute(final JobConfiguration jobConfig, final ShardingContexts sh |
140 | 141 | private void process(final JobConfiguration jobConfig, final ShardingContexts shardingContexts, final ExecutionSource executionSource) { |
141 | 142 | Collection<Integer> items = shardingContexts.getShardingItemParameters().keySet(); |
142 | 143 | Queue<Future<Boolean>> futures = new LinkedList<>(); |
| 144 | + CountDownLatch latch = new CountDownLatch(items.size()); |
143 | 145 | for (int each : items) { |
144 | 146 | JobExecutionEvent jobExecutionEvent = new JobExecutionEvent(IpUtils.getHostName(), IpUtils.getIp(), shardingContexts.getTaskId(), jobConfig.getJobName(), executionSource, each); |
145 | 147 | ExecutorService executorService = executorContext.get(ExecutorService.class); |
146 | 148 | if (executorService.isShutdown()) { |
147 | 149 | return; |
148 | 150 | } |
149 | 151 | Future<Boolean> future = executorService.submit(() -> { |
150 | | - process(jobConfig, shardingContexts, each, jobExecutionEvent); |
| 152 | + try { |
| 153 | + process(jobConfig, shardingContexts, each, jobExecutionEvent); |
| 154 | + } finally { |
| 155 | + latch.countDown(); |
| 156 | + } |
151 | 157 | return true; |
152 | 158 | }); |
153 | 159 | futures.offer(future); |
@@ -178,7 +184,12 @@ private void process(final JobConfiguration jobConfig, final ShardingContexts sh |
178 | 184 | } finally { |
179 | 185 | futures.poll(); |
180 | 186 | } |
| 187 | + } |
181 | 188 |
|
| 189 | + try { |
| 190 | + latch.await(); |
| 191 | + } catch (final InterruptedException ex) { |
| 192 | + Thread.currentThread().interrupt(); |
182 | 193 | } |
183 | 194 | } |
184 | 195 |
|
|
0 commit comments