1
1
package redis .clients .jedis ;
2
2
3
- import java .util .Iterator ;
4
3
import java .util .LinkedHashMap ;
5
4
import java .util .LinkedList ;
6
5
import java .util .List ;
7
6
import java .util .Map ;
8
7
import java .util .Queue ;
9
- import java .util .concurrent .CountDownLatch ;
8
+ import java .util .concurrent .CompletableFuture ;
9
+ import java .util .concurrent .ExecutionException ;
10
10
import java .util .concurrent .ExecutorService ;
11
11
import java .util .concurrent .Executors ;
12
12
@@ -31,6 +31,7 @@ public abstract class MultiNodePipelineBase extends PipelineBase {
31
31
32
32
private final Map <HostAndPort , Queue <Response <?>>> pipelinedResponses ;
33
33
private final Map <HostAndPort , Connection > connections ;
34
+ private ExecutorService executorService ;
34
35
private volatile boolean syncing = false ;
35
36
36
37
public MultiNodePipelineBase (CommandObjects commandObjects ) {
@@ -39,6 +40,13 @@ public MultiNodePipelineBase(CommandObjects commandObjects) {
39
40
connections = new LinkedHashMap <>();
40
41
}
41
42
43
+ public MultiNodePipelineBase (CommandObjects commandObjects , ExecutorService executorService ) {
44
+ super (commandObjects );
45
+ this .executorService = executorService ;
46
+ pipelinedResponses = new LinkedHashMap <>();
47
+ connections = new LinkedHashMap <>();
48
+ }
49
+
42
50
/**
43
51
* Sub-classes must call this method, if graph commands are going to be used.
44
52
* @param connectionProvider connection provider
@@ -96,44 +104,47 @@ public final void sync() {
96
104
return ;
97
105
}
98
106
syncing = true ;
99
-
100
- ExecutorService executorService = Executors .newFixedThreadPool (MULTI_NODE_PIPELINE_SYNC_WORKERS );
101
-
102
- CountDownLatch countDownLatch = new CountDownLatch (pipelinedResponses .size ());
103
- Iterator <Map .Entry <HostAndPort , Queue <Response <?>>>> pipelinedResponsesIterator
104
- = pipelinedResponses .entrySet ().iterator ();
105
- while (pipelinedResponsesIterator .hasNext ()) {
106
- Map .Entry <HostAndPort , Queue <Response <?>>> entry = pipelinedResponsesIterator .next ();
107
- HostAndPort nodeKey = entry .getKey ();
108
- Queue <Response <?>> queue = entry .getValue ();
109
- Connection connection = connections .get (nodeKey );
110
- executorService .submit (() -> {
111
- try {
112
- List <Object > unformatted = connection .getMany (queue .size ());
113
- for (Object o : unformatted ) {
114
- queue .poll ().set (o );
115
- }
116
- } catch (JedisConnectionException jce ) {
117
- log .error ("Error with connection to " + nodeKey , jce );
118
- // cleanup the connection
119
- pipelinedResponsesIterator .remove ();
120
- connections .remove (nodeKey );
121
- IOUtils .closeQuietly (connection );
122
- } finally {
123
- countDownLatch .countDown ();
124
- }
125
- });
126
- }
127
-
107
+ ExecutorService executorService = getExecutorService ();
108
+ CompletableFuture [] futures
109
+ = pipelinedResponses .entrySet ().stream ()
110
+ .map (e -> CompletableFuture .runAsync (() -> closeConnection (e ), executorService ))
111
+ .toArray (CompletableFuture []::new );
112
+ CompletableFuture awaitAllCompleted = CompletableFuture .allOf (futures );
128
113
try {
129
- countDownLatch .await ();
114
+ awaitAllCompleted .get ();
115
+ if (executorService != this .executorService ) {
116
+ executorService .shutdown ();
117
+ }
118
+ } catch (ExecutionException e ) {
119
+ log .error ("Failed execution." , e );
130
120
} catch (InterruptedException e ) {
131
121
log .error ("Thread is interrupted during sync." , e );
122
+ Thread .currentThread ().interrupt ();
132
123
}
124
+ syncing = false ;
125
+ }
133
126
134
- executorService .shutdownNow ();
127
+ private ExecutorService getExecutorService () {
128
+ if (executorService == null ) {
129
+ return Executors .newFixedThreadPool (Math .min (this .pipelinedResponses .size (), MULTI_NODE_PIPELINE_SYNC_WORKERS ));
130
+ }
131
+ return executorService ;
132
+ }
135
133
136
- syncing = false ;
134
+ private void closeConnection (Map .Entry <HostAndPort , Queue <Response <?>>> entry ) {
135
+ HostAndPort nodeKey = entry .getKey ();
136
+ Queue <Response <?>> queue = entry .getValue ();
137
+ Connection connection = connections .get (nodeKey );
138
+ try {
139
+ List <Object > unformatted = connection .getMany (queue .size ());
140
+ for (Object o : unformatted ) {
141
+ queue .poll ().set (o );
142
+ }
143
+ } catch (JedisConnectionException jce ) {
144
+ log .error ("Error with connection to " + nodeKey , jce );
145
+ connections .remove (nodeKey );
146
+ IOUtils .closeQuietly (connection );
147
+ }
137
148
}
138
149
139
150
@ Deprecated
0 commit comments