diff --git a/core/src/main/java/tech/ydb/core/impl/BaseGrpcTransport.java b/core/src/main/java/tech/ydb/core/impl/BaseGrpcTransport.java index bdd3504f5..e15eb5c5a 100644 --- a/core/src/main/java/tech/ydb/core/impl/BaseGrpcTransport.java +++ b/core/src/main/java/tech/ydb/core/impl/BaseGrpcTransport.java @@ -119,13 +119,14 @@ public CompletableFuture> unaryCall( String traceId = settings.getTraceId(); try { - GrpcChannel channel = getChannel(settings); - EndpointRecord endpoint = channel.getEndpoint(); CallOptions options = prepareCallOptions(settings); if (options == null) { return CompletableFuture.completedFuture(deadlineExpiredResult(method, settings)); } + GrpcChannel channel = getChannel(settings); + EndpointRecord endpoint = channel.getEndpoint(); + ClientCall call = channel.getReadyChannel().newCall(method, options); ChannelStatusHandler handler = new ChannelStatusHandler(channel, settings); @@ -158,13 +159,14 @@ public GrpcReadStream readStreamCall( String traceId = settings.getTraceId(); try { - GrpcChannel channel = getChannel(settings); - EndpointRecord endpoint = channel.getEndpoint(); CallOptions options = prepareCallOptions(settings); if (options == null) { return new EmptyStream<>(deadlineExpiredStatus(method, settings)); } + GrpcChannel channel = getChannel(settings); + EndpointRecord endpoint = channel.getEndpoint(); + ClientCall call = channel.getReadyChannel().newCall(method, options); ChannelStatusHandler handler = new ChannelStatusHandler(channel, settings); @@ -200,13 +202,14 @@ public GrpcReadWriteStream readWriteStreamCall( String traceId = settings.getTraceId(); try { - GrpcChannel channel = getChannel(settings); - EndpointRecord endpoint = channel.getEndpoint(); CallOptions options = prepareCallOptions(settings); if (options == null) { return new EmptyStream<>(deadlineExpiredStatus(method, settings)); } + GrpcChannel channel = getChannel(settings); + EndpointRecord endpoint = channel.getEndpoint(); + ClientCall call = channel.getReadyChannel().newCall(method, options); ChannelStatusHandler hdlr = new ChannelStatusHandler(channel, settings); diff --git a/table/src/main/java/tech/ydb/table/impl/pool/WaitingQueue.java b/table/src/main/java/tech/ydb/table/impl/pool/WaitingQueue.java index 494e6a21a..a82572118 100644 --- a/table/src/main/java/tech/ydb/table/impl/pool/WaitingQueue.java +++ b/table/src/main/java/tech/ydb/table/impl/pool/WaitingQueue.java @@ -60,6 +60,8 @@ default void destroy(T object, PoolMetrics.Reason reason) { /** Size of waiting acquires queue */ private final AtomicInteger waitingAcqueireCount = new AtomicInteger(); + private final ThreadLocal localGuard = new ThreadLocal<>(); + @VisibleForTesting WaitingQueue(Handler handler, int maxSize, int waitingsLimit) { Preconditions.checkArgument(maxSize > 0, "WaitingQueue max size (%s) must be positive", maxSize); @@ -110,6 +112,11 @@ public void release(T object) { return; } + WaitingGuard guard = localGuard.get(); + if (guard != null && !guard.isAllowed(object)) { + return; + } + // Try to complete waiting request if (!tryToCompleteWaiting(object)) { // if queue is overflowed @@ -252,18 +259,20 @@ private boolean tryToCompleteWaiting(T object) { return false; } - CompletableFuture next = waitingAcquires.poll(); - while (next != null) { - waitingAcqueireCount.decrementAndGet(); + try (WaitingGuard guard = new WaitingGuard(object)) { + CompletableFuture next = waitingAcquires.poll(); + while (next != null) { + guard.init(); + waitingAcqueireCount.decrementAndGet(); + if (safeAcquireObject(next, object) && !guard.isBroken()) { + return true; + } - if (safeAcquireObject(next, object)) { - return true; + next = waitingAcquires.poll(); } - next = waitingAcquires.poll(); + return false; } - - return false; } private void checkNextWaitingAcquire() { @@ -307,6 +316,39 @@ private void clear() { } } + private final class WaitingGuard implements AutoCloseable { + private final WaitingGuard prev; + private final Object obj; + private boolean isBroken = false; + + WaitingGuard(Object obj) { + this.prev = localGuard.get(); + this.obj = obj; + } + + public void init() { + localGuard.set(this); + isBroken = false; + } + + @Override + public void close() { + localGuard.set(prev); + } + + private boolean isAllowed(Object object) { + if (object != obj) { + return true; + } + isBroken = true; + return false; + } + + public boolean isBroken() { + return isBroken; + } + } + private static class Limits { private final int maxSize; private final int waitingsLimit; diff --git a/table/src/test/java/tech/ydb/table/impl/pool/WaitingQueueTest.java b/table/src/test/java/tech/ydb/table/impl/pool/WaitingQueueTest.java index 76233a173..9616ebf2e 100644 --- a/table/src/test/java/tech/ydb/table/impl/pool/WaitingQueueTest.java +++ b/table/src/test/java/tech/ydb/table/impl/pool/WaitingQueueTest.java @@ -28,6 +28,10 @@ public static class Resource { Resource(int id) { this.id = id; } + + public int getId() { + return id; + } } private static class ResourceHandler implements WaitingQueue.Handler { @@ -576,6 +580,77 @@ public void canceledWaitingTest() { check(rs).requestsCount(0).activeCount(0); } + @Test + public void immediatellyCompletedWaitingsTest() { + ResourceHandler rs = new ResourceHandler(); + WaitingQueue queue = new WaitingQueue<>(rs, 1, 9999); + + @SuppressWarnings("unchecked") + CompletableFuture[] wa = (CompletableFuture[]) new CompletableFuture[10000]; + for (int idx = 0; idx < 10000; idx++) { + wa[idx] = new CompletableFuture<>(); + queue.acquire(wa[idx]); + wa[idx].thenAccept(queue::release); + } + + check(queue).queueSize(1).idleSize(0).waitingsCount(9999); + rs.completeNext(); // all 10000 acquires must be completed immediatelly + check(queue).queueSize(1).idleSize(1).waitingsCount(0); + + queue.close(); + check(queue).queueSize(0).idleSize(0).waitingsCount(0); + check(rs).requestsCount(0).activeCount(0); + } + + @Test + public void immediatellyCompletedOtherTest() { + ResourceHandler rs = new ResourceHandler(); + WaitingQueue queue = new WaitingQueue<>(rs, 2, 3); + + CompletableFuture r1 = pendingFuture(acquire(queue)); + CompletableFuture r2 = pendingFuture(acquire(queue)); + check(queue).queueSize(2).idleSize(0).waitingsCount(0); + + rs.completeNext().completeNext(); + Resource a = pendingIsReady(r1); + Resource b = pendingIsReady(r2); + check(queue).queueSize(2).idleSize(0).waitingsCount(0); + + CompletableFuture w1 = pendingFuture(acquire(queue)); + CompletableFuture w2 = pendingFuture(acquire(queue)); + w1.thenAccept(ignored -> queue.release(b)); + check(queue).queueSize(2).idleSize(0).waitingsCount(2); + + queue.release(a); + + Assert.assertSame(a, pendingIsReady(w1)); + Assert.assertSame(b, pendingIsReady(w2)); + check(queue).queueSize(2).idleSize(0).waitingsCount(0); + } + + @Test + public void currentReleaseBeforeOtherMustNotBeLost() { + ResourceHandler rs = new ResourceHandler(); + WaitingQueue queue = new WaitingQueue<>(rs, 2, 3); + + CompletableFuture first = pendingFuture(acquire(queue)); + CompletableFuture second = pendingFuture(acquire(queue)); + rs.completeNext().completeNext(); + Resource a = pendingIsReady(first); + Resource b = pendingIsReady(second); + + CompletableFuture waiting = pendingFuture(acquire(queue)); + waiting.thenAccept(ignored -> { + queue.release(a); + queue.release(b); + }); + + queue.release(a); + + Assert.assertSame(a, pendingIsReady(waiting)); + check(queue).queueSize(2).idleSize(2).waitingsCount(0); + } + @Test public void checkWaitingAfterDeleteTest() { ResourceHandler rs = new ResourceHandler();