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