Skip to content

Commit

Permalink
Fix TransportResponse reference counting in DirectResponseChannel (#9…
Browse files Browse the repository at this point in the history 
…1289) (#91315)

In #76474 we fixed a circuit breaker leak in TransportActionProxy 
by incrementing a reference on the TransportResponse that is later 
decremented by the OutboundHandler.

This works well for all cases except when the request targets the 
node which is also the proxy node. In that case the reference is 
incremented but will never be decremented as the local execution 
(using TransportService#localNodeConnection and 
DirectResponseChannel) bypasses the OutboundHandler.

This change fixes the ref counting by also decrementing the 
TransportResponse in DirectResponseChannel.

This will also have the consequence to correctly decrement used 
bytes of the request circuit breaker when 
GetCcrRestoreFileChunkResponse are executed on a node that 
is also a proxy node.
  • Loading branch information
@tlrx
tlrx committed Nov 4, 2022
1 parent 297a0e2 commit 55cde1a
Show file tree
Hide file tree
Showing 3 changed files with 132 additions and 29 deletions.
5 changes: 5 additions & 0 deletions docs/changelog/91289.yaml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
pr: 91289
summary: Fix `TransportActionProxy` for local execution
area: Network
type: bug
issues: []
60 changes: 35 additions & 25 deletions server/src/main/java/org/elasticsearch/transport/TransportService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1411,33 +1411,43 @@ public String getProfileName() {

@Override
public void sendResponse(TransportResponse response) throws IOException {
service.onResponseSent(requestId, action, response);
try (var shutdownBlock = service.pendingDirectHandlers.withRef()) {
if (shutdownBlock == null) {
// already shutting down, the handler will be completed by sendRequestInternal or doStop
return;
}
final TransportResponseHandler<?> handler = service.responseHandlers.onResponseReceived(requestId, service);
if (handler == null) {
// handler already completed, likely by a timeout which is logged elsewhere
return;
}
final String executor = handler.executor();
if (ThreadPool.Names.SAME.equals(executor)) {
processResponse(handler, response);
} else {
threadPool.executor(executor).execute(new ForkingResponseHandlerRunnable(handler, null) {
@Override
protected void doRun() {
processResponse(handler, response);
}
try {
service.onResponseSent(requestId, action, response);
try (var shutdownBlock = service.pendingDirectHandlers.withRef()) {
if (shutdownBlock == null) {
// already shutting down, the handler will be completed by sendRequestInternal or doStop
return;
}
final TransportResponseHandler<?> handler = service.responseHandlers.onResponseReceived(requestId, service);
if (handler == null) {
// handler already completed, likely by a timeout which is logged elsewhere
return;
}
final String executor = handler.executor();
if (ThreadPool.Names.SAME.equals(executor)) {
processResponse(handler, response);
} else {
response.incRef();
threadPool.executor(executor).execute(new ForkingResponseHandlerRunnable(handler, null) {
@Override
protected void doRun() {
processResponse(handler, response);
}

@Override
public String toString() {
return "delivery of response to [" + requestId + "][" + action + "]: " + response;
}
});
@Override
public void onAfter() {
response.decRef();
}

@Override
public String toString() {
return "delivery of response to [" + requestId + "][" + action + "]: " + response;
}
});
}
}
} finally {
response.decRef();
}
}

Expand Down
96 changes: 92 additions & 4 deletions server/src/test/java/org/elasticsearch/transport/TransportActionProxyTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.core.AbstractRefCounted;
import org.elasticsearch.core.IOUtils;
import org.elasticsearch.core.RefCounted;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
Expand All @@ -27,8 +29,10 @@
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;

import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.notNullValue;

public class TransportActionProxyTests extends ESTestCase {
protected ThreadPool threadPool;
Expand Down Expand Up @@ -76,8 +80,9 @@ private MockTransportService buildService(final Version version) {
public void testSendMessage() throws InterruptedException {
serviceA.registerRequestHandler("internal:test", ThreadPool.Names.SAME, SimpleTestRequest::new, (request, channel, task) -> {
assertEquals(request.sourceNode, "TS_A");
SimpleTestResponse response = new SimpleTestResponse("TS_A");
final SimpleTestResponse response = new SimpleTestResponse("TS_A");
channel.sendResponse(response);
assertThat(response.hasReferences(), equalTo(false));
});
final boolean cancellable = randomBoolean();
TransportActionProxy.registerProxyAction(serviceA, "internal:test", cancellable, SimpleTestResponse::new);
Expand All @@ -86,21 +91,24 @@ public void testSendMessage() throws InterruptedException {
serviceB.registerRequestHandler("internal:test", ThreadPool.Names.SAME, SimpleTestRequest::new, (request, channel, task) -> {
assertThat(task instanceof CancellableTask, equalTo(cancellable));
assertEquals(request.sourceNode, "TS_A");
SimpleTestResponse response = new SimpleTestResponse("TS_B");
final SimpleTestResponse response = new SimpleTestResponse("TS_B");
channel.sendResponse(response);
assertThat(response.hasReferences(), equalTo(false));
});
TransportActionProxy.registerProxyAction(serviceB, "internal:test", cancellable, SimpleTestResponse::new);
AbstractSimpleTransportTestCase.connectToNode(serviceB, nodeC);
serviceC.registerRequestHandler("internal:test", ThreadPool.Names.SAME, SimpleTestRequest::new, (request, channel, task) -> {
assertThat(task instanceof CancellableTask, equalTo(cancellable));
assertEquals(request.sourceNode, "TS_A");
SimpleTestResponse response = new SimpleTestResponse("TS_C");
final SimpleTestResponse response = new SimpleTestResponse("TS_C");
channel.sendResponse(response);
assertThat(response.hasReferences(), equalTo(false));
});

TransportActionProxy.registerProxyAction(serviceC, "internal:test", cancellable, SimpleTestResponse::new);

CountDownLatch latch = new CountDownLatch(1);
final CountDownLatch latch = new CountDownLatch(1);
// Node A -> Node B -> Node C
serviceA.sendRequest(
nodeB,
TransportActionProxy.getProxyAction("internal:test"),
Expand Down Expand Up @@ -133,6 +141,61 @@ public void handleException(TransportException exp) {
latch.await();
}

public void testSendLocalRequest() throws InterruptedException {
final AtomicReference<SimpleTestResponse> response = new AtomicReference<>();
final boolean cancellable = randomBoolean();
serviceB.registerRequestHandler(
"internal:test",
randomFrom(ThreadPool.Names.SAME, ThreadPool.Names.GENERIC),
SimpleTestRequest::new,
(request, channel, task) -> {
assertThat(task instanceof CancellableTask, equalTo(cancellable));
assertEquals(request.sourceNode, "TS_A");
final SimpleTestResponse responseB = new SimpleTestResponse("TS_B");
channel.sendResponse(responseB);
response.set(responseB);
}
);
TransportActionProxy.registerProxyAction(serviceB, "internal:test", cancellable, SimpleTestResponse::new);
AbstractSimpleTransportTestCase.connectToNode(serviceA, nodeB);

final CountDownLatch latch = new CountDownLatch(1);
// Node A -> Proxy Node B (Local execution)
serviceA.sendRequest(
nodeB,
TransportActionProxy.getProxyAction("internal:test"),
TransportActionProxy.wrapRequest(nodeB, new SimpleTestRequest("TS_A", cancellable)), // Request
new TransportResponseHandler<SimpleTestResponse>() {
@Override
public SimpleTestResponse read(StreamInput in) throws IOException {
return new SimpleTestResponse(in);
}

@Override
public void handleResponse(SimpleTestResponse response) {
try {
assertEquals("TS_B", response.targetNode);
} finally {
latch.countDown();
}
}

@Override
public void handleException(TransportException exp) {
try {
throw new AssertionError(exp);
} finally {
latch.countDown();
}
}
}
);
latch.await();

assertThat(response.get(), notNullValue());
assertThat(response.get().hasReferences(), equalTo(false));
}

public void testException() throws InterruptedException {
boolean cancellable = randomBoolean();
serviceA.registerRequestHandler("internal:test", ThreadPool.Names.SAME, SimpleTestRequest::new, (request, channel, task) -> {
Expand Down Expand Up @@ -230,7 +293,12 @@ public boolean shouldCancelChildrenOnCancellation() {
}

public static class SimpleTestResponse extends TransportResponse {

final String targetNode;
final RefCounted refCounted = new AbstractRefCounted() {
@Override
protected void closeInternal() {}
};

SimpleTestResponse(String targetNode) {
this.targetNode = targetNode;
Expand All @@ -245,6 +313,26 @@ public static class SimpleTestResponse extends TransportResponse {
public void writeTo(StreamOutput out) throws IOException {
out.writeString(targetNode);
}

@Override
public void incRef() {
refCounted.incRef();
}

@Override
public boolean tryIncRef() {
return refCounted.tryIncRef();
}

@Override
public boolean decRef() {
return refCounted.decRef();
}

@Override
public boolean hasReferences() {
return refCounted.hasReferences();
}
}

public void testGetAction() {
Expand Down

0 comments on commit 55cde1a

Please sign in to comment.