diff --git a/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsTableBenchmark.java b/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsTableBenchmark.java index d75f46c89eef..3b85e24f9d1a 100644 --- a/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsTableBenchmark.java +++ b/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsTableBenchmark.java @@ -35,12 +35,12 @@ import org.apache.druid.client.InternalQueryConfig; import org.apache.druid.client.TimelineServerView; import org.apache.druid.client.coordinator.NoopCoordinatorClient; +import org.apache.druid.collections.ResourceHolder; +import org.apache.druid.collections.StupidResourceHolder; import org.apache.druid.error.NotYetImplemented; import org.apache.druid.jackson.DefaultObjectMapper; -import org.apache.druid.java.util.common.CloseableIterators; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.java.util.common.parsers.CloseableIterator; import org.apache.druid.segment.join.JoinableFactory; import org.apache.druid.segment.metadata.CentralizedDatasourceSchemaConfig; import org.apache.druid.server.QueryLifecycleFactory; @@ -73,6 +73,7 @@ import java.util.ArrayList; import java.util.Comparator; import java.util.EnumMap; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; @@ -180,12 +181,12 @@ public void setup() final NoopCoordinatorClient coordinatorClient = new NoopCoordinatorClient() { @Override - public ListenableFuture> fetchAllUsedSegmentsWithOvershadowedStatus( + public ListenableFuture>> fetchAllUsedSegmentsWithOvershadowedStatus( Set watchedDataSources, boolean includeRealtimeSegments ) { - return Futures.immediateFuture(CloseableIterators.withEmptyBaggage(publishedSegments.iterator())); + return Futures.immediateFuture(StupidResourceHolder.create(publishedSegments.iterator())); } }; diff --git a/docs/configuration/index.md b/docs/configuration/index.md index f46f919e4546..f550a4dc4307 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -706,6 +706,7 @@ All Druid components can communicate with each other over HTTP. |`druid.global.http.connectTimeout`|Connect timeout for the HTTP client used for most direct RPC between Druid services. This covers, among other things, Overlord-to-task and supervisor-to-task calls in the indexing service, Coordinator lookup management, dynamic config sync between services, MSQ tasks reading from data servers, and general Coordinator/Overlord/Broker service clients. Does not affect Broker-to-Historical query dispatch (see `druid.broker.http.connectTimeout`) or request forwarding (see `clientConnectTimeout`).|`PT10S`| |`druid.global.http.allocator`|Netty memory allocator used by the direct-RPC HTTP client. Accepts `adaptive` (adaptive between `pooled` and `unpooled` based on load), `pooled`, or `unpooled`.|`adaptive`| |`druid.global.http.poolImplementation`|How the connection pool tracks demand, never exceeding `numConnections` either way. With `adaptive`, a request discards every stale or broken connection it walks past and opens a new one only once none is left, so the pool falls back to the number of connections the traffic actually needs. With `retaining`, the pool holds on to every connection it has opened, replacing a stale or broken one by a fresh one, one for one, so it stays at its high-water mark.|`adaptive`| +|`druid.client.coordinator.maxAttempts`|Maximum number of attempts, including the first one, that a service makes for a request to the Coordinator before giving up on a retryable error. Must be at least 1.|`15`| ### Common endpoints configuration diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/CoordinatorClientTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/CoordinatorClientTest.java index 065963f25a37..980b2b256abf 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/CoordinatorClientTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/CoordinatorClientTest.java @@ -21,11 +21,11 @@ import org.apache.druid.client.ImmutableSegmentLoadInfo; import org.apache.druid.client.coordinator.CoordinatorClient; +import org.apache.druid.collections.ResourceHolder; import org.apache.druid.common.utils.IdUtils; import org.apache.druid.indexing.common.task.Task; import org.apache.druid.indexing.common.task.TaskBuilder; import org.apache.druid.java.util.common.Intervals; -import org.apache.druid.java.util.common.parsers.CloseableIterator; import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.server.coordinator.rules.ForeverBroadcastDistributionRule; import org.apache.druid.server.coordinator.rules.Rule; @@ -44,9 +44,9 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; -import java.io.IOException; import java.net.URI; import java.util.ArrayList; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; @@ -154,13 +154,14 @@ public void test_fetchUsedSegments() @Test @Timeout(20) - public void test_fetchAllUsedSegmentsWithOvershadowedStatus() throws IOException + public void test_fetchAllUsedSegmentsWithOvershadowedStatus() { runIndexTask(); - try (CloseableIterator iterator = cluster.callApi().onLeaderCoordinator( + try (ResourceHolder> segments = cluster.callApi().onLeaderCoordinator( c -> c.fetchAllUsedSegmentsWithOvershadowedStatus(Set.of(dataSource), true)) ) { + final Iterator iterator = segments.get(); Assertions.assertTrue(iterator.hasNext()); SegmentStatusInCluster segmentStatus = iterator.next(); Assertions.assertEquals(dataSource, segmentStatus.getDataSegment().getDataSource()); diff --git a/processing/src/main/java/org/apache/druid/java/util/http/client/NettyHttpClient.java b/processing/src/main/java/org/apache/druid/java/util/http/client/NettyHttpClient.java index ff4f64af3947..7455eb487b5c 100644 --- a/processing/src/main/java/org/apache/druid/java/util/http/client/NettyHttpClient.java +++ b/processing/src/main/java/org/apache/druid/java/util/http/client/NettyHttpClient.java @@ -240,6 +240,7 @@ public ListenableFuture go( private final Object watermarkLock = new Object(); private long suspendWatermark = -1; private long resumeWatermark = -1; + private boolean returnedToPool = false; @Override protected void channelRead0(ChannelHandlerContext ctx, HttpObject msg) @@ -274,6 +275,9 @@ protected void channelRead0(ChannelHandlerContext ctx, HttpObject msg) public long resume(long resumeChunkNum) { synchronized (watermarkLock) { + if (returnedToPool) { + return 0; + } resumeWatermark = Math.max(resumeWatermark, resumeChunkNum); if (suspendWatermark >= 0 && resumeWatermark >= suspendWatermark) { @@ -292,8 +296,12 @@ public long resume(long resumeChunkNum) @Override public void abort() { - log.debug("[%s] Aborted connection at caller's request.", requestDesc); - channel.close(); + // On the event loop, so that it is ordered against finishRequest handing the channel back. + if (channel.eventLoop().inEventLoop()) { + closeUnlessReturnedToPool(); + } else { + channel.eventLoop().execute(() -> closeUnlessReturnedToPool()); + } } }; @@ -388,9 +396,22 @@ private void finishRequest() } removeHandlers(); channel.config().setAutoRead(true); + synchronized (watermarkLock) { + returnedToPool = true; + } channelResourceContainer.returnResource(); } + private void closeUnlessReturnedToPool() + { + synchronized (watermarkLock) { + if (!returnedToPool) { + log.debug("[%s] Aborted connection at caller's request.", requestDesc); + channel.close(); + } + } + } + @Override public void exceptionCaught(ChannelHandlerContext context, Throwable cause) { diff --git a/processing/src/main/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStream.java b/processing/src/main/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStream.java index 5c147e776df4..a9ca04d84565 100644 --- a/processing/src/main/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStream.java +++ b/processing/src/main/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStream.java @@ -37,7 +37,7 @@ public class AppendableByteArrayInputStream extends InputStream private volatile boolean done = false; private volatile Throwable throwable; - private volatile int available = 0; + private volatile long available = 0; private byte[] curr = new byte[]{}; private int currIndex = 0; @@ -49,6 +49,9 @@ public void add(byte[] bytesToAdd) } synchronized (singleByteReaderDoer) { + if (done) { + return; + } bytes.addLast(bytesToAdd); available += bytesToAdd.length; singleByteReaderDoer.notify(); @@ -68,10 +71,18 @@ public void exceptionCaught(Throwable t) synchronized (singleByteReaderDoer) { done = true; throwable = t; + bytes.clear(); + available = 0; singleByteReaderDoer.notifyAll(); } } + @Override + public void close() + { + exceptionCaught(new IOException("Stream closed")); + } + @Override public int read() throws IOException { @@ -145,7 +156,7 @@ private long scanThroughBytesAndDoSomething(long numToScan, Doer doer) throws IO break; } try { - available -= numPulled; + releaseUnreadBytes(numPulled); numPulled = 0; singleByteReaderDoer.wait(); } @@ -181,16 +192,24 @@ private long scanThroughBytesAndDoSomething(long numToScan, Doer doer) throws IO } synchronized (singleByteReaderDoer) { - available -= numPulled; + releaseUnreadBytes(numPulled); } return numScanned; } + private void releaseUnreadBytes(long numPulled) + { + // exceptionCaught zeroes the count, bytes pulled concurrently included. + if (throwable == null) { + available -= numPulled; + } + } + @Override public int available() { - return available; + return (int) Math.min(available, Integer.MAX_VALUE); } private interface Doer diff --git a/processing/src/main/java/org/apache/druid/java/util/http/client/response/HttpResponseHandler.java b/processing/src/main/java/org/apache/druid/java/util/http/client/response/HttpResponseHandler.java index 4b8127b2d926..cc880046a4e9 100644 --- a/processing/src/main/java/org/apache/druid/java/util/http/client/response/HttpResponseHandler.java +++ b/processing/src/main/java/org/apache/druid/java/util/http/client/response/HttpResponseHandler.java @@ -87,6 +87,10 @@ ClientResponse handleChunk( void exceptionCaught(ClientResponse clientResponse, Throwable e); + /** + * Flow control over the connection carrying one response. Once that response has completed, calls have no effect: + * the connection may already be carrying the response of another request. + */ interface TrafficCop { /** diff --git a/processing/src/test/java/org/apache/druid/java/util/http/client/FriendlyServersTest.java b/processing/src/test/java/org/apache/druid/java/util/http/client/FriendlyServersTest.java index 098843c7d422..629174b26858 100644 --- a/processing/src/test/java/org/apache/druid/java/util/http/client/FriendlyServersTest.java +++ b/processing/src/test/java/org/apache/druid/java/util/http/client/FriendlyServersTest.java @@ -32,6 +32,8 @@ import org.apache.druid.java.util.http.client.response.HttpResponseHandler; import org.apache.druid.java.util.http.client.response.StatusResponseHandler; import org.apache.druid.java.util.http.client.response.StatusResponseHolder; +import org.apache.druid.java.util.http.client.response.StringFullResponseHandler; +import org.apache.druid.java.util.http.client.response.StringFullResponseHolder; import org.eclipse.jetty.server.Connector; import org.eclipse.jetty.server.HttpConfiguration; import org.eclipse.jetty.server.HttpConnectionFactory; @@ -58,11 +60,13 @@ import java.nio.charset.StandardCharsets; import java.util.Locale; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; @@ -71,6 +75,9 @@ */ public class FriendlyServersTest { + private static final String CHUNKED_OK = "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n"; + private static final String LAST_CHUNK = "0\r\n\r\n"; + @Test public void testFriendlyHttpServer() throws Exception { @@ -753,4 +760,198 @@ public void exceptionCaught(ClientResponse clientResponse, Throwable e) lifecycle.stop(); } } + + /** + * An abort that arrives once the response is complete must not close the channel: it is back in the pool, and the + * next request may already be on it. + */ + @Test + public void testAbortAfterCompletionLeavesThePooledChannelOpen() throws Exception + { + final ExecutorService exec = Executors.newSingleThreadExecutor(); + final ServerSocket serverSocket = new ServerSocket(0); + final BlockingQueue wire = new LinkedBlockingQueue<>(); + serveChunkedFromWire(exec, serverSocket, wire); + + final Lifecycle lifecycle = new Lifecycle(); + try { + final HttpClient client = HttpClientInit.createClient( + HttpClientConfig.builder().withNumConnections(1).build(), + lifecycle + ); + final Request request = new Request( + HttpMethod.GET, + new URL(StringUtils.format("http://localhost:%d/", serverSocket.getLocalPort())) + ); + + final CapturingTrafficCop first = + new CapturingTrafficCop<>(new StringFullResponseHandler(StandardCharsets.UTF_8)); + final ListenableFuture firstResponse = client.go(request, first); + wire.put(CHUNKED_OK + chunk("a") + LAST_CHUNK); + Assertions.assertEquals("a", firstResponse.get(30, TimeUnit.SECONDS).getContent()); + + first.trafficCop.get().abort(); + + // The server answers on its only connection, so the second request succeeds only if it reuses the first's. + final ListenableFuture secondResponse = + client.go(request, new StringFullResponseHandler(StandardCharsets.UTF_8)); + wire.put(CHUNKED_OK + chunk("b") + LAST_CHUNK); + Assertions.assertEquals("b", secondResponse.get(30, TimeUnit.SECONDS).getContent()); + } + finally { + exec.shutdownNow(); + serverSocket.close(); + lifecycle.stop(); + } + } + + /** + * A resume that arrives once the response is complete must not touch the channel: it is back in the pool, and the + * next request on it may have suspended reads. + */ + @Test + public void testResumeAfterCompletionLeavesTheNextRequestSuspended() throws Exception + { + final ExecutorService exec = Executors.newSingleThreadExecutor(); + final ServerSocket serverSocket = new ServerSocket(0); + final BlockingQueue wire = new LinkedBlockingQueue<>(); + serveChunkedFromWire(exec, serverSocket, wire); + + final Lifecycle lifecycle = new Lifecycle(); + try { + final HttpClient client = HttpClientInit.createClient( + HttpClientConfig.builder().withNumConnections(1).build(), + lifecycle + ); + final Request request = new Request( + HttpMethod.GET, + new URL(StringUtils.format("http://localhost:%d/", serverSocket.getLocalPort())) + ); + + // The first request suspends too: only then does its late resume pass its own watermark and reach the channel. + final SuspendAfterFirstChunk first = + new SuspendAfterFirstChunk<>(new StringFullResponseHandler(StandardCharsets.UTF_8)); + final ListenableFuture firstResponse = client.go(request, first); + wire.put(CHUNKED_OK + chunk("a") + LAST_CHUNK); + Assertions.assertEquals("a", firstResponse.get(30, TimeUnit.SECONDS).getContent()); + + final SuspendAfterFirstChunk second = + new SuspendAfterFirstChunk<>(new StringFullResponseHandler(StandardCharsets.UTF_8)); + final ListenableFuture secondResponse = client.go(request, second); + wire.put(CHUNKED_OK + chunk("b")); + Assertions.assertTrue(second.suspended.await(30, TimeUnit.SECONDS)); + + first.trafficCop.get().resume(1); + wire.put(chunk("c") + LAST_CHUNK); + Assertions.assertThrows(TimeoutException.class, () -> secondResponse.get(1, TimeUnit.SECONDS)); + + second.trafficCop.get().resume(1); + Assertions.assertEquals("bc", secondResponse.get(30, TimeUnit.SECONDS).getContent()); + } + finally { + exec.shutdownNow(); + serverSocket.close(); + lifecycle.stop(); + } + } + + private static String chunk(String data) + { + return StringUtils.format("%x\r\n%s\r\n", data.length(), data); + } + + /** + * Serves one keep-alive connection. Each request is answered with the segments taken from {@code wire}, up to the + * one that ends the chunked body, so the test decides when each part of a response goes out. + */ + private static void serveChunkedFromWire(ExecutorService exec, ServerSocket serverSocket, BlockingQueue wire) + { + exec.submit( + () -> { + try ( + Socket s = serverSocket.accept(); + BufferedReader in = new BufferedReader(new InputStreamReader(s.getInputStream(), StandardCharsets.UTF_8)); + OutputStream out = s.getOutputStream() + ) { + String line; + while ((line = in.readLine()) != null) { + if (line.isEmpty()) { + String segment; + do { + segment = wire.take(); + out.write(StringUtils.toUtf8(segment)); + out.flush(); + } while (!segment.endsWith(LAST_CHUNK)); + } + } + } + catch (Exception ignored) { + // suppress + } + } + ); + } + + /** + * Exposes the {@link HttpResponseHandler.TrafficCop} of the response it handles. + */ + private static class CapturingTrafficCop implements HttpResponseHandler + { + final HttpResponseHandler delegate; + final AtomicReference trafficCop = new AtomicReference<>(); + + CapturingTrafficCop(HttpResponseHandler delegate) + { + this.delegate = delegate; + } + + @Override + public ClientResponse handleResponse(HttpResponse response, TrafficCop trafficCop) + { + this.trafficCop.set(trafficCop); + return delegate.handleResponse(response, trafficCop); + } + + @Override + public ClientResponse handleChunk(ClientResponse clientResponse, HttpContent chunk, long chunkNum) + { + return delegate.handleChunk(clientResponse, chunk, chunkNum); + } + + @Override + public ClientResponse done(ClientResponse clientResponse) + { + return delegate.done(clientResponse); + } + + @Override + public void exceptionCaught(ClientResponse clientResponse, Throwable e) + { + delegate.exceptionCaught(clientResponse, e); + } + } + + /** + * Suspends reads after the first chunk; resume them through {@link #trafficCop}. + */ + private static class SuspendAfterFirstChunk extends CapturingTrafficCop + { + final CountDownLatch suspended = new CountDownLatch(1); + + SuspendAfterFirstChunk(HttpResponseHandler delegate) + { + super(delegate); + } + + @Override + public ClientResponse handleChunk(ClientResponse clientResponse, HttpContent chunk, long chunkNum) + { + final ClientResponse next = super.handleChunk(clientResponse, chunk, chunkNum); + if (chunkNum != 1) { + return next; + } + suspended.countDown(); + return new ClientResponse<>(next.isFinished(), false, next.getObj()); + } + } } diff --git a/processing/src/test/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStreamTest.java b/processing/src/test/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStreamTest.java index c337f7d9f70d..6d4d430f382c 100644 --- a/processing/src/test/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStreamTest.java +++ b/processing/src/test/java/org/apache/druid/java/util/http/client/io/AppendableByteArrayInputStreamTest.java @@ -211,6 +211,55 @@ public void testAddEmptyByteArray() throws Exception Assertions.assertEquals(-1, in.read()); } + @Test + public void testAvailableSaturatesInsteadOfOverflowing() throws Exception + { + final AppendableByteArrayInputStream in = new AppendableByteArrayInputStream(); + + // The same array is queued repeatedly: the stream retains references, so this accounts for 2 GiB of queued + // bytes while allocating only one chunk. + final byte[] oneMebibyte = new byte[1024 * 1024]; + for (int i = 0; i < 2048; i++) { + in.add(oneMebibyte); + } + + Assertions.assertEquals(Integer.MAX_VALUE, in.available()); + + in.read(new byte[oneMebibyte.length]); + + Assertions.assertEquals((2048L - 1) * oneMebibyte.length, (long) in.available()); + } + + @Test + public void testExceptionCaughtReleasesQueuedBytes() throws IOException + { + final AppendableByteArrayInputStream in = new AppendableByteArrayInputStream(); + in.add(new byte[10]); + in.add(new byte[8192]); + Assertions.assertEquals(5, in.read(new byte[5])); + + in.exceptionCaught(new IOException("connection reset")); + + Assertions.assertEquals(0, in.available()); + Assertions.assertEquals(5, in.read(new byte[5]), "the chunk being read is still handed out"); + Assertions.assertEquals(0, in.available()); + Assertions.assertThrows(IOException.class, () -> in.read(new byte[8192])); + } + + @Test + public void testCloseDiscardsQueuedAndLaterChunks() + { + final AppendableByteArrayInputStream in = new AppendableByteArrayInputStream(); + in.add(new byte[8192]); + in.add(new byte[8192]); + + in.close(); + in.add(new byte[8192]); + + Assertions.assertEquals(0, in.available()); + Assertions.assertThrows(IOException.class, in::read); + } + @Test public void testExceptionUnblocks() throws InterruptedException { diff --git a/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java b/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java index db8db5079b8b..37f73d74ea94 100644 --- a/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java +++ b/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java @@ -19,23 +19,13 @@ package org.apache.druid.client; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.annotations.VisibleForTesting; import com.google.inject.Inject; import jakarta.validation.constraints.NotNull; -import org.apache.druid.client.coordinator.Coordinator; import org.apache.druid.client.coordinator.CoordinatorClient; -import org.apache.druid.client.coordinator.CoordinatorClientImpl; -import org.apache.druid.discovery.NodeRole; -import org.apache.druid.guice.annotations.EscalatedGlobal; -import org.apache.druid.guice.annotations.Json; import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.QueryConfigProvider; import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryContexts; -import org.apache.druid.rpc.ServiceClientFactory; -import org.apache.druid.rpc.ServiceLocator; -import org.apache.druid.rpc.StandardRetryPolicy; import org.apache.druid.server.broker.BrokerDynamicConfig; import org.apache.druid.server.broker.QueryConfigSnapshot; @@ -64,27 +54,6 @@ public class BrokerViewOfBrokerConfig extends BaseBrokerViewOfConfig sourceCloneServers = Set.of(); @Inject - public BrokerViewOfCoordinatorConfig( - @Json final ObjectMapper jsonMapper, - @EscalatedGlobal final ServiceClientFactory clientFactory, - @Coordinator final ServiceLocator serviceLocator - ) - { - this.coordinatorClient = - new CoordinatorClientImpl( - clientFactory.makeClient( - NodeRole.COORDINATOR.getJsonName(), - serviceLocator, - StandardRetryPolicy.builder().maxAttempts(15).build() - ), - jsonMapper - ); - } - - @VisibleForTesting public BrokerViewOfCoordinatorConfig(CoordinatorClient coordinatorClient) { this.coordinatorClient = coordinatorClient; diff --git a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java index 2f963e0b270c..858930cc0d02 100644 --- a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java +++ b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java @@ -24,7 +24,6 @@ import com.fasterxml.jackson.dataformat.smile.SmileConstants; import com.fasterxml.jackson.dataformat.smile.SmileFactory; import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes; -import com.google.common.base.Preconditions; import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; @@ -222,8 +221,6 @@ private Failure(String message, Throwable cause) } } - private final AtomicReference trafficCopRef = new AtomicReference<>(); - private QueryMetrics> queryMetrics; private long responseStartTimeNs; @@ -257,7 +254,7 @@ private boolean enqueue(ByteBuf buffer, long chunkNum) throws InterruptedExcepti return !usingBackpressure || currentQueuedByteCount < maxQueuedBytes; } - private InputStream dequeue() throws InterruptedException + private InputStream dequeue(TrafficCop trafficCop) throws InterruptedException { final InputStreamHolder holder = queue.poll(checkQueryTimeout(), TimeUnit.MILLISECONDS); if (holder == null) { @@ -266,8 +263,7 @@ private InputStream dequeue() throws InterruptedException final long currentQueuedByteCount = queuedByteCount.addAndGet(-holder.getLength()); if (usingBackpressure && currentQueuedByteCount < maxQueuedBytes) { - long backPressureTime = Preconditions.checkNotNull(trafficCopRef.get(), "No TrafficCop, how can this be?") - .resume(holder.getChunkNum()); + long backPressureTime = trafficCop.resume(holder.getChunkNum()); channelSuspendedTime.addAndGet(backPressureTime); } @@ -424,7 +420,6 @@ private void throwForNonJsonBody( @Override public ClientResponse handleResponse(HttpResponse response, TrafficCop trafficCop) { - trafficCopRef.set(trafficCop); checkQueryTimeout(); // Netty 4: the initial HttpResponse carries no body, so the status and Content-Type are recorded // here and the body itself is inspected on the first HttpContent chunk. The goal is to detect a @@ -511,7 +506,7 @@ public InputStream nextElement() } try { - return dequeue(); + return dequeue(trafficCop); } catch (InterruptedException e) { Thread.currentThread().interrupt(); @@ -529,20 +524,8 @@ public InputStream nextElement() @Override public void close() { - final TrafficCop trafficCop; - synchronized (done) { - if (done.get()) { - return; - } - // Stop buffering further chunks (see enqueue()) and drop anything already buffered so the - // underlying Netty ChannelBuffers can be released. - discard.set(true); - queue.clear(); - trafficCop = trafficCopRef.get(); - } - if (trafficCop == null) { - return; - } + discard.set(true); + queue.clear(); trafficCop.abort(); } }, diff --git a/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClient.java b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClient.java index fb23fadf4894..bc2e0a71511c 100644 --- a/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClient.java +++ b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClient.java @@ -22,7 +22,7 @@ import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.client.BootstrapSegmentsResponse; import org.apache.druid.client.ImmutableSegmentLoadInfo; -import org.apache.druid.java.util.common.parsers.CloseableIterator; +import org.apache.druid.collections.ResourceHolder; import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.lookup.LookupExtractorFactoryContainer; import org.apache.druid.rpc.ServiceRetryPolicy; @@ -37,6 +37,7 @@ import javax.annotation.Nullable; import java.net.URI; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; @@ -137,12 +138,14 @@ public interface CoordinatorClient /** * Returns an iterator over the metadata segments of multiple datasources in the cluster, fetching them in one go. *

+ * The caller is responsible for closing the holder. + *

* API: {@code GET /druid/coordinator/v1/metadata/segments?includeOvershadowedStatus} * * @param watchedDataSources Optional datasources to filter the segments by. If null or empty, all segments are returned. * @param includeRealtimeSegments If true, includes realtime segments in the result. */ - ListenableFuture> fetchAllUsedSegmentsWithOvershadowedStatus( + ListenableFuture>> fetchAllUsedSegmentsWithOvershadowedStatus( @Nullable Set watchedDataSources, boolean includeRealtimeSegments ); diff --git a/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientConfig.java b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientConfig.java new file mode 100644 index 000000000000..448727c6931e --- /dev/null +++ b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientConfig.java @@ -0,0 +1,41 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.client.coordinator; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.Min; + +/** + * Configuration of the client that other services use to call the Coordinator. + */ +public class CoordinatorClientConfig +{ + /** + * Attempts per request to the Coordinator, the first one included, before giving up on a retryable error. + */ + @JsonProperty + @Min(1) + private int maxAttempts = 15; + + public int getMaxAttempts() + { + return maxAttempts; + } +} diff --git a/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientImpl.java b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientImpl.java index 1fed06decc92..9dfa0bab5b73 100644 --- a/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientImpl.java +++ b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientImpl.java @@ -28,10 +28,10 @@ import org.apache.druid.client.BootstrapSegmentsResponse; import org.apache.druid.client.ImmutableSegmentLoadInfo; import org.apache.druid.client.JsonParserIterator; +import org.apache.druid.collections.ResourceHolder; import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.jackson.JacksonUtils; -import org.apache.druid.java.util.common.parsers.CloseableIterator; import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler; import org.apache.druid.java.util.http.client.response.BytesFullResponseHolder; import org.apache.druid.java.util.http.client.response.InputStreamResponseHandler; @@ -52,6 +52,7 @@ import org.apache.druid.server.coordinator.rules.Rule; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentStatusInCluster; +import org.apache.druid.utils.CloseableUtils; import org.joda.time.Interval; import javax.annotation.Nullable; @@ -59,6 +60,7 @@ import java.net.URISyntaxException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; @@ -321,7 +323,7 @@ public Map fetchLookupsForTierSync(Stri } @Override - public ListenableFuture> fetchAllUsedSegmentsWithOvershadowedStatus( + public ListenableFuture>> fetchAllUsedSegmentsWithOvershadowedStatus( @Nullable Set watchedDataSources, boolean includeRealtimeSegments ) @@ -346,11 +348,25 @@ public ListenableFuture> fetchAllUsedS new InputStreamResponseHandler() ), inputStream -> { - return new JsonParserIterator<>( + final JsonParserIterator segments = new JsonParserIterator<>( jsonMapper.getTypeFactory().constructType(SegmentStatusInCluster.class), Futures.immediateFuture(inputStream), jsonMapper ); + return new ResourceHolder>() + { + @Override + public Iterator get() + { + return segments; + } + + @Override + public void close() + { + CloseableUtils.closeAndWrapExceptions(() -> CloseableUtils.closeAll(segments, inputStream)); + } + }; } ); } diff --git a/server/src/main/java/org/apache/druid/client/coordinator/NoopCoordinatorClient.java b/server/src/main/java/org/apache/druid/client/coordinator/NoopCoordinatorClient.java index b8aa55ec1e4a..24881e6994ba 100644 --- a/server/src/main/java/org/apache/druid/client/coordinator/NoopCoordinatorClient.java +++ b/server/src/main/java/org/apache/druid/client/coordinator/NoopCoordinatorClient.java @@ -22,7 +22,7 @@ import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.client.BootstrapSegmentsResponse; import org.apache.druid.client.ImmutableSegmentLoadInfo; -import org.apache.druid.java.util.common.parsers.CloseableIterator; +import org.apache.druid.collections.ResourceHolder; import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.lookup.LookupExtractorFactoryContainer; import org.apache.druid.rpc.ServiceRetryPolicy; @@ -37,6 +37,7 @@ import javax.annotation.Nullable; import java.net.URI; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; @@ -131,7 +132,7 @@ public Map fetchLookupsForTierSync( } @Override - public ListenableFuture> fetchAllUsedSegmentsWithOvershadowedStatus( + public ListenableFuture>> fetchAllUsedSegmentsWithOvershadowedStatus( @Nullable Set watchedDataSources, boolean includeOvershadowed ) diff --git a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java index 203886f99fea..1dbca72c24b2 100644 --- a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java +++ b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java @@ -27,10 +27,12 @@ import org.apache.druid.client.broker.BrokerClientImpl; import org.apache.druid.client.coordinator.Coordinator; import org.apache.druid.client.coordinator.CoordinatorClient; +import org.apache.druid.client.coordinator.CoordinatorClientConfig; import org.apache.druid.client.coordinator.CoordinatorClientImpl; import org.apache.druid.client.indexing.IndexingService; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; import org.apache.druid.discovery.NodeRole; +import org.apache.druid.guice.JsonConfigProvider; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.guice.annotations.EscalatedGlobal; @@ -57,7 +59,7 @@ public class ServiceClientModule implements DruidModule @Override public void configure(Binder binder) { - // Nothing to do. + JsonConfigProvider.bind(binder, "druid.client.coordinator", CoordinatorClientConfig.class); } @Provides @@ -112,13 +114,14 @@ public ServiceLocator makeCoordinatorServiceLocator(final DruidNodeDiscoveryProv @Coordinator public ServiceClient makeServiceClientForCoordinator( @EscalatedGlobal final ServiceClientFactory clientFactory, - @Coordinator final ServiceLocator serviceLocator + @Coordinator final ServiceLocator serviceLocator, + final CoordinatorClientConfig config ) { return clientFactory.makeClient( NodeRole.COORDINATOR.getJsonName(), serviceLocator, - StandardRetryPolicy.builder().maxAttempts(CLIENT_MAX_ATTEMPTS).build() + StandardRetryPolicy.builder().maxAttempts(config.getMaxAttempts()).build() ); } diff --git a/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java b/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java index 206e5df84462..d7c3c5a8e0bf 100644 --- a/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java +++ b/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java @@ -35,14 +35,15 @@ import org.apache.druid.client.BootstrapSegmentsResponse; import org.apache.druid.client.DruidServer; import org.apache.druid.client.ImmutableSegmentLoadInfo; +import org.apache.druid.collections.ResourceHolder; import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.guice.StartupInjectorBuilder; import org.apache.druid.initialization.CoreInjectorBuilder; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.java.util.common.parsers.CloseableIterator; import org.apache.druid.java.util.http.client.response.StringFullResponseHolder; +import org.apache.druid.query.QueryInterruptedException; import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.lookup.LookupExtractorFactory; import org.apache.druid.query.lookup.LookupExtractorFactoryContainer; @@ -78,6 +79,7 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Collections; +import java.util.Iterator; import java.util.LinkedHashSet; import java.util.List; import java.util.Map; @@ -575,10 +577,10 @@ public void test_fetchAllUsedSegmentsWithOvershadowedStatus_includeRealtime() th jsonMapper.writeValueAsBytes(segments) ); - CloseableIterator iterator = FutureUtils.getUnchecked( + Iterator iterator = FutureUtils.getUnchecked( coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus(null, true), true - ); + ).get(); List actualSegments = new ArrayList<>(); while (iterator.hasNext()) { actualSegments.add(iterator.next()); @@ -606,10 +608,10 @@ public void test_fetchAllUsedSegmentsWithOvershadowedStatus_noParams() throws Js jsonMapper.writeValueAsBytes(segments) ); - CloseableIterator iterator = FutureUtils.getUnchecked( + Iterator iterator = FutureUtils.getUnchecked( coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus(null, false), true - ); + ).get(); List actualSegments = new ArrayList<>(); while (iterator.hasNext()) { actualSegments.add(iterator.next()); @@ -622,6 +624,26 @@ public void test_fetchAllUsedSegmentsWithOvershadowedStatus_noParams() throws Js ); } + @Test + public void test_fetchAllUsedSegmentsWithOvershadowedStatus_closeBeforeIteratingReleasesTheResponse() + throws JsonProcessingException + { + serviceClient.expectAndRespond( + new RequestBuilder(HttpMethod.GET, "/druid/coordinator/v1/metadata/segments?includeOvershadowedStatus"), + HttpResponseStatus.OK, + ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), + jsonMapper.writeValueAsBytes(ImmutableList.of(SEGMENT1)) + ); + + final ResourceHolder> segments = FutureUtils.getUnchecked( + coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus(null, false), + true + ); + segments.close(); + + Assertions.assertThrows(QueryInterruptedException.class, () -> segments.get().hasNext()); + } + @Test public void test_fetchAllUsedSegmentsWithOvershadowedStatus_filterByDataSource() throws Exception { @@ -635,10 +657,10 @@ public void test_fetchAllUsedSegmentsWithOvershadowedStatus_filterByDataSource() jsonMapper.writeValueAsBytes(ImmutableList.of(SEGMENT3)) ); - CloseableIterator iterator = FutureUtils.getUnchecked( + Iterator iterator = FutureUtils.getUnchecked( coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus(Set.of("abc"), true), true - ); + ).get(); List actualSegments = new ArrayList<>(); while (iterator.hasNext()) { @@ -666,10 +688,10 @@ public void test_fetchAllUsedSegmentsWithOvershadowedStatus_filterByDataSources( ); Set dataSources = new LinkedHashSet<>(List.of("xyz", "abc")); - CloseableIterator iterator = FutureUtils.getUnchecked( + Iterator iterator = FutureUtils.getUnchecked( coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus(dataSources, true), true - ); + ).get(); List actualSegments = new ArrayList<>(); while (iterator.hasNext()) { @@ -696,10 +718,10 @@ public void test_fetchAllUsedSegmentsWithOvershadowedStatus_filterByDataSourceOn jsonMapper.writeValueAsBytes(List.of(SEGMENT3)) ); - CloseableIterator iterator = FutureUtils.getUnchecked( + Iterator iterator = FutureUtils.getUnchecked( coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus(Set.of("abc"), false), true - ); + ).get(); List actualSegments = new ArrayList<>(); while (iterator.hasNext()) { diff --git a/server/src/test/java/org/apache/druid/rpc/guice/ServiceClientModuleTest.java b/server/src/test/java/org/apache/druid/rpc/guice/ServiceClientModuleTest.java index f667cb8db5fa..033485dc575d 100644 --- a/server/src/test/java/org/apache/druid/rpc/guice/ServiceClientModuleTest.java +++ b/server/src/test/java/org/apache/druid/rpc/guice/ServiceClientModuleTest.java @@ -23,12 +23,16 @@ import com.google.inject.Guice; import com.google.inject.Injector; import com.google.inject.Key; +import com.google.inject.ProvisionException; import org.apache.druid.client.broker.Broker; import org.apache.druid.client.broker.BrokerClient; import org.apache.druid.client.coordinator.Coordinator; import org.apache.druid.client.coordinator.CoordinatorClient; +import org.apache.druid.client.coordinator.CoordinatorClientConfig; import org.apache.druid.client.indexing.IndexingService; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.guice.ConfigModule; import org.apache.druid.guice.DruidGuiceExtensions; import org.apache.druid.guice.LifecycleModule; import org.apache.druid.guice.annotations.EscalatedGlobal; @@ -37,16 +41,22 @@ import org.apache.druid.rpc.ServiceClient; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.rpc.ServiceLocator; +import org.apache.druid.rpc.ServiceRetryPolicy; import org.apache.druid.rpc.indexing.OverlordClient; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.ArgumentMatchers; import org.mockito.Mock; +import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; +import java.util.Properties; + @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.WARN) public class ServiceClientModuleTest @@ -68,9 +78,17 @@ public class ServiceClientModuleTest @BeforeEach public void setUp() { - injector = Guice.createInjector( + injector = makeInjector("4"); + } + + private Injector makeInjector(final String coordinatorMaxAttempts) + { + final Properties properties = new Properties(); + properties.setProperty("druid.client.coordinator.maxAttempts", coordinatorMaxAttempts); + return Guice.createInjector( ImmutableList.of( new DruidGuiceExtensions(), + new ConfigModule(), new LifecycleModule(), new JacksonModule(), new ServiceClientModule(), @@ -79,6 +97,7 @@ public void setUp() binder.bind(ServiceLocator.class).toInstance(serviceLocator); binder.bind(DruidNodeDiscoveryProvider.class).toInstance(discoveryProvider); binder.bind(ServiceClientFactory.class).toInstance(serviceClientFactory); + binder.bind(Properties.class).toInstance(properties); } ) ); @@ -102,6 +121,37 @@ public void testGetCoordinatorClient() Assertions.assertNotNull(injector.getInstance(CoordinatorClient.class)); } + @Test + public void testCoordinatorClientConfigRejectsNonPositiveMaxAttempts() + { + Assertions.assertThrows( + ProvisionException.class, + () -> makeInjector("0").getInstance(CoordinatorClientConfig.class) + ); + Assertions.assertThrows( + ProvisionException.class, + () -> makeInjector("-1").getInstance(CoordinatorClientConfig.class) + ); + } + + @Test + public void testCoordinatorRetryPolicyUsesConfiguredMaxAttempts() + { + new ServiceClientModule().makeServiceClientForCoordinator( + serviceClientFactory, + serviceLocator, + injector.getInstance(CoordinatorClientConfig.class) + ); + + final ArgumentCaptor retryPolicy = ArgumentCaptor.forClass(ServiceRetryPolicy.class); + Mockito.verify(serviceClientFactory).makeClient( + ArgumentMatchers.eq(NodeRole.COORDINATOR.getJsonName()), + ArgumentMatchers.eq(serviceLocator), + retryPolicy.capture() + ); + Assertions.assertEquals(4, retryPolicy.getValue().maxAttempts()); + } + @Test public void testGetBrokerClient() { diff --git a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java index 932b82cb41a3..c8f3c24eee62 100644 --- a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java +++ b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java @@ -28,6 +28,7 @@ import org.apache.druid.client.broker.BrokerClientImpl; import org.apache.druid.client.coordinator.Coordinator; import org.apache.druid.client.coordinator.CoordinatorClient; +import org.apache.druid.client.coordinator.CoordinatorClientConfig; import org.apache.druid.client.indexing.IndexingService; import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.discovery.NodeRole; @@ -105,7 +106,8 @@ private EmbeddedServiceClient(EmbeddedDruidCluster cluster, Escalator escalator) ); this.coordinatorServiceClient = module.makeServiceClientForCoordinator( factory, - anyServer.bindings().getInstance(ServiceLocator.class, Coordinator.class) + anyServer.bindings().getInstance(ServiceLocator.class, Coordinator.class), + new CoordinatorClientConfig() ); } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java index 77a5d01dcaea..0a7ca689a605 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java @@ -23,6 +23,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Joiner; import com.google.common.base.Preconditions; import com.google.common.base.Supplier; @@ -59,6 +60,8 @@ import org.apache.druid.error.InvalidSqlInput; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.guava.BaseSequence; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.SequenceWrapper; import org.apache.druid.java.util.common.guava.Sequences; import org.apache.druid.java.util.emitter.EmittingLogger; import org.apache.druid.query.Query; @@ -333,45 +336,48 @@ private PlannerResult planWithBindableConvention() planner.getTypeFactory(), plannerContext.getParameters() ); - final Supplier> resultsSupplier = () -> { - final Enumerable enumerable = theRel.bind(dataContext); - final Enumerator enumerator = enumerable.enumerator(); - return QueryResponse.withEmptyContext( - Sequences.withBaggage(new BaseSequence<>( - new BaseSequence.IteratorMaker>() - { - @Override - public QueryHandler.EnumeratorIterator make() - { - return new QueryHandler.EnumeratorIterator<>(new Iterator<>() - { - @Override - public boolean hasNext() - { - return enumerator.moveNext(); - } - - @Override - public Object[] next() - { - return (Object[]) enumerator.current(); - } - }); - } - - @Override - public void cleanup(QueryHandler.EnumeratorIterator iterFromMake) - { - - } - } - ), enumerator::close) - ); - }; + final Supplier> resultsSupplier = + () -> QueryResponse.withEmptyContext(enumerate(theRel.bind(dataContext))); return new PlannerResult(resultsSupplier, rootQueryRel.validatedRowType); } } + /** + * Rows of a bound {@link BindableRel}; closing the sequence releases everything the binding holds. + */ + @VisibleForTesting + static Sequence enumerate(final Enumerable enumerable) + { + final Sequence rows = new BaseSequence<>( + new BaseSequence.IteratorMaker() + { + @Override + public EnumeratorIterator make() + { + return new EnumeratorIterator(enumerable.enumerator()); + } + + @Override + public void cleanup(EnumeratorIterator iterFromMake) + { + iterFromMake.enumerator.close(); + } + } + ); + + if (!(enumerable instanceof AutoCloseable closeable)) { + return rows; + } + return Sequences.wrap(rows, new SequenceWrapper() + { + @Override + public void after(boolean isDone, Throwable thrown) throws Exception + { + closeable.close(); + } + }); + } + /** * Construct a {@link PlannerResult} for an 'explain' query from a {@link RelNode} and root {@link RelRoot} */ @@ -750,25 +756,25 @@ protected QueryMaker buildQueryMaker(final RelRoot rootQueryRel) throws Validati } } - private static class EnumeratorIterator implements Iterator + private static class EnumeratorIterator implements Iterator { - private final Iterator it; + private final Enumerator enumerator; - EnumeratorIterator(Iterator it) + EnumeratorIterator(Enumerator enumerator) { - this.it = it; + this.enumerator = enumerator; } @Override public boolean hasNext() { - return it.hasNext(); + return enumerator.moveNext(); } @Override - public T next() + public Object[] next() { - return it.next(); + return (Object[]) enumerator.current(); } } } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/MetadataSegmentView.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/MetadataSegmentView.java index b292e5e8e04a..e623efedbef9 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/MetadataSegmentView.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/MetadataSegmentView.java @@ -29,9 +29,11 @@ import org.apache.druid.client.BrokerSegmentWatcherConfig; import org.apache.druid.client.DataSegmentInterner; import org.apache.druid.client.coordinator.CoordinatorClient; +import org.apache.druid.collections.ResourceHolder; import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.concurrent.LifecycleLock; import org.apache.druid.guice.ManageLifecycle; +import org.apache.druid.java.util.common.CloseableIterators; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Stopwatch; import org.apache.druid.java.util.common.concurrent.Execs; @@ -47,6 +49,7 @@ import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; import org.apache.druid.timeline.SegmentStatusInCluster; +import org.apache.druid.utils.CloseableUtils; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.joda.time.Duration; @@ -154,26 +157,35 @@ private void poll() { log.info("Polling segments from coordinator"); final Stopwatch syncTime = Stopwatch.createStarted(); - final CloseableIterator metadataSegments = fetchSegmentMetadataFromCoordinator(); final ImmutableSortedSet.Builder builder = ImmutableSortedSet.naturalOrder(); - while (metadataSegments.hasNext()) { - final SegmentStatusInCluster segment = metadataSegments.next(); - final DataSegment interned = DataSegmentInterner.intern(segment.getDataSegment()); - Integer replicationFactor = segment.getReplicationFactor(); - if (replicationFactor == null) { - replicationFactor = segmentIdToReplicationFactor.getIfPresent(segment.getDataSegment().getId()); - } else { - segmentIdToReplicationFactor.put(segment.getDataSegment().getId(), segment.getReplicationFactor()); + final ResourceHolder> metadataSegments = fetchSegmentMetadataFromCoordinator(); + try { + final Iterator segments = metadataSegments.get(); + while (segments.hasNext()) { + final SegmentStatusInCluster segment = segments.next(); + final DataSegment interned = DataSegmentInterner.intern(segment.getDataSegment()); + Integer replicationFactor = segment.getReplicationFactor(); + if (replicationFactor == null) { + replicationFactor = segmentIdToReplicationFactor.getIfPresent(segment.getDataSegment().getId()); + } else { + segmentIdToReplicationFactor.put(segment.getDataSegment().getId(), segment.getReplicationFactor()); + } + final SegmentStatusInCluster segmentStatusInCluster = new SegmentStatusInCluster( + interned, + segment.isOvershadowed(), + replicationFactor, + segment.getNumRows(), + segment.isRealtime() + ); + builder.add(segmentStatusInCluster); } - final SegmentStatusInCluster segmentStatusInCluster = new SegmentStatusInCluster( - interned, - segment.isOvershadowed(), - replicationFactor, - segment.getNumRows(), - segment.isRealtime() + } + finally { + CloseableUtils.closeAndSuppressExceptions( + metadataSegments, + e -> log.warn(e, "Failed to close the segment metadata response from the Coordinator") ); - builder.add(segmentStatusInCluster); } publishedSegments = builder.build(); cachePopulated.countDown(); @@ -187,7 +199,7 @@ private void poll() * {@link BrokerSegmentMetadataCacheConfig#isMetadataSegmentCacheEnable()}) * OR by querying the Coordinator on the fly. */ - Iterator getSegments() + CloseableIterator getSegments() { return getSegments(null); } @@ -196,23 +208,27 @@ Iterator getSegments() * Returns published (and, with centralized schema, realtime) segment metadata, optionally * restricted to {@code dataSources}. */ - Iterator getSegments(@Nullable Set dataSources) + CloseableIterator getSegments(@Nullable Set dataSources) { - final Iterator base; + final CloseableIterator base; if (isCacheEnabled) { Uninterruptibles.awaitUninterruptibly(cachePopulated); - base = publishedSegments.iterator(); + base = CloseableIterators.withEmptyBaggage(publishedSegments.iterator()); } else { // Cache disabled: the Coordinator returns all used segments; filter client-side to preserve semantics. - base = fetchSegmentMetadataFromCoordinator(); + final ResourceHolder> fetched = fetchSegmentMetadataFromCoordinator(); + base = CloseableIterators.wrap(fetched.get(), fetched); } return dataSources == null ? base - : Iterators.filter(base, s -> dataSources.contains(s.getDataSegment().getDataSource())); + : CloseableIterators.wrap( + Iterators.filter(base, s -> dataSources.contains(s.getDataSegment().getDataSource())), + base + ); } // Note that coordinator must be up to get segments - private CloseableIterator fetchSegmentMetadataFromCoordinator() + private ResourceHolder> fetchSegmentMetadataFromCoordinator() { // includeRealtimeSegments flag would additionally request realtime segments // note that realtime segments are returned only when druid.centralizedDatasourceSchema.enabled is set on the Coordinator diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java index 885fc1ca3522..09e175d24151 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java @@ -55,6 +55,7 @@ import org.apache.druid.indexing.overlord.supervisor.SupervisorStatus; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.java.util.common.parsers.CloseableIterator; import org.apache.druid.java.util.http.client.HttpClient; @@ -441,10 +442,11 @@ public Enumerable scan( : new HashSet<>(); // Get segments from metadata segment cache (if enabled in SQL planner config), else directly from - // Coordinator. This may include both published and realtime segments. - final Iterator metadataStoreSegments = metadataView.getSegments(dataSourceFilter); + // Coordinator. This may include both published and realtime segments. Fetched on first read, so a scan that is + // never read holds no Coordinator response. + final Closer closer = Closer.create(); final FluentIterable publishedSegments = FluentIterable - .from(() -> getAuthorizedPublishedSegments(metadataStoreSegments)) + .from(() -> getAuthorizedPublishedSegments(closer.register(metadataView.getSegments(dataSourceFilter)))) .transform(val -> { final DataSegment segment = val.getDataSegment(); final AvailableSegmentMetadata availableSegmentMetadata = @@ -549,7 +551,7 @@ public Enumerable scan( Iterables.concat(publishedSegments, availableSegments) ); - return Linq4j.asEnumerable(allSegments) + return Linq4j.asEnumerable(() -> wrap(allSegments.iterator(), closer)) .where(Objects::nonNull) .select(row -> projectSegmentsRow(row, projects, jsonMapper)); } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/planner/QueryHandlerTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/planner/QueryHandlerTest.java new file mode 100644 index 000000000000..93d5f35c68b3 --- /dev/null +++ b/sql/src/test/java/org/apache/druid/sql/calcite/planner/QueryHandlerTest.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.sql.calcite.planner; + +import org.apache.calcite.interpreter.Interpreter; +import org.apache.calcite.linq4j.Enumerator; +import org.apache.calcite.linq4j.Linq4j; +import org.apache.druid.java.util.common.ISE; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.List; + +public class QueryHandlerTest +{ + private final Interpreter interpreter = Mockito.mock(Interpreter.class); + + @Test + public void testEnumerateClosesTheEnumeratorAndTheInterpreterOnceTheResultsAreConsumed() + { + final Enumerator enumerator = Mockito.spy(Linq4j.enumerator(List.of(new Object[]{1L}))); + Mockito.when(interpreter.enumerator()).thenReturn(enumerator); + + Assertions.assertEquals(1, QueryHandler.enumerate(interpreter).toList().size()); + + Mockito.verify(enumerator).close(); + Mockito.verify(interpreter).close(); + } + + @Test + public void testEnumerateClosesTheInterpreterIfEnumerationFails() + { + Mockito.when(interpreter.enumerator()).thenThrow(new ISE("Interpreter node failed")); + + Assertions.assertThrows(ISE.class, () -> QueryHandler.enumerate(interpreter).toList()); + + Mockito.verify(interpreter).close(); + } +} diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/schema/MetadataSegmentViewTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/schema/MetadataSegmentViewTest.java index 08b8d111b961..b84ef6818fea 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/schema/MetadataSegmentViewTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/schema/MetadataSegmentViewTest.java @@ -23,7 +23,9 @@ import com.google.common.util.concurrent.Futures; import org.apache.druid.client.BrokerSegmentWatcherConfig; import org.apache.druid.client.coordinator.CoordinatorClient; -import org.apache.druid.java.util.common.CloseableIterators; +import org.apache.druid.collections.ResourceHolder; +import org.apache.druid.collections.StupidResourceHolder; +import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.metrics.NoopTaskHolder; import org.apache.druid.metadata.segment.cache.Metric; import org.apache.druid.segment.TestDataSource; @@ -31,14 +33,23 @@ import org.apache.druid.server.metrics.LatchableEmitter; import org.apache.druid.server.metrics.LatchableEmitterConfig; import org.apache.druid.timeline.SegmentStatusInCluster; +import org.apache.druid.utils.CloseableUtils; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; +import java.io.Closeable; +import java.io.IOException; import java.util.ArrayList; +import java.util.Iterator; import java.util.List; +import java.util.NoSuchElementException; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; public class MetadataSegmentViewTest { @@ -84,7 +95,7 @@ public void test_start_triggersSegmentPollFromCoordinator_ifCacheIsEnabled() ArgumentMatchers.eq(true) ) ).thenReturn( - Futures.immediateFuture(CloseableIterators.withEmptyBaggage(expectedSegments.iterator())) + Futures.immediateFuture(StupidResourceHolder.create(expectedSegments.iterator())) ); // Start the test target and wait for it to sync with the Coordinator @@ -96,4 +107,124 @@ public void test_start_triggersSegmentPollFromCoordinator_ifCacheIsEnabled() Assertions.assertEquals(10, observedSegments.size()); Assertions.assertEquals(expectedSegments, observedSegments); } + + @Test + public void test_poll_releasesTheCoordinatorResponse() + { + final AtomicBoolean closed = new AtomicBoolean(false); + + Mockito.when( + coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus( + ArgumentMatchers.eq(null), + ArgumentMatchers.eq(true) + ) + ).thenAnswer( + invocation -> Futures.immediateFuture( + holder(List.of().iterator(), () -> closed.set(true)) + ) + ); + + segmentView.start(); + emitter.waitForEvent(event -> event.hasMetricName(Metric.SYNC_DURATION_MILLIS)); + + Assertions.assertTrue(closed.get(), "poll must close the streamed response"); + } + + @Test + public void test_poll_releasesTheCoordinatorResponse_whenIterationFails() throws Exception + { + final CountDownLatch closed = new CountDownLatch(1); + + final Iterator failsMidStream = new Iterator<>() + { + @Override + public boolean hasNext() + { + throw new ISE("Coordinator response ended mid-stream"); + } + + @Override + public SegmentStatusInCluster next() + { + throw new NoSuchElementException(); + } + }; + + Mockito.when( + coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus( + ArgumentMatchers.eq(null), + ArgumentMatchers.eq(true) + ) + ).thenAnswer( + invocation -> Futures.immediateFuture(holder(failsMidStream, closed::countDown)) + ); + + segmentView.start(); + + try { + Assertions.assertTrue( + closed.await(30, TimeUnit.SECONDS), + "a failed poll must still release the streamed response" + ); + } + finally { + segmentView.stop(); + } + } + + @Test + public void test_getSegments_releasesTheCoordinatorResponse_ifCacheIsDisabledAndFiltered() throws IOException + { + final AtomicBoolean closed = new AtomicBoolean(false); + + Mockito.when( + coordinatorClient.fetchAllUsedSegmentsWithOvershadowedStatus( + ArgumentMatchers.eq(null), + ArgumentMatchers.eq(true) + ) + ).thenReturn( + Futures.immediateFuture( + holder(List.of().iterator(), () -> closed.set(true)) + ) + ); + + final MetadataSegmentView uncachedView = new MetadataSegmentView( + coordinatorClient, + new BrokerSegmentWatcherConfig(), + new BrokerSegmentMetadataCacheConfig() + { + @Override + public boolean isMetadataSegmentCacheEnable() + { + return false; + } + }, + emitter + ); + + uncachedView.getSegments(Set.of(TestDataSource.WIKI)).close(); + + Assertions.assertTrue(closed.get()); + } + + private static ResourceHolder> holder( + final Iterator segments, + final Closeable response + ) + { + return new ResourceHolder<>() + { + @Override + public Iterator get() + { + return segments; + } + + @Override + public void close() + { + CloseableUtils.closeAndWrapExceptions(response); + } + }; + } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java index ea3e981a8d90..28945c6c066b 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java @@ -35,6 +35,7 @@ import org.apache.calcite.DataContext; import org.apache.calcite.adapter.java.JavaTypeFactory; import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.linq4j.Enumerator; import org.apache.calcite.linq4j.QueryProvider; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeField; @@ -153,6 +154,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; public class SystemSchemaTest extends CalciteTestBase { @@ -707,7 +709,9 @@ public void testSegmentsTable() throws Exception new SegmentStatusInCluster(segment2, false, 0, null, false) )); - EasyMock.expect(metadataView.getSegments(EasyMock.anyObject())).andReturn(publishedSegments.iterator()).once(); + EasyMock.expect(metadataView.getSegments(EasyMock.anyObject())) + .andReturn(CloseableIterators.withEmptyBaggage(publishedSegments.iterator())) + .once(); EasyMock.replay(request, responseHolder, responseHandler, metadataView); DataContext dataContext = createDataContext(); @@ -828,7 +832,9 @@ public void testSegmentsTableWithProjection() throws JsonProcessingException new SegmentStatusInCluster(segment2, false, 0, null, false) )); - EasyMock.expect(metadataView.getSegments(EasyMock.anyObject())).andReturn(publishedSegments.iterator()).once(); + EasyMock.expect(metadataView.getSegments(EasyMock.anyObject())) + .andReturn(CloseableIterators.withEmptyBaggage(publishedSegments.iterator())) + .once(); EasyMock.replay(request, responseHolder, responseHandler, metadataView); DataContext dataContext = createDataContext(); @@ -884,6 +890,42 @@ public void testSegmentsTableWithProjection() throws JsonProcessingException ); } + @Test + public void testSegmentsTableReleasesTheCoordinatorResponseWhenClosedEarly() + { + final SegmentsTable segmentsTable = + new SegmentsTable(segmentMetadataCache, metadataView, MAPPER, authMapper, createAuthResult(Users.SUPER)); + final List publishedSegments = List.of( + new SegmentStatusInCluster(segment1, true, 2, null, false), + new SegmentStatusInCluster(segment2, false, 0, null, false) + ); + final AtomicBoolean closed = new AtomicBoolean(false); + + EasyMock.expect(metadataView.getSegments(EasyMock.anyObject())) + .andReturn(CloseableIterators.wrap(publishedSegments.iterator(), () -> closed.set(true))) + .once(); + + EasyMock.replay(request, responseHolder, responseHandler, metadataView); + final Enumerator rows = + segmentsTable.scan(createDataContext(), Collections.emptyList(), null).enumerator(); + Assertions.assertTrue(rows.moveNext()); + rows.close(); + + Assertions.assertTrue(closed.get()); + } + + @Test + public void testSegmentsTableDefersTheCoordinatorFetchUntilRowsAreRead() + { + final SegmentsTable segmentsTable = + new SegmentsTable(segmentMetadataCache, metadataView, MAPPER, authMapper, createAuthResult(Users.SUPER)); + + EasyMock.replay(request, responseHolder, responseHandler, metadataView); + segmentsTable.scan(createDataContext(), Collections.emptyList(), null).enumerator().close(); + + EasyMock.verify(metadataView); + } + @Test public void testServersTable() throws URISyntaxException { diff --git a/website/.spelling b/website/.spelling index 5c7d8ba7d5c1..7ca0afdf5332 100644 --- a/website/.spelling +++ b/website/.spelling @@ -2630,6 +2630,7 @@ SRE reachability ownership_mismatch giveup +retryable - ../docs/development/extensions-contrib/spectator-histogram.md SpectatorHistogram