diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 36e19fd..201e00e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -63,6 +63,21 @@ jobs: || mvn --projects ingester-grpc test || mvn --projects ingester-grpc test) + test_ingester_bulk_protocol: + needs: check + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-java@v4 + with: + java-version: '8' + distribution: 'zulu' + - name: Maven Test + run: mvn clean install -DskipTests=true -Dmaven.javadoc.skip=true -B -V + && (mvn --projects ingester-bulk-protocol test + || mvn --projects ingester-bulk-protocol test + || mvn --projects ingester-bulk-protocol test) + test_ingester_protocol: needs: check runs-on: ubuntu-latest @@ -79,7 +94,7 @@ jobs: || mvn --projects ingester-protocol test) integration_tests: - needs: [test_ingester_common, test_ingester_grpc, test_ingester_protocol] + needs: [test_ingester_common, test_ingester_grpc, test_ingester_bulk_protocol, test_ingester_protocol] runs-on: ubuntu-latest timeout-minutes: 10 steps: diff --git a/.github/workflows/mvn_publish.yml b/.github/workflows/mvn_publish.yml index c60475b..a203aac 100644 --- a/.github/workflows/mvn_publish.yml +++ b/.github/workflows/mvn_publish.yml @@ -17,6 +17,7 @@ name: Maven Publish on: release: types: [created] + workflow_dispatch: jobs: publish: @@ -33,6 +34,24 @@ jobs: gpg-private-key: ${{ secrets.MAVEN_GPG_PRIVATE_KEY }} # Value of the GPG private key to import gpg-passphrase: MAVEN_GPG_PASSPHRASE # env variable for GPG private key passphrase + - name: Verify snapshot version + if: github.event_name == 'workflow_dispatch' + run: | + version="$(mvn help:evaluate -Dexpression=project.version -q -DforceStdout)" + if [[ "$version" != *-SNAPSHOT ]]; then + echo "Manual publishing requires a -SNAPSHOT version, got: $version" + exit 1 + fi + + - name: Verify release version + if: github.event_name == 'release' + run: | + version="$(mvn help:evaluate -Dexpression=project.version -q -DforceStdout)" + if [[ "$version" == *-SNAPSHOT ]]; then + echo "Release publishing requires a non-SNAPSHOT version, got: $version" + exit 1 + fi + - name: Build with Maven run: mvn clean deploy --batch-mode -DskipTests -P release -B -U -e env: diff --git a/ingester-all/pom.xml b/ingester-all/pom.xml index 59628af..75e311d 100644 --- a/ingester-all/pom.xml +++ b/ingester-all/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-all diff --git a/ingester-bulk-protocol/pom.xml b/ingester-bulk-protocol/pom.xml index 2746261..8cc7e88 100644 --- a/ingester-bulk-protocol/pom.xml +++ b/ingester-bulk-protocol/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-bulk-protocol @@ -118,4 +118,25 @@ test + + + + java9-plus + + [9,) + + + + + org.apache.maven.plugins + maven-surefire-plugin + 2.12.4 + + --add-opens=java.base/java.nio=ALL-UNNAMED + + + + + + diff --git a/ingester-bulk-protocol/src/main/java/io/greptime/BulkWriteService.java b/ingester-bulk-protocol/src/main/java/io/greptime/BulkWriteService.java index 610b97a..3ceeb30 100644 --- a/ingester-bulk-protocol/src/main/java/io/greptime/BulkWriteService.java +++ b/ingester-bulk-protocol/src/main/java/io/greptime/BulkWriteService.java @@ -18,11 +18,16 @@ import com.google.protobuf.ByteString; import io.greptime.common.TimeoutCompletableFuture; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; import org.apache.arrow.flight.BulkFlightClient.ClientStreamListener; import org.apache.arrow.flight.BulkFlightClient.PutListener; @@ -60,6 +65,7 @@ public class BulkWriteService implements AutoCloseable { private final VectorSchemaRoot root; private final ClientStreamListener listener; private final AsyncPutListener metadataListener; + private final String tableName; private final long timeoutMs; /** @@ -86,9 +92,17 @@ public BulkWriteService( this.root = manager.createSchemaRoot(schema); this.metadataListener = new AsyncPutListener(); this.listener = manager.startPut(descriptor, this.metadataListener, maxRequestsInFlight, options); + this.tableName = diagnosticName(descriptor); this.timeoutMs = timeoutMs; } + static String diagnosticName(FlightDescriptor descriptor) { + if (descriptor.isCommand() || descriptor.getPath().isEmpty()) { + return "unknown"; + } + return String.join("/", descriptor.getPath()); + } + /** * Starts the bulk write stream with default IPC options. */ @@ -146,32 +160,73 @@ public boolean isStreamReady() { */ public PutStage putNext() { long id = nextId(); - long totalRowCount = this.root.getRowCount(); + int totalRowCount = this.root.getRowCount(); LOG.debug("Starting putNext operation [id={}], total row count: {}", id, totalRowCount); // Create future with timeout and attach to listener - IdentifiableCompletableFuture future = new IdentifiableCompletableFuture(id, this.timeoutMs); - this.metadataListener.attach(id, future); - - // Prepare metadata buffer - byte[] metadata = new Metadata.RequestMetadata(id).toJsonBytesUtf8(); + IdentifiableCompletableFuture future = + new IdentifiableCompletableFuture(id, this.timeoutMs, this.tableName, totalRowCount, LOG); + ArrowBuf metadataBuf = null; + Throwable putFailure = null; try { - ArrowBuf metadataBuf = this.allocator.buffer(metadata.length); + if (!this.metadataListener.attach(id, future)) { + return new PutStage(future, this.metadataListener.numInFlight()); + } + + // Prepare metadata buffer + byte[] metadata = new Metadata.RequestMetadata(id).toJsonBytesUtf8(); + metadataBuf = this.allocator.buffer(metadata.length); metadataBuf.writeBytes(metadata); // Send data to the server LOG.debug("Sending data to server [id={}]", id); this.listener.putNext(metadataBuf); + metadataBuf = null; // Ownership transfers to the Flight writer. int inFlightCount = this.metadataListener.numInFlight(); LOG.debug("Data sent successfully [id={}], in-flight requests: {}", id, inFlightCount); return new PutStage(future, inFlightCount); + } catch (RuntimeException | Error e) { + Throwable failure = e; + if (!future.completeExceptionally(e)) { + try { + future.join(); + } catch (CompletionException completedFailure) { + if (completedFailure.getCause() != null) { + failure = completedFailure.getCause(); + } + } + } + putFailure = failure; + // Flight closes metadata after accepting ownership; a positive refcount here means the + // failure occurred before that handoff completed and this method still owns the buffer. + if (metadataBuf != null && metadataBuf.refCnt() > 0) { + try { + metadataBuf.close(); + } catch (RuntimeException closeError) { + failure.addSuppressed(closeError); + } + } + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } + throw new CompletionException(failure); } finally { // Clear the root to prepare for next batch - this.root.clear(); - LOG.debug("Cleared root for next batch [id={}], previous row count: {}", id, totalRowCount); + try { + this.root.clear(); + LOG.debug("Cleared root for next batch [id={}], previous row count: {}", id, totalRowCount); + } catch (RuntimeException | Error clearError) { + if (putFailure == null) { + throw clearError; + } + putFailure.addSuppressed(clearError); + } } } @@ -244,6 +299,19 @@ public CompletableFuture future() { public int numInFlight() { return this.numInFlight; } + + /** + * Logs diagnostic context when a caller times out waiting for this request. + * + * @param timeout the timeout raised to the caller + * @param value the caller's timeout value + * @param unit the caller's timeout unit + */ + public void logTimeout(Throwable timeout, long value, TimeUnit unit) { + if (this.future instanceof IdentifiableCompletableFuture) { + ((IdentifiableCompletableFuture) this.future).logTimeout(timeout, unit.toMillis(value)); + } + } } /** @@ -252,6 +320,11 @@ public int numInFlight() { */ static class IdentifiableCompletableFuture extends TimeoutCompletableFuture { private final long id; + private final long timeoutMs; + private final String tableName; + private final int rows; + private final Logger logger; + private final AtomicBoolean timeoutLogged = new AtomicBoolean(); /** * Creates a new IdentifiableCompletableFuture. @@ -260,8 +333,16 @@ static class IdentifiableCompletableFuture extends TimeoutCompletableFuture futuresInFlight; private final CompletableFuture completed; + private final Object terminalLock; + private boolean terminal; + private Throwable terminalFailure; /** * Creates a new AsyncPutListener. @@ -288,16 +395,7 @@ static class AsyncPutListener implements PutListener { AsyncPutListener() { this.futuresInFlight = new ConcurrentHashMap<>(); this.completed = new CompletableFuture<>(); - this.completed.whenComplete((r, t) -> { - if (t != null) { - // Also complete all the futures with the same exception - for (IdentifiableCompletableFuture future : this.futuresInFlight.values()) { - future.completeExceptionally(t); - } - } - // When completed, clear the futuresInFlight - this.futuresInFlight.clear(); - }); + this.terminalLock = new Object(); } /** @@ -305,16 +403,33 @@ static class AsyncPutListener implements PutListener { * * @param id The unique identifier for the request * @param future The future to track + * @return true if the future was attached, false if the stream was already terminal */ - public void attach(long id, IdentifiableCompletableFuture future) { - this.futuresInFlight.put(id, future); + public boolean attach(long id, IdentifiableCompletableFuture future) { + Throwable failure; + synchronized (this.terminalLock) { + if (this.terminal) { + failure = this.terminalFailure; + } else { + this.futuresInFlight.put(id, future); + failure = null; + } + } + + if (failure != null) { + future.completeExceptionally(failure); + return false; + } + future.whenComplete((r, t) -> { // Remove the future from the map when it's completed - this.futuresInFlight.remove(id); + this.futuresInFlight.remove(id, future); if (t != null) { - LOG.error("Put operation failed [id={}]: {}", id, t.getMessage(), t); - if (!(t instanceof TimeoutCompletableFuture.FutureDeadlineExceededException)) { + if (t instanceof TimeoutCompletableFuture.FutureDeadlineExceededException) { + future.logTimeout(t, future.timeoutMs); + } else { + LOG.error("Put operation failed [id={}]: {}", id, t.getMessage(), t); // If a put next operation fails, we complete the future with the exception // and the stream will be terminated immediately to prevent further operations onError(t); @@ -323,11 +438,13 @@ public void attach(long id, IdentifiableCompletableFuture future) { LOG.debug("Put operation succeeded [id={}], affected rows: {}", id, r); } }); + future.scheduleTimeout(); if (LOG.isDebugEnabled()) { LOG.debug("Attached future [id={}], current in-flight count: {}", id, this.futuresInFlight.size()); } + return true; } /** @@ -364,14 +481,52 @@ public void onNext(PutResult val) { @Override public void onError(Throwable t) { - LOG.error("Stream error occurred: {}", t.getMessage(), t); - this.completed.completeExceptionally(StatusUtils.fromThrowable(t)); + Throwable failure = StatusUtils.fromThrowable(t); + if (terminate(failure, false)) { + LOG.error("Stream error occurred: {}", t.getMessage(), t); + } } @Override public final void onCompleted() { - LOG.info("Server signaled stream completion"); - this.completed.complete(null); + if (terminate(null, true)) { + LOG.info("Server signaled stream completion"); + } + } + + private boolean terminate(Throwable failure, boolean normalCompletion) { + List pending; + Throwable completionFailure; + synchronized (this.terminalLock) { + if (this.terminal) { + return false; + } + + pending = new ArrayList<>(this.futuresInFlight.values()); + this.futuresInFlight.clear(); + if (normalCompletion && pending.isEmpty()) { + completionFailure = null; + this.terminalFailure = new IllegalStateException("The bulk write stream is already completed"); + } else if (failure != null) { + completionFailure = failure; + this.terminalFailure = failure; + } else { + completionFailure = new IllegalStateException( + "The bulk write stream completed before all put responses were received"); + this.terminalFailure = completionFailure; + } + this.terminal = true; + } + + if (completionFailure != null) { + this.completed.completeExceptionally(completionFailure); + for (IdentifiableCompletableFuture future : pending) { + future.completeExceptionally(completionFailure); + } + } else { + this.completed.complete(null); + } + return true; } @Override diff --git a/ingester-bulk-protocol/src/test/java/io/greptime/BulkWriteServiceTest.java b/ingester-bulk-protocol/src/test/java/io/greptime/BulkWriteServiceTest.java new file mode 100644 index 0000000..3af3738 --- /dev/null +++ b/ingester-bulk-protocol/src/test/java/io/greptime/BulkWriteServiceTest.java @@ -0,0 +1,361 @@ +/* + * Copyright 2023 Greptime Team + * + * Licensed 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 io.greptime; + +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.arrow.flight.BulkFlightClient.ClientStreamListener; +import org.apache.arrow.flight.BulkFlightClient.PutListener; +import org.apache.arrow.flight.CallOption; +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightRuntimeException; +import org.apache.arrow.flight.FlightStatusCode; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; +import org.slf4j.Logger; + +public class BulkWriteServiceTest { + + @Test + public void testDiagnosticNameDerivedFromDescriptor() { + Assert.assertEquals( + "catalog/schema/metrics", + BulkWriteService.diagnosticName(FlightDescriptor.path("catalog", "schema", "metrics"))); + Assert.assertEquals("unknown", BulkWriteService.diagnosticName(FlightDescriptor.command(new byte[] {1}))); + } + + @Test + public void testAttachAfterStreamErrorFailsImmediately() throws Exception { + BulkWriteService.AsyncPutListener listener = Mockito.spy(new BulkWriteService.AsyncPutListener()); + listener.onError(CallStatus.UNAVAILABLE.toRuntimeException()); + + BulkWriteService.IdentifiableCompletableFuture future = newFuture(1L); + listener.attach(1L, future); + + Throwable failure = getFailure(future); + Assert.assertTrue(failure instanceof FlightRuntimeException); + Assert.assertEquals( + FlightStatusCode.UNAVAILABLE, + ((FlightRuntimeException) failure).status().code()); + Assert.assertEquals(0, listener.numInFlight()); + Mockito.verify(listener, Mockito.times(1)).onError(Mockito.any()); + } + + @Test + public void testAttachAfterNormalCompletionFailsImmediately() throws Exception { + BulkWriteService.AsyncPutListener listener = new BulkWriteService.AsyncPutListener(); + listener.onCompleted(); + + BulkWriteService.IdentifiableCompletableFuture future = newFuture(1L); + listener.attach(1L, future); + + Assert.assertTrue(getFailure(future) instanceof IllegalStateException); + Assert.assertEquals(0, listener.numInFlight()); + } + + @Test + public void testNormalCompletionFailsPendingFuture() throws Exception { + BulkWriteService.AsyncPutListener listener = new BulkWriteService.AsyncPutListener(); + BulkWriteService.IdentifiableCompletableFuture future = newFuture(1L); + listener.attach(1L, future); + + listener.onCompleted(); + + Assert.assertTrue(getFailure(future) instanceof IllegalStateException); + Assert.assertEquals(0, listener.numInFlight()); + Assert.assertTrue(listener.isCompletedExceptionally()); + } + + @Test + public void testConcurrentAttachAndStreamErrorAlwaysCompletesFuture() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 100; i++) { + BulkWriteService.AsyncPutListener listener = new BulkWriteService.AsyncPutListener(); + BulkWriteService.IdentifiableCompletableFuture future = newFuture(i + 1L); + CountDownLatch start = new CountDownLatch(1); + + Future attach = executor.submit(() -> { + await(start); + listener.attach(future.getId(), future); + }); + Future terminate = executor.submit(() -> { + await(start); + listener.onError(CallStatus.UNAVAILABLE.toRuntimeException()); + }); + + start.countDown(); + attach.get(1, TimeUnit.SECONDS); + terminate.get(1, TimeUnit.SECONDS); + + Assert.assertTrue("Future was left unresolved", future.isCompletedExceptionally()); + Assert.assertEquals(0, listener.numInFlight()); + } + } finally { + executor.shutdownNow(); + } + } + + @Test + public void testStreamResultCompletesBeforePendingFutureCallbacksRun() throws Exception { + BulkWriteService.AsyncPutListener listener = new BulkWriteService.AsyncPutListener(); + BulkWriteService.IdentifiableCompletableFuture future = newFuture(1L); + AtomicBoolean callbackFinished = new AtomicBoolean(); + AtomicBoolean callbackObservedStreamError = new AtomicBoolean(); + future.whenComplete((r, t) -> { + try { + listener.getResult(); + } catch (FlightRuntimeException e) { + callbackObservedStreamError.set(true); + } + callbackFinished.set(true); + }); + listener.attach(1L, future); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future terminate = executor.submit(() -> listener.onError(CallStatus.UNAVAILABLE.toRuntimeException())); + terminate.get(1, TimeUnit.SECONDS); + } finally { + executor.shutdownNow(); + } + + Assert.assertTrue(callbackFinished.get()); + Assert.assertTrue(callbackObservedStreamError.get()); + } + + @Test + public void testPutNextDoesNotSendAfterStreamTermination() throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ServiceFixture fixture = newServiceFixture(allocator); + try (BulkWriteService service = fixture.service) { + fixture.metadataListener.onError(CallStatus.UNAVAILABLE.toRuntimeException()); + + BulkWriteService.PutStage stage = service.putNext(); + + Throwable failure = getFailure(stage.future()); + Assert.assertTrue(failure instanceof FlightRuntimeException); + Assert.assertEquals( + FlightStatusCode.UNAVAILABLE, + ((FlightRuntimeException) failure).status().code()); + Mockito.verify(fixture.stream, Mockito.never()).putNext(Mockito.any()); + Assert.assertEquals(0, fixture.metadataListener.numInFlight()); + } + } + } + + @Test + public void testPutNextCleansUpFutureWhenSendFailsSynchronously() throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ServiceFixture fixture = newServiceFixture(allocator); + RuntimeException sendFailure = new RuntimeException("send failed"); + Mockito.doThrow(sendFailure).when(fixture.stream).putNext(Mockito.any()); + + try (BulkWriteService service = fixture.service) { + try { + service.putNext(); + Assert.fail("Expected putNext to fail"); + } catch (RuntimeException e) { + Assert.assertSame(sendFailure, e); + } + + Assert.assertEquals(0, fixture.metadataListener.numInFlight()); + Assert.assertTrue(fixture.metadataListener.isCompletedExceptionally()); + } + } + } + + @Test + public void testRootCleanupDoesNotMaskSynchronousSendFailure() throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + BulkWriteManager manager = Mockito.mock(BulkWriteManager.class); + ClientStreamListener stream = Mockito.mock(ClientStreamListener.class); + VectorSchemaRoot root = Mockito.mock(VectorSchemaRoot.class); + Schema schema = new Schema(Collections.emptyList()); + FlightDescriptor descriptor = FlightDescriptor.path("metrics"); + RuntimeException sendFailure = new RuntimeException("send failed"); + RuntimeException cleanupFailure = new RuntimeException("cleanup failed"); + + Mockito.when(manager.createSchemaRoot(schema)).thenReturn(root); + Mockito.when(manager.startPut( + Mockito.eq(descriptor), + Mockito.any(PutListener.class), + Mockito.eq(1L), + Mockito.any())) + .thenReturn(stream); + Mockito.doThrow(sendFailure).when(stream).putNext(Mockito.any()); + Mockito.doThrow(cleanupFailure).when(root).clear(); + + try (BulkWriteService service = new BulkWriteService(manager, allocator, schema, descriptor, 60000L, 1)) { + try { + service.putNext(); + Assert.fail("Expected putNext to fail"); + } catch (RuntimeException e) { + Assert.assertSame(sendFailure, e); + Assert.assertArrayEquals(new Throwable[] {cleanupFailure}, e.getSuppressed()); + } + } + } + } + + @Test + public void testPutNextDoesNotDoubleCloseConsumedMetadataWhenWriterFails() throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ServiceFixture fixture = newServiceFixture(allocator); + RuntimeException sendFailure = new RuntimeException("send failed"); + Mockito.doAnswer(invocation -> { + ArrowBuf metadata = invocation.getArgument(0); + metadata.close(); + Assert.assertEquals(0, metadata.refCnt()); + throw sendFailure; + }) + .when(fixture.stream) + .putNext(Mockito.any()); + + try (BulkWriteService service = fixture.service) { + try { + service.putNext(); + Assert.fail("Expected putNext to fail"); + } catch (RuntimeException e) { + Assert.assertSame(sendFailure, e); + } + } + } + } + + @Test + public void testTerminalFailureWinsOverLaterSynchronousWriterFailure() throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ServiceFixture fixture = newServiceFixture(allocator); + RuntimeException writerFailure = new RuntimeException("writer failed"); + Mockito.doAnswer(invocation -> { + fixture.metadataListener.onError(CallStatus.UNAVAILABLE.toRuntimeException()); + throw writerFailure; + }) + .when(fixture.stream) + .putNext(Mockito.any()); + + try (BulkWriteService service = fixture.service) { + try { + service.putNext(); + Assert.fail("Expected putNext to fail"); + } catch (FlightRuntimeException e) { + Assert.assertEquals(FlightStatusCode.UNAVAILABLE, e.status().code()); + Assert.assertNotSame(writerFailure, e); + } + } + } + } + + @Test + public void testTimedGetLogsBulkWriteRequestContextOnce() throws Exception { + Logger logger = Mockito.mock(Logger.class); + BulkWriteService.IdentifiableCompletableFuture future = + new BulkWriteService.IdentifiableCompletableFuture(42L, 60000L, "metrics", 100, logger); + + TimeoutException timeout = null; + try { + future.get(1, TimeUnit.MILLISECONDS); + Assert.fail("Expected timed get to fail"); + } catch (TimeoutException e) { + timeout = e; + } + future.logTimeout(timeout, 60000L); + + Mockito.verify(logger) + .warn( + "Bulk write timed out - table={}, request-id={}, rows={}, timeout={}ms", + "metrics", + 42L, + 100, + 1L, + timeout); + } + + private static BulkWriteService.IdentifiableCompletableFuture newFuture(long id) { + return new BulkWriteService.IdentifiableCompletableFuture(id, TimeUnit.MINUTES.toMillis(1)); + } + + private static Throwable getFailure(java.util.concurrent.CompletableFuture future) throws Exception { + try { + future.get(1, TimeUnit.SECONDS); + Assert.fail("Expected future to fail"); + return null; + } catch (ExecutionException e) { + return e.getCause(); + } + } + + private static void await(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + private static ServiceFixture newServiceFixture(BufferAllocator allocator) { + BulkWriteManager manager = Mockito.mock(BulkWriteManager.class); + ClientStreamListener stream = Mockito.mock(ClientStreamListener.class); + Schema schema = new Schema(Collections.emptyList()); + FlightDescriptor descriptor = FlightDescriptor.path("metrics"); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + ArgumentCaptor metadataListener = ArgumentCaptor.forClass(PutListener.class); + + Mockito.when(manager.createSchemaRoot(schema)).thenReturn(root); + Mockito.when(manager.startPut( + Mockito.eq(descriptor), + metadataListener.capture(), + Mockito.eq(1L), + Mockito.any())) + .thenReturn(stream); + + BulkWriteService service = new BulkWriteService(manager, allocator, schema, descriptor, 60000L, 1); + return new ServiceFixture(service, stream, (BulkWriteService.AsyncPutListener) metadataListener.getValue()); + } + + private static class ServiceFixture { + private final BulkWriteService service; + private final ClientStreamListener stream; + private final BulkWriteService.AsyncPutListener metadataListener; + + private ServiceFixture( + BulkWriteService service, + ClientStreamListener stream, + BulkWriteService.AsyncPutListener metadataListener) { + this.service = service; + this.stream = stream; + this.metadataListener = metadataListener; + } + } +} diff --git a/ingester-common/pom.xml b/ingester-common/pom.xml index 39ad459..057fc64 100644 --- a/ingester-common/pom.xml +++ b/ingester-common/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-common diff --git a/ingester-example/pom.xml b/ingester-example/pom.xml index d574f89..442536d 100644 --- a/ingester-example/pom.xml +++ b/ingester-example/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-example diff --git a/ingester-grpc/pom.xml b/ingester-grpc/pom.xml index 2e7e8a3..225c7d3 100644 --- a/ingester-grpc/pom.xml +++ b/ingester-grpc/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-grpc diff --git a/ingester-integration-tests/pom.xml b/ingester-integration-tests/pom.xml index bcf0873..38a205c 100644 --- a/ingester-integration-tests/pom.xml +++ b/ingester-integration-tests/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-integration-tests diff --git a/ingester-prometheus-metrics/pom.xml b/ingester-prometheus-metrics/pom.xml index 2222178..146ec59 100644 --- a/ingester-prometheus-metrics/pom.xml +++ b/ingester-prometheus-metrics/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-prometheus-metrics ${project.groupId}:${project.artifactId} diff --git a/ingester-protocol/pom.xml b/ingester-protocol/pom.xml index 05c3f38..1d03587 100644 --- a/ingester-protocol/pom.xml +++ b/ingester-protocol/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-protocol diff --git a/ingester-protocol/src/main/java/io/greptime/BulkWriteClient.java b/ingester-protocol/src/main/java/io/greptime/BulkWriteClient.java index fc94aac..8d5c681 100644 --- a/ingester-protocol/src/main/java/io/greptime/BulkWriteClient.java +++ b/ingester-protocol/src/main/java/io/greptime/BulkWriteClient.java @@ -38,8 +38,10 @@ import io.greptime.rpc.TlsOptions; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicReference; import org.apache.arrow.flight.FlightCallHeaders; import org.apache.arrow.flight.HeaderCallOption; @@ -236,11 +238,13 @@ public CompletableFuture writeNext() throws Exception { "Stream busy with pending requests. Check `isStreamReady()` before calling `writeNext()` to avoid busy-waiting."); } - return this.pipelineWriteLimiter.acquireAndDo(null, () -> { + AtomicReference putStage = new AtomicReference<>(); + CompletableFuture result = this.pipelineWriteLimiter.acquireAndDo(null, () -> { Clock clock = Clock.defaultClock(); long startPut = clock.getTick(); BulkWriteService.PutStage stage = this.writer.putNext(); + putStage.set(stage); InnerMetricHelper.prepareTime().update(clock.duration(startPut), TimeUnit.MILLISECONDS); long startCall = clock.getTick(); @@ -262,6 +266,7 @@ public CompletableFuture writeNext() throws Exception { return future; }); + return new TimeoutLoggingFuture(result, putStage.get()); } @Override @@ -282,6 +287,32 @@ public void close() throws Exception { } } + static class TimeoutLoggingFuture extends CompletableFuture { + private final BulkWriteService.PutStage putStage; + + TimeoutLoggingFuture(CompletableFuture delegate, BulkWriteService.PutStage putStage) { + this.putStage = putStage; + delegate.whenComplete((r, t) -> { + if (t == null) { + complete(r); + } else { + completeExceptionally(t); + } + }); + } + + @Override + public Integer get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + try { + return super.get(timeout, unit); + } catch (TimeoutException e) { + this.putStage.logTimeout(e, timeout, unit); + throw e; + } + } + } + /** * Limiter that controls the number of concurrent bulk write operations. * Uses a blocking policy to ensure the maximum number of in-flight requests is not exceeded. diff --git a/ingester-protocol/src/main/java/io/greptime/limit/AbstractLimiter.java b/ingester-protocol/src/main/java/io/greptime/limit/AbstractLimiter.java index 1bb5c88..7ebc0f0 100644 --- a/ingester-protocol/src/main/java/io/greptime/limit/AbstractLimiter.java +++ b/ingester-protocol/src/main/java/io/greptime/limit/AbstractLimiter.java @@ -68,7 +68,12 @@ public CompletableFuture acquireAndDo(In in, Supplier release(permits)); + try { + return action.get().whenComplete((r, e) -> release(permits)); + } catch (RuntimeException | Error e) { + release(permits); + throw e; + } } return Util.completedCf(rejected(in, acquirePermits, maxPermits)); } finally { diff --git a/ingester-protocol/src/test/java/io/greptime/BulkWriteClientTest.java b/ingester-protocol/src/test/java/io/greptime/BulkWriteClientTest.java new file mode 100644 index 0000000..f597956 --- /dev/null +++ b/ingester-protocol/src/test/java/io/greptime/BulkWriteClientTest.java @@ -0,0 +1,43 @@ +/* + * Copyright 2023 Greptime Team + * + * Licensed 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 io.greptime; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +public class BulkWriteClientTest { + + @Test + public void testTimedGetReportsCallerTimeoutToPutStage() throws Exception { + BulkWriteService.PutStage stage = Mockito.mock(BulkWriteService.PutStage.class); + CompletableFuture future = new BulkWriteClient.TimeoutLoggingFuture(new CompletableFuture<>(), stage); + + TimeoutException timeout = null; + try { + future.get(1, TimeUnit.MILLISECONDS); + Assert.fail("Expected timed get to fail"); + } catch (TimeoutException e) { + timeout = e; + } + + Mockito.verify(stage).logTimeout(timeout, 1, TimeUnit.MILLISECONDS); + } +} diff --git a/ingester-protocol/src/test/java/io/greptime/WriteLimitTest.java b/ingester-protocol/src/test/java/io/greptime/WriteLimitTest.java index 274d7ac..459d3e4 100644 --- a/ingester-protocol/src/test/java/io/greptime/WriteLimitTest.java +++ b/ingester-protocol/src/test/java/io/greptime/WriteLimitTest.java @@ -25,9 +25,10 @@ import io.greptime.models.WriteOk; import java.util.Collection; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Assert; import org.junit.Test; @@ -61,6 +62,24 @@ public void discardWriteLimitTest() throws ExecutionException, InterruptedExcept Assert.assertEquals(Result.FLOW_CONTROL, ret.getErr().getCode()); } + @Test + public void supplierFailureReleasesWritePermit() throws Exception { + WriteLimiter limiter = new WriteClient.DefaultWriteLimiter(1, new LimitedPolicy.AbortPolicy()); + Collection rows = TestUtil.testTable("test1", 1); + RuntimeException failure = new RuntimeException("write failed"); + + try { + limiter.acquireAndDo(rows, () -> { + throw failure; + }); + Assert.fail("Expected write action to fail"); + } catch (RuntimeException e) { + Assert.assertSame(failure, e); + } + + Assert.assertTrue(limiter.acquireAndDo(rows, this::emptyOk).get().isOk()); + } + @Test public void blockingWriteLimitTest() throws InterruptedException { WriteLimiter limiter = new WriteClient.DefaultWriteLimiter(1, new LimitedPolicy.BlockingPolicy()); @@ -69,24 +88,25 @@ public void blockingWriteLimitTest() throws InterruptedException { // consume the permits limiter.acquireAndDo(rows, CompletableFuture::new); - final AtomicBoolean alwaysFalse = new AtomicBoolean(); + CountDownLatch acquiring = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); final Thread t = new Thread(() -> { + acquiring.countDown(); try { limiter.acquireAndDo(rows, this::emptyOk); - alwaysFalse.set(true); } catch (Throwable err) { - // noinspection ConstantConditions - Assert.assertTrue(err instanceof InterruptedException); + failure.set(err); } }); t.start(); - Assert.assertFalse(alwaysFalse.get()); - Thread.sleep(1000); - Assert.assertFalse(alwaysFalse.get()); + acquiring.await(); t.interrupt(); - Assert.assertFalse(alwaysFalse.get()); - Assert.assertTrue(t.isInterrupted()); + t.join(TimeUnit.SECONDS.toMillis(1)); + + Assert.assertFalse("Limiter thread did not stop after interruption", t.isAlive()); + Assert.assertTrue(failure.get() instanceof LimitedException); + Assert.assertTrue(failure.get().getCause() instanceof InterruptedException); } @Test diff --git a/ingester-rpc/pom.xml b/ingester-rpc/pom.xml index d8f16cc..0f8bc78 100644 --- a/ingester-rpc/pom.xml +++ b/ingester-rpc/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT ingester-rpc diff --git a/pom.xml b/pom.xml index 4714f05..882df0d 100644 --- a/pom.xml +++ b/pom.xml @@ -21,7 +21,7 @@ io.greptime greptimedb-ingester - 0.15.0 + 0.15.1-SNAPSHOT pom ${project.groupId}:${project.artifactId} @@ -284,10 +284,11 @@ org.sonatype.central central-publishing-maven-plugin - 0.4.0 + 0.11.0 true central + https://central.sonatype.com/repository/maven-snapshots/ true