|
19 | 19 |
|
20 | 20 | package org.apache.druid.msq.dart.worker; |
21 | 21 |
|
22 | | -import com.fasterxml.jackson.databind.ObjectMapper; |
23 | | -import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes; |
24 | 22 | import com.google.common.util.concurrent.ListenableFuture; |
25 | | -import com.google.errorprone.annotations.concurrent.GuardedBy; |
26 | | -import it.unimi.dsi.fastutil.Pair; |
27 | | -import org.apache.druid.error.DruidException; |
28 | | -import org.apache.druid.java.util.common.logger.Logger; |
29 | | -import org.apache.druid.java.util.http.client.response.HttpResponseHandler; |
30 | 23 | import org.apache.druid.msq.dart.controller.DartWorkerManager; |
31 | | -import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; |
32 | | -import org.apache.druid.msq.dart.worker.http.DartWorkerResource; |
33 | 24 | import org.apache.druid.msq.exec.WorkerClient; |
34 | | -import org.apache.druid.msq.rpc.BaseWorkerClientImpl; |
35 | | -import org.apache.druid.rpc.FixedServiceLocator; |
36 | | -import org.apache.druid.rpc.IgnoreHttpResponseHandler; |
37 | | -import org.apache.druid.rpc.RequestBuilder; |
38 | | -import org.apache.druid.rpc.ServiceClient; |
39 | | -import org.apache.druid.rpc.ServiceClientFactory; |
40 | | -import org.apache.druid.rpc.ServiceLocation; |
41 | | -import org.apache.druid.rpc.ServiceRetryPolicy; |
42 | | -import org.apache.druid.utils.CloseableUtils; |
43 | | -import org.jboss.netty.handler.codec.http.HttpMethod; |
44 | 25 |
|
45 | | -import javax.annotation.Nullable; |
46 | | -import java.io.Closeable; |
47 | | -import java.net.URI; |
48 | | -import java.util.HashMap; |
49 | | -import java.util.Map; |
50 | | - |
51 | | -/** |
52 | | - * Dart implementation of {@link WorkerClient}. Uses the same {@link BaseWorkerClientImpl} as the task-based engine. |
53 | | - * Each instance of this class is scoped to a single query. |
54 | | - */ |
55 | | -public class DartWorkerClient extends BaseWorkerClientImpl |
| 26 | +public interface DartWorkerClient extends WorkerClient |
56 | 27 | { |
57 | | - private static final Logger log = new Logger(DartWorkerClient.class); |
58 | | - |
59 | | - private final String queryId; |
60 | | - private final ServiceClientFactory clientFactory; |
61 | | - private final ServiceRetryPolicy retryPolicy; |
62 | | - |
63 | | - @Nullable |
64 | | - private final String controllerHost; |
65 | | - |
66 | | - @GuardedBy("clientMap") |
67 | | - private final Map<String, Pair<ServiceClient, Closeable>> clientMap = new HashMap<>(); |
68 | | - |
69 | | - /** |
70 | | - * Create a worker client. |
71 | | - * |
72 | | - * @param queryId dart query ID. see {@link DartSqlEngine#CTX_DART_QUERY_ID} |
73 | | - * @param clientFactory service client factor |
74 | | - * @param smileMapper Smile object mapper |
75 | | - * @param controllerHost Controller host (see {@link DartWorkerResource#HEADER_CONTROLLER_HOST}) if this is a |
76 | | - * controller-to-worker client. Null if this is a worker-to-worker client. |
77 | | - */ |
78 | | - public DartWorkerClient( |
79 | | - final String queryId, |
80 | | - final ServiceClientFactory clientFactory, |
81 | | - final ObjectMapper smileMapper, |
82 | | - @Nullable final String controllerHost |
83 | | - ) |
84 | | - { |
85 | | - super(smileMapper, SmileMediaTypes.APPLICATION_JACKSON_SMILE); |
86 | | - this.queryId = queryId; |
87 | | - this.clientFactory = clientFactory; |
88 | | - this.controllerHost = controllerHost; |
89 | | - |
90 | | - if (controllerHost == null) { |
91 | | - // worker -> worker client. Retry HTTP 503 in case worker A starts up before worker B, and needs to |
92 | | - // contact it immediately. |
93 | | - this.retryPolicy = new DartWorkerRetryPolicy(true); |
94 | | - } else { |
95 | | - // controller -> worker client. Do not retry any HTTP error codes. If we retry HTTP 503 for controller -> worker, |
96 | | - // we can get stuck trying to contact workers that have exited. |
97 | | - this.retryPolicy = new DartWorkerRetryPolicy(false); |
98 | | - } |
99 | | - } |
100 | | - |
101 | | - @Override |
102 | | - protected ServiceClient getClient(final String workerIdString) |
103 | | - { |
104 | | - final WorkerId workerId = WorkerId.fromString(workerIdString); |
105 | | - if (!queryId.equals(workerId.getQueryId())) { |
106 | | - throw DruidException.defensive("Unexpected queryId[%s]. Expected queryId[%s]", workerId.getQueryId(), queryId); |
107 | | - } |
108 | | - |
109 | | - synchronized (clientMap) { |
110 | | - return clientMap.computeIfAbsent(workerId.getHostAndPort(), ignored -> makeNewClient(workerId)).left(); |
111 | | - } |
112 | | - } |
113 | | - |
114 | 28 | /** |
115 | 29 | * Close a single worker's clients. Used when that worker fails, so we stop trying to contact it. |
116 | 30 | * |
117 | 31 | * @param workerHost worker host:port |
118 | 32 | */ |
119 | | - public void closeClient(final String workerHost) |
120 | | - { |
121 | | - synchronized (clientMap) { |
122 | | - final Pair<ServiceClient, Closeable> clientPair = clientMap.remove(workerHost); |
123 | | - if (clientPair != null) { |
124 | | - CloseableUtils.closeAndWrapExceptions(clientPair.right()); |
125 | | - } |
126 | | - } |
127 | | - } |
128 | | - |
129 | | - /** |
130 | | - * Close all outstanding clients. |
131 | | - */ |
132 | | - @Override |
133 | | - public void close() |
134 | | - { |
135 | | - synchronized (clientMap) { |
136 | | - for (Map.Entry<String, Pair<ServiceClient, Closeable>> entry : clientMap.entrySet()) { |
137 | | - CloseableUtils.closeAndSuppressExceptions( |
138 | | - entry.getValue().right(), |
139 | | - e -> log.warn(e, "Failed to close client[%s]", entry.getKey()) |
140 | | - ); |
141 | | - } |
142 | | - |
143 | | - clientMap.clear(); |
144 | | - } |
145 | | - } |
| 33 | + void closeClient(String hostAndPort); |
146 | 34 |
|
147 | 35 | /** |
148 | 36 | * Stops a worker. Dart-only API, used by the {@link DartWorkerManager}. |
149 | 37 | */ |
150 | | - public ListenableFuture<?> stopWorker(String workerId) |
151 | | - { |
152 | | - return getClient(workerId).asyncRequest( |
153 | | - new RequestBuilder(HttpMethod.POST, "/stop"), |
154 | | - IgnoreHttpResponseHandler.INSTANCE |
155 | | - ); |
156 | | - } |
157 | | - |
158 | | - /** |
159 | | - * Create a new client. Called by {@link #getClient(String)} if a new one is needed. |
160 | | - */ |
161 | | - private Pair<ServiceClient, Closeable> makeNewClient(final WorkerId workerId) |
162 | | - { |
163 | | - final URI uri = workerId.toUri(); |
164 | | - final FixedServiceLocator locator = new FixedServiceLocator(ServiceLocation.fromUri(uri)); |
165 | | - final ServiceClient baseClient = |
166 | | - clientFactory.makeClient(workerId.toString(), locator, retryPolicy); |
167 | | - final ServiceClient client; |
168 | | - |
169 | | - if (controllerHost != null) { |
170 | | - client = new ControllerDecoratedClient(baseClient, controllerHost); |
171 | | - } else { |
172 | | - client = baseClient; |
173 | | - } |
174 | | - |
175 | | - return Pair.of(client, locator); |
176 | | - } |
177 | | - |
178 | | - /** |
179 | | - * Service client that adds the {@link DartWorkerResource#HEADER_CONTROLLER_HOST} header. |
180 | | - */ |
181 | | - private static class ControllerDecoratedClient implements ServiceClient |
182 | | - { |
183 | | - private final ServiceClient delegate; |
184 | | - private final String controllerHost; |
185 | | - |
186 | | - ControllerDecoratedClient(final ServiceClient delegate, final String controllerHost) |
187 | | - { |
188 | | - this.delegate = delegate; |
189 | | - this.controllerHost = controllerHost; |
190 | | - } |
191 | | - |
192 | | - @Override |
193 | | - public <IntermediateType, FinalType> ListenableFuture<FinalType> asyncRequest( |
194 | | - final RequestBuilder requestBuilder, |
195 | | - final HttpResponseHandler<IntermediateType, FinalType> handler |
196 | | - ) |
197 | | - { |
198 | | - return delegate.asyncRequest( |
199 | | - requestBuilder.header(DartWorkerResource.HEADER_CONTROLLER_HOST, controllerHost), |
200 | | - handler |
201 | | - ); |
202 | | - } |
203 | | - |
204 | | - @Override |
205 | | - public ServiceClient withRetryPolicy(final ServiceRetryPolicy retryPolicy) |
206 | | - { |
207 | | - return new ControllerDecoratedClient(delegate.withRetryPolicy(retryPolicy), controllerHost); |
208 | | - } |
209 | | - } |
| 38 | + ListenableFuture<?> stopWorker(String workerId); |
210 | 39 | } |
0 commit comments