ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
lidavidm wants to merge 3 commits into from
Closed

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut #8010

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c3a66cc
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…
lidavidm Aug 19, 2020
1e535b7
ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation
lidavidm Aug 31, 2020
80abecb
Fix TestLeak taking the full 60 second timeout
lidavidm Aug 31, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
final ClientStreamListener writer = new PutObserver(
descriptor, observer, stream.completed::isDone,
descriptor, observer, stream.cancelled::isDone,
() -> {
try {
stream.completed.get();
Expand Down
2 changes: 1 addition & 1 deletion java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
@Override
public String toString() {
String s = getClass().getName();
return String.format("%s: %s", s, status);
return String.format("%s: %s: %s", s, status.code(), status.description());
}
}
53 changes: 24 additions & 29 deletions java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
responseObserver.disableAutoInboundFlowControl();
responseObserver.request(1);

final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
}, responseObserver::request);
final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
.wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
final FlightStream fs = new FlightStream(
allocator,
PENDING_REQUESTS,
/* server-upload streams are not cancellable */null,
responseObserver::request);
// When the ackStream is completed, the FlightStream will be closed with it
ackStream.setAutoCloseable(fs);
final StreamObserver<ArrowMessage> observer = fs.asObserver();
executors.submit(() -> {
final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
.wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
try {
producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
} catch (Exception ex) {
ackStream.onError(ex);
} finally {
// Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
try {
fs.close();
} catch (Exception e) {
handleExceptionWithMiddleware(e);
}
// ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
// We don't do this for other streams since the implementation may be asynchronous
ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
*/
private void handleExceptionWithMiddleware(Throwable t) {
final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
if (middleware == null) {
if (middleware == null || middleware.isEmpty()) {
logger.error("Uncaught exception in Flight method body", t);
return;
}
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

/** Ensures that other resources are cleaned up when the service finishes its call. */
private static class ExchangeListener extends GetListener {
private final AutoCloseable resource;

private AutoCloseable resource;
private boolean closed = false;
private Runnable onCancelHandler = null;

public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
AutoCloseable resource) {
public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
super(responseObserver, errorHandler);
this.resource = resource;
this.resource = null;
super.setOnCancelHandler(() -> {
try {
if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
}
closed = true;
try {
this.resource.close();
AutoCloseables.close(resource);
} catch (Exception e) {
throw CallStatus.INTERNAL
.withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
final ServerCallStreamObserver<ArrowMessage> responseObserver =
(ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
}, responseObserver::request);
// When service completes the call, this cleans up the FlightStream
final ExchangeListener listener = new ExchangeListener(
responseObserver,
this::handleExceptionWithMiddleware,
() -> {
// Force the stream to "complete" so it will close without incident. At this point, we don't care since
// we are about to end the call. (Normally it will raise an error.)
fs.completed.complete(null);
fs.close();
});
this::handleExceptionWithMiddleware);
final FlightStream fs = new FlightStream(
allocator,
PENDING_REQUESTS,
/* server-upload streams are not cancellable */null,
responseObserver::request);
// When service completes the call, this cleans up the FlightStream
listener.resource = fs;
responseObserver.disableAutoInboundFlowControl();
responseObserver.request(1);
final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down
121 changes: 85 additions & 36 deletions java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
* An adaptor between protobuf streams and flight data streams.
*/
public class FlightStream implements AutoCloseable {

// Use AutoCloseable sentinel objects to simplify logic in #close
private final AutoCloseable DONE = () -> {
};
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
private final int pendingTarget;
private final Requestor requestor;
// The completion flags.
// This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
// data and closed the stream
final CompletableFuture<Void> completed;
// This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
// we don't block forever trying to write to a server that has rejected a call.
final CompletableFuture<Void> cancelled;

private volatile int pending = 1;
private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
*
* @param allocator The allocator to use for creating/reallocating buffers for Vectors.
* @param pendingTarget Target number of messages to receive.
* @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
* @param cancellable Used to cancel mid-stream requests.
* @param requestor A callback to determine how many pending items there are.
*/
public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
Objects.requireNonNull(allocator);
Objects.requireNonNull(requestor);
this.allocator = allocator;
this.pendingTarget = pendingTarget;
this.cancellable = cancellable;
this.requestor = requestor;
this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
this.completed = new CompletableFuture<>();
this.cancelled = new CompletableFuture<>();
}

/**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
/**
* Closes the stream (freeing any existing resources).
*
* <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
* <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
*/
public void close() throws Exception {
final List<AutoCloseable> closeables = new ArrayList<>();
// cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
closeables.add(() -> {
if (!completed.isDone() && cancellable != null) {
cancel("Stream closed before end.", /* no exception to report */ null);
Throwable suppressor = null;
if (cancellable != null) {
// Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
// On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
// may never run https://github.com/grpc/grpc-java/issues/5882
try {
synchronized (cancellable) {
if (!cancelled.isDone()) {
// Only cancel if the call is not done on the gRPC side
cancellable.cancel("Stream closed before end", /* no exception to report */null);
}
}
// Drain the stream without the lock (as next() implicitly needs the lock)
while (next()) { }
} catch (FlightRuntimeException e) {
suppressor = e;
}
});
if (fulfilledRoot != null) {
closeables.add(fulfilledRoot);
}
closeables.add(applicationMetadata);
closeables.addAll(queue);
if (dictionaries != null) {
dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
// Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the

@rymurrrymurrSep 5, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification

@lidavidmlidavidmSep 5, 2020

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.

// middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
// by the lambda above.
synchronized (completed) {
try {
if (fulfilledRoot != null) {
closeables.add(fulfilledRoot);
}
closeables.add(applicationMetadata);
closeables.addAll(queue);
if (dictionaries != null) {
dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
}
if (suppressor != null) {
AutoCloseables.close(suppressor, closeables);
} else {
AutoCloseables.close(closeables);
}
} finally {
// The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
// No-op if already complete
completed.complete(null);
}
}

AutoCloseables.close(closeables);
// Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
// No-op if already complete; do this after the check in the AutoCloseable lambda above
completed.complete(null);
}

/**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
super();
}

/** Helper to add an item to the queue under the appropriate lock. */
private void enqueue(AutoCloseable message) {
synchronized (completed) {
if (completed.isDone()) {
// The stream is already closed (RPC ended), discard the message
AutoCloseables.closeNoChecked(message);
} else {
queue.add(message);
}
}
}

@Override
public void onNext(ArrowMessage msg) {
// Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
// close().
requestOutstanding();
switch (msg.getMessageType()) {
case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
descriptor.set(new FlightDescriptor(msg.getDescriptor()));
}
if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
}
break;
}
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
try {
MetadataV4UnionChecker.checkRead(schema, metadataVersion);
} catch (IOException e) {
queue.add(DONE_EX);
ex = e;
enqueue(DONE_EX);
break;
}

fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
loader = new VectorLoader(fulfilledRoot);
if (msg.getDescriptor() != null) {
descriptor.set(new FlightDescriptor(msg.getDescriptor()));
synchronized (completed) {
if (!completed.isDone()) {
fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
loader = new VectorLoader(fulfilledRoot);
if (msg.getDescriptor() != null) {
descriptor.set(new FlightDescriptor(msg.getDescriptor()));
}
root.set(fulfilledRoot);
}
}
root.set(fulfilledRoot);
break;
}
case RECORD_BATCH:
queue.add(msg);
break;
case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
break;
case TENSOR:
default:
queue.add(DONE_EX);
ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
enqueue(DONE_EX);
}
}

@Override
public void onError(Throwable t) {
ex = StatusUtils.fromThrowable(t);
queue.add(DONE_EX);
cancelled.complete(null);
root.setException(ex);
}

@Override
public void onCompleted() {
// Depends on gRPC calling onNext and onCompleted non-concurrently
cancelled.complete(null);
queue.add(DONE);
}
}

/**
* Cancels sending the stream to a client.
*
* @throws UnsupportedOperationException on a stream being uploaded from the client.
* <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
* received and to wait for the underlying transport to acknowledge cancellation.
*/
public void cancel(String message, Throwable exception) {
completed.completeExceptionally(
CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
if (cancellable != null) {
cancellable.cancel(message, exception);
} else {
if (cancellable == null) {
throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
"Instead, server should reject incoming messages.");
}
cancellable.cancel(message, exception);
// Do not mark the stream as completed, as gRPC may still be delivering messages.
}

StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

lidavidm
wants to merge
3
commits into

from

 Closed


ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut
#8010
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


3 commits


Select commit
Hold shift + click to select a range
c3a66cc

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…


lidavidm Aug 19, 2020
1e535b7

ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation


lidavidm Aug 31, 2020
80abecb

Fix TestLeak taking the full 60 second timeout


lidavidm Aug 31, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
 final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
 ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
 final ClientStreamListener writer = new PutObserver(
 descriptor, observer, stream.completed::isDone,
 descriptor, observer, stream.cancelled::isDone,
 () -> {
 try {
 stream.completed.get();
Expand Down

2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
 @Override
 public String toString() {
 String s = getClass().getName();
 return String.format("%s: %s", s, status);
 return String.format("%s: %s: %s", s, status.code(), status.description());
 }
}

53 changes: 24 additions & 29 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);

 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When the ackStream is completed, the FlightStream will be closed with it
 ackStream.setAutoCloseable(fs);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
 executors.submit(() -> {
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 try {
 producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
 } catch (Exception ex) {
 ackStream.onError(ex);
 } finally {
 // Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
 try {
 fs.close();
 } catch (Exception e) {
 handleExceptionWithMiddleware(e);
 }
 // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
 // We don't do this for other streams since the implementation may be asynchronous
 ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
 */
 private void handleExceptionWithMiddleware(Throwable t) {
 final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
 if (middleware == null) {
 if (middleware == null || middleware.isEmpty()) {
 logger.error("Uncaught exception in Flight method body", t);
 return;
 }
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

 /** Ensures that other resources are cleaned up when the service finishes its call. */
 private static class ExchangeListener extends GetListener {
 private final AutoCloseable resource;

 private AutoCloseable resource;
 private boolean closed = false;
 private Runnable onCancelHandler = null;

 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
 AutoCloseable resource) {
 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
 super(responseObserver, errorHandler);
 this.resource = resource;
 this.resource = null;
 super.setOnCancelHandler(() -> {
 try {
 if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
 }
 closed = true;
 try {
this.resource.close();
AutoCloseables.close(resource);
 } catch (Exception e) {
 throw CallStatus.INTERNAL
 .withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
 public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
 final ServerCallStreamObserver<ArrowMessage> responseObserver =
 (ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 final ExchangeListener listener = new ExchangeListener(
 responseObserver,
 this::handleExceptionWithMiddleware,
 () -> {
 // Force the stream to "complete" so it will close without incident. At this point, we don't care since
 // we are about to end the call. (Normally it will raise an error.)
 fs.completed.complete(null);
 fs.close();
 });
 this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 listener.resource = fs;
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down

121 changes: 85 additions & 36 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
 * An adaptor between protobuf streams and flight data streams.
 */
public class FlightStream implements AutoCloseable {

 // Use AutoCloseable sentinel objects to simplify logic in #close
 private final AutoCloseable DONE = () -> {
 };
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
 private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
 private final int pendingTarget;
 private final Requestor requestor;
 // The completion flags.
 // This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
 // data and closed the stream
 final CompletableFuture<Void> completed;
 // This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
 // we don't block forever trying to write to a server that has rejected a call.
 final CompletableFuture<Void> cancelled;

 private volatile int pending = 1;
 private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
 *
 * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
 * @param pendingTarget Target number of messages to receive.
 * @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
 * @param cancellable Used to cancel mid-stream requests.
 * @param requestor A callback to determine how many pending items there are.
 */
 public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
 Objects.requireNonNull(allocator);
 Objects.requireNonNull(requestor);
 this.allocator = allocator;
 this.pendingTarget = pendingTarget;
 this.cancellable = cancellable;
 this.requestor = requestor;
 this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
 this.completed = new CompletableFuture<>();
 this.cancelled = new CompletableFuture<>();
 }

 /**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
 /**
 * Closes the stream (freeing any existing resources).
 *
 * <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
 * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
 */
 public void close() throws Exception {
 final List<AutoCloseable> closeables = new ArrayList<>();
 // cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
 closeables.add(() -> {
 if (!completed.isDone() && cancellable != null) {
 cancel("Stream closed before end.", /* no exception to report */ null);
 Throwable suppressor = null;
 if (cancellable != null) {
 // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
 // On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
 // may never run https://github.com/grpc/grpc-java/issues/5882
 try {
 synchronized (cancellable) {
 if (!cancelled.isDone()) {
 // Only cancel if the call is not done on the gRPC side
 cancellable.cancel("Stream closed before end", /* no exception to report */null);
 }
 }
 // Drain the stream without the lock (as next() implicitly needs the lock)
 while (next()) { }
 } catch (FlightRuntimeException e) {
 suppressor = e;
 }
 });
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 // Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the
@rymurrrymurrSep 5, 2020

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification
@lidavidmlidavidmSep 5, 2020

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.
 // middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
 // by the lambda above.
 synchronized (completed) {
 try {
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 }
 if (suppressor != null) {
 AutoCloseables.close(suppressor, closeables);
 } else {
 AutoCloseables.close(closeables);
 }
 } finally {
 // The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
 // No-op if already complete
 completed.complete(null);
 }
 }

 AutoCloseables.close(closeables);
 // Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
 // No-op if already complete; do this after the check in the AutoCloseable lambda above
 completed.complete(null);
 }

 /**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
 super();
 }

 /** Helper to add an item to the queue under the appropriate lock. */
 private void enqueue(AutoCloseable message) {
 synchronized (completed) {
 if (completed.isDone()) {
 // The stream is already closed (RPC ended), discard the message
 AutoCloseables.closeNoChecked(message);
 } else {
 queue.add(message);
 }
 }
 }

 @Override
 public void onNext(ArrowMessage msg) {
 // Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
 // close().
 requestOutstanding();
 switch (msg.getMessageType()) {
 case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
 }
 break;
 }
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
 try {
 MetadataV4UnionChecker.checkRead(schema, metadataVersion);
 } catch (IOException e) {
 queue.add(DONE_EX);
 ex = e;
 enqueue(DONE_EX);
 break;
 }

 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 synchronized (completed) {
 if (!completed.isDone()) {
 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 root.set(fulfilledRoot);
 }
 }
 root.set(fulfilledRoot);
 break;
 }
 case RECORD_BATCH:
 queue.add(msg);
 break;
 case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
 break;
 case TENSOR:
 default:
 queue.add(DONE_EX);
 ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 enqueue(DONE_EX);
 }
 }

 @Override
 public void onError(Throwable t) {
 ex = StatusUtils.fromThrowable(t);
 queue.add(DONE_EX);
 cancelled.complete(null);
 root.setException(ex);
 }

 @Override
 public void onCompleted() {
 // Depends on gRPC calling onNext and onCompleted non-concurrently
 cancelled.complete(null);
 queue.add(DONE);
 }
 }

 /**
 * Cancels sending the stream to a client.
 *
 * @throws UnsupportedOperationException on a stream being uploaded from the client.
 * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
 * received and to wait for the underlying transport to acknowledge cancellation.
 */
 public void cancel(String message, Throwable exception) {
 completed.completeExceptionally(
 CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
 if (cancellable != null) {
 cancellable.cancel(message, exception);
 } else {
 if (cancellable == null) {
 throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
 "Instead, server should reject incoming messages.");
 }
 cancellable.cancel(message, exception);
 // Do not mark the stream as completed, as gRPC may still be delivering messages.
 }

 StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

lidavidm
wants to merge
3
commits into

from

 Closed


ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut
#8010
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


3 commits


Select commit
Hold shift + click to select a range
c3a66cc

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…


lidavidm Aug 19, 2020
1e535b7

ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation


lidavidm Aug 31, 2020
80abecb

Fix TestLeak taking the full 60 second timeout


lidavidm Aug 31, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
 final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
 ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
 final ClientStreamListener writer = new PutObserver(
 descriptor, observer, stream.completed::isDone,
 descriptor, observer, stream.cancelled::isDone,
 () -> {
 try {
 stream.completed.get();
Expand Down

2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
 @Override
 public String toString() {
 String s = getClass().getName();
 return String.format("%s: %s", s, status);
 return String.format("%s: %s: %s", s, status.code(), status.description());
 }
}

53 changes: 24 additions & 29 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);

 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When the ackStream is completed, the FlightStream will be closed with it
 ackStream.setAutoCloseable(fs);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
 executors.submit(() -> {
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 try {
 producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
 } catch (Exception ex) {
 ackStream.onError(ex);
 } finally {
 // Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
 try {
 fs.close();
 } catch (Exception e) {
 handleExceptionWithMiddleware(e);
 }
 // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
 // We don't do this for other streams since the implementation may be asynchronous
 ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
 */
 private void handleExceptionWithMiddleware(Throwable t) {
 final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
 if (middleware == null) {
 if (middleware == null || middleware.isEmpty()) {
 logger.error("Uncaught exception in Flight method body", t);
 return;
 }
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

 /** Ensures that other resources are cleaned up when the service finishes its call. */
 private static class ExchangeListener extends GetListener {
 private final AutoCloseable resource;

 private AutoCloseable resource;
 private boolean closed = false;
 private Runnable onCancelHandler = null;

 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
 AutoCloseable resource) {
 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
 super(responseObserver, errorHandler);
 this.resource = resource;
 this.resource = null;
 super.setOnCancelHandler(() -> {
 try {
 if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
 }
 closed = true;
 try {
this.resource.close();
AutoCloseables.close(resource);
 } catch (Exception e) {
 throw CallStatus.INTERNAL
 .withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
 public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
 final ServerCallStreamObserver<ArrowMessage> responseObserver =
 (ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 final ExchangeListener listener = new ExchangeListener(
 responseObserver,
 this::handleExceptionWithMiddleware,
 () -> {
 // Force the stream to "complete" so it will close without incident. At this point, we don't care since
 // we are about to end the call. (Normally it will raise an error.)
 fs.completed.complete(null);
 fs.close();
 });
 this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 listener.resource = fs;
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down

121 changes: 85 additions & 36 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
 * An adaptor between protobuf streams and flight data streams.
 */
public class FlightStream implements AutoCloseable {

 // Use AutoCloseable sentinel objects to simplify logic in #close
 private final AutoCloseable DONE = () -> {
 };
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
 private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
 private final int pendingTarget;
 private final Requestor requestor;
 // The completion flags.
 // This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
 // data and closed the stream
 final CompletableFuture<Void> completed;
 // This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
 // we don't block forever trying to write to a server that has rejected a call.
 final CompletableFuture<Void> cancelled;

 private volatile int pending = 1;
 private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
 *
 * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
 * @param pendingTarget Target number of messages to receive.
 * @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
 * @param cancellable Used to cancel mid-stream requests.
 * @param requestor A callback to determine how many pending items there are.
 */
 public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
 Objects.requireNonNull(allocator);
 Objects.requireNonNull(requestor);
 this.allocator = allocator;
 this.pendingTarget = pendingTarget;
 this.cancellable = cancellable;
 this.requestor = requestor;
 this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
 this.completed = new CompletableFuture<>();
 this.cancelled = new CompletableFuture<>();
 }

 /**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
 /**
 * Closes the stream (freeing any existing resources).
 *
 * <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
 * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
 */
 public void close() throws Exception {
 final List<AutoCloseable> closeables = new ArrayList<>();
 // cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
 closeables.add(() -> {
 if (!completed.isDone() && cancellable != null) {
 cancel("Stream closed before end.", /* no exception to report */ null);
 Throwable suppressor = null;
 if (cancellable != null) {
 // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
 // On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
 // may never run https://github.com/grpc/grpc-java/issues/5882
 try {
 synchronized (cancellable) {
 if (!cancelled.isDone()) {
 // Only cancel if the call is not done on the gRPC side
 cancellable.cancel("Stream closed before end", /* no exception to report */null);
 }
 }
 // Drain the stream without the lock (as next() implicitly needs the lock)
 while (next()) { }
 } catch (FlightRuntimeException e) {
 suppressor = e;
 }
 });
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 // Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the
@rymurrrymurrSep 5, 2020

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification
@lidavidmlidavidmSep 5, 2020

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.
 // middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
 // by the lambda above.
 synchronized (completed) {
 try {
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 }
 if (suppressor != null) {
 AutoCloseables.close(suppressor, closeables);
 } else {
 AutoCloseables.close(closeables);
 }
 } finally {
 // The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
 // No-op if already complete
 completed.complete(null);
 }
 }

 AutoCloseables.close(closeables);
 // Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
 // No-op if already complete; do this after the check in the AutoCloseable lambda above
 completed.complete(null);
 }

 /**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
 super();
 }

 /** Helper to add an item to the queue under the appropriate lock. */
 private void enqueue(AutoCloseable message) {
 synchronized (completed) {
 if (completed.isDone()) {
 // The stream is already closed (RPC ended), discard the message
 AutoCloseables.closeNoChecked(message);
 } else {
 queue.add(message);
 }
 }
 }

 @Override
 public void onNext(ArrowMessage msg) {
 // Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
 // close().
 requestOutstanding();
 switch (msg.getMessageType()) {
 case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
 }
 break;
 }
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
 try {
 MetadataV4UnionChecker.checkRead(schema, metadataVersion);
 } catch (IOException e) {
 queue.add(DONE_EX);
 ex = e;
 enqueue(DONE_EX);
 break;
 }

 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 synchronized (completed) {
 if (!completed.isDone()) {
 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 root.set(fulfilledRoot);
 }
 }
 root.set(fulfilledRoot);
 break;
 }
 case RECORD_BATCH:
 queue.add(msg);
 break;
 case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
 break;
 case TENSOR:
 default:
 queue.add(DONE_EX);
 ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 enqueue(DONE_EX);
 }
 }

 @Override
 public void onError(Throwable t) {
 ex = StatusUtils.fromThrowable(t);
 queue.add(DONE_EX);
 cancelled.complete(null);
 root.setException(ex);
 }

 @Override
 public void onCompleted() {
 // Depends on gRPC calling onNext and onCompleted non-concurrently
 cancelled.complete(null);
 queue.add(DONE);
 }
 }

 /**
 * Cancels sending the stream to a client.
 *
 * @throws UnsupportedOperationException on a stream being uploaded from the client.
 * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
 * received and to wait for the underlying transport to acknowledge cancellation.
 */
 public void cancel(String message, Throwable exception) {
 completed.completeExceptionally(
 CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
 if (cancellable != null) {
 cancellable.cancel(message, exception);
 } else {
 if (cancellable == null) {
 throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
 "Instead, server should reject incoming messages.");
 }
 cancellable.cancel(message, exception);
 // Do not mark the stream as completed, as gRPC may still be delivering messages.
 }

 StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

lidavidm
wants to merge
3
commits into

from

 Closed


ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut
#8010
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


3 commits


Select commit
Hold shift + click to select a range
c3a66cc

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…


lidavidm Aug 19, 2020
1e535b7

ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation


lidavidm Aug 31, 2020
80abecb

Fix TestLeak taking the full 60 second timeout


lidavidm Aug 31, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
 final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
 ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
 final ClientStreamListener writer = new PutObserver(
 descriptor, observer, stream.completed::isDone,
 descriptor, observer, stream.cancelled::isDone,
 () -> {
 try {
 stream.completed.get();
Expand Down

2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
 @Override
 public String toString() {
 String s = getClass().getName();
 return String.format("%s: %s", s, status);
 return String.format("%s: %s: %s", s, status.code(), status.description());
 }
}

53 changes: 24 additions & 29 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);

 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When the ackStream is completed, the FlightStream will be closed with it
 ackStream.setAutoCloseable(fs);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
 executors.submit(() -> {
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 try {
 producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
 } catch (Exception ex) {
 ackStream.onError(ex);
 } finally {
 // Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
 try {
 fs.close();
 } catch (Exception e) {
 handleExceptionWithMiddleware(e);
 }
 // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
 // We don't do this for other streams since the implementation may be asynchronous
 ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
 */
 private void handleExceptionWithMiddleware(Throwable t) {
 final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
 if (middleware == null) {
 if (middleware == null || middleware.isEmpty()) {
 logger.error("Uncaught exception in Flight method body", t);
 return;
 }
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

 /** Ensures that other resources are cleaned up when the service finishes its call. */
 private static class ExchangeListener extends GetListener {
 private final AutoCloseable resource;

 private AutoCloseable resource;
 private boolean closed = false;
 private Runnable onCancelHandler = null;

 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
 AutoCloseable resource) {
 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
 super(responseObserver, errorHandler);
 this.resource = resource;
 this.resource = null;
 super.setOnCancelHandler(() -> {
 try {
 if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
 }
 closed = true;
 try {
this.resource.close();
AutoCloseables.close(resource);
 } catch (Exception e) {
 throw CallStatus.INTERNAL
 .withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
 public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
 final ServerCallStreamObserver<ArrowMessage> responseObserver =
 (ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 final ExchangeListener listener = new ExchangeListener(
 responseObserver,
 this::handleExceptionWithMiddleware,
 () -> {
 // Force the stream to "complete" so it will close without incident. At this point, we don't care since
 // we are about to end the call. (Normally it will raise an error.)
 fs.completed.complete(null);
 fs.close();
 });
 this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 listener.resource = fs;
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down

121 changes: 85 additions & 36 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
 * An adaptor between protobuf streams and flight data streams.
 */
public class FlightStream implements AutoCloseable {

 // Use AutoCloseable sentinel objects to simplify logic in #close
 private final AutoCloseable DONE = () -> {
 };
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
 private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
 private final int pendingTarget;
 private final Requestor requestor;
 // The completion flags.
 // This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
 // data and closed the stream
 final CompletableFuture<Void> completed;
 // This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
 // we don't block forever trying to write to a server that has rejected a call.
 final CompletableFuture<Void> cancelled;

 private volatile int pending = 1;
 private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
 *
 * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
 * @param pendingTarget Target number of messages to receive.
 * @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
 * @param cancellable Used to cancel mid-stream requests.
 * @param requestor A callback to determine how many pending items there are.
 */
 public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
 Objects.requireNonNull(allocator);
 Objects.requireNonNull(requestor);
 this.allocator = allocator;
 this.pendingTarget = pendingTarget;
 this.cancellable = cancellable;
 this.requestor = requestor;
 this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
 this.completed = new CompletableFuture<>();
 this.cancelled = new CompletableFuture<>();
 }

 /**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
 /**
 * Closes the stream (freeing any existing resources).
 *
 * <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
 * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
 */
 public void close() throws Exception {
 final List<AutoCloseable> closeables = new ArrayList<>();
 // cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
 closeables.add(() -> {
 if (!completed.isDone() && cancellable != null) {
 cancel("Stream closed before end.", /* no exception to report */ null);
 Throwable suppressor = null;
 if (cancellable != null) {
 // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
 // On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
 // may never run https://github.com/grpc/grpc-java/issues/5882
 try {
 synchronized (cancellable) {
 if (!cancelled.isDone()) {
 // Only cancel if the call is not done on the gRPC side
 cancellable.cancel("Stream closed before end", /* no exception to report */null);
 }
 }
 // Drain the stream without the lock (as next() implicitly needs the lock)
 while (next()) { }
 } catch (FlightRuntimeException e) {
 suppressor = e;
 }
 });
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 // Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the
@rymurrrymurrSep 5, 2020

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification
@lidavidmlidavidmSep 5, 2020

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.
 // middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
 // by the lambda above.
 synchronized (completed) {
 try {
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 }
 if (suppressor != null) {
 AutoCloseables.close(suppressor, closeables);
 } else {
 AutoCloseables.close(closeables);
 }
 } finally {
 // The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
 // No-op if already complete
 completed.complete(null);
 }
 }

 AutoCloseables.close(closeables);
 // Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
 // No-op if already complete; do this after the check in the AutoCloseable lambda above
 completed.complete(null);
 }

 /**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
 super();
 }

 /** Helper to add an item to the queue under the appropriate lock. */
 private void enqueue(AutoCloseable message) {
 synchronized (completed) {
 if (completed.isDone()) {
 // The stream is already closed (RPC ended), discard the message
 AutoCloseables.closeNoChecked(message);
 } else {
 queue.add(message);
 }
 }
 }

 @Override
 public void onNext(ArrowMessage msg) {
 // Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
 // close().
 requestOutstanding();
 switch (msg.getMessageType()) {
 case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
 }
 break;
 }
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
 try {
 MetadataV4UnionChecker.checkRead(schema, metadataVersion);
 } catch (IOException e) {
 queue.add(DONE_EX);
 ex = e;
 enqueue(DONE_EX);
 break;
 }

 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 synchronized (completed) {
 if (!completed.isDone()) {
 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 root.set(fulfilledRoot);
 }
 }
 root.set(fulfilledRoot);
 break;
 }
 case RECORD_BATCH:
 queue.add(msg);
 break;
 case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
 break;
 case TENSOR:
 default:
 queue.add(DONE_EX);
 ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 enqueue(DONE_EX);
 }
 }

 @Override
 public void onError(Throwable t) {
 ex = StatusUtils.fromThrowable(t);
 queue.add(DONE_EX);
 cancelled.complete(null);
 root.setException(ex);
 }

 @Override
 public void onCompleted() {
 // Depends on gRPC calling onNext and onCompleted non-concurrently
 cancelled.complete(null);
 queue.add(DONE);
 }
 }

 /**
 * Cancels sending the stream to a client.
 *
 * @throws UnsupportedOperationException on a stream being uploaded from the client.
 * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
 * received and to wait for the underlying transport to acknowledge cancellation.
 */
 public void cancel(String message, Throwable exception) {
 completed.completeExceptionally(
 CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
 if (cancellable != null) {
 cancellable.cancel(message, exception);
 } else {
 if (cancellable == null) {
 throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
 "Instead, server should reject incoming messages.");
 }
 cancellable.cancel(message, exception);
 // Do not mark the stream as completed, as gRPC may still be delivering messages.
 }

 StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

lidavidm
wants to merge
3
commits into

from

 Closed


ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut
#8010
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


3 commits


Select commit
Hold shift + click to select a range
c3a66cc

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…


lidavidm Aug 19, 2020
1e535b7

ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation


lidavidm Aug 31, 2020
80abecb

Fix TestLeak taking the full 60 second timeout


lidavidm Aug 31, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
 final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
 ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
 final ClientStreamListener writer = new PutObserver(
 descriptor, observer, stream.completed::isDone,
 descriptor, observer, stream.cancelled::isDone,
 () -> {
 try {
 stream.completed.get();
Expand Down

2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
 @Override
 public String toString() {
 String s = getClass().getName();
 return String.format("%s: %s", s, status);
 return String.format("%s: %s: %s", s, status.code(), status.description());
 }
}

53 changes: 24 additions & 29 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);

 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When the ackStream is completed, the FlightStream will be closed with it
 ackStream.setAutoCloseable(fs);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
 executors.submit(() -> {
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 try {
 producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
 } catch (Exception ex) {
 ackStream.onError(ex);
 } finally {
 // Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
 try {
 fs.close();
 } catch (Exception e) {
 handleExceptionWithMiddleware(e);
 }
 // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
 // We don't do this for other streams since the implementation may be asynchronous
 ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
 */
 private void handleExceptionWithMiddleware(Throwable t) {
 final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
 if (middleware == null) {
 if (middleware == null || middleware.isEmpty()) {
 logger.error("Uncaught exception in Flight method body", t);
 return;
 }
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

 /** Ensures that other resources are cleaned up when the service finishes its call. */
 private static class ExchangeListener extends GetListener {
 private final AutoCloseable resource;

 private AutoCloseable resource;
 private boolean closed = false;
 private Runnable onCancelHandler = null;

 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
 AutoCloseable resource) {
 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
 super(responseObserver, errorHandler);
 this.resource = resource;
 this.resource = null;
 super.setOnCancelHandler(() -> {
 try {
 if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
 }
 closed = true;
 try {
this.resource.close();
AutoCloseables.close(resource);
 } catch (Exception e) {
 throw CallStatus.INTERNAL
 .withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
 public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
 final ServerCallStreamObserver<ArrowMessage> responseObserver =
 (ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 final ExchangeListener listener = new ExchangeListener(
 responseObserver,
 this::handleExceptionWithMiddleware,
 () -> {
 // Force the stream to "complete" so it will close without incident. At this point, we don't care since
 // we are about to end the call. (Normally it will raise an error.)
 fs.completed.complete(null);
 fs.close();
 });
 this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 listener.resource = fs;
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down

121 changes: 85 additions & 36 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
 * An adaptor between protobuf streams and flight data streams.
 */
public class FlightStream implements AutoCloseable {

 // Use AutoCloseable sentinel objects to simplify logic in #close
 private final AutoCloseable DONE = () -> {
 };
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
 private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
 private final int pendingTarget;
 private final Requestor requestor;
 // The completion flags.
 // This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
 // data and closed the stream
 final CompletableFuture<Void> completed;
 // This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
 // we don't block forever trying to write to a server that has rejected a call.
 final CompletableFuture<Void> cancelled;

 private volatile int pending = 1;
 private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
 *
 * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
 * @param pendingTarget Target number of messages to receive.
 * @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
 * @param cancellable Used to cancel mid-stream requests.
 * @param requestor A callback to determine how many pending items there are.
 */
 public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
 Objects.requireNonNull(allocator);
 Objects.requireNonNull(requestor);
 this.allocator = allocator;
 this.pendingTarget = pendingTarget;
 this.cancellable = cancellable;
 this.requestor = requestor;
 this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
 this.completed = new CompletableFuture<>();
 this.cancelled = new CompletableFuture<>();
 }

 /**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
 /**
 * Closes the stream (freeing any existing resources).
 *
 * <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
 * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
 */
 public void close() throws Exception {
 final List<AutoCloseable> closeables = new ArrayList<>();
 // cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
 closeables.add(() -> {
 if (!completed.isDone() && cancellable != null) {
 cancel("Stream closed before end.", /* no exception to report */ null);
 Throwable suppressor = null;
 if (cancellable != null) {
 // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
 // On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
 // may never run https://github.com/grpc/grpc-java/issues/5882
 try {
 synchronized (cancellable) {
 if (!cancelled.isDone()) {
 // Only cancel if the call is not done on the gRPC side
 cancellable.cancel("Stream closed before end", /* no exception to report */null);
 }
 }
 // Drain the stream without the lock (as next() implicitly needs the lock)
 while (next()) { }
 } catch (FlightRuntimeException e) {
 suppressor = e;
 }
 });
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 // Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the
@rymurrrymurrSep 5, 2020

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification
@lidavidmlidavidmSep 5, 2020

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.
 // middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
 // by the lambda above.
 synchronized (completed) {
 try {
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 }
 if (suppressor != null) {
 AutoCloseables.close(suppressor, closeables);
 } else {
 AutoCloseables.close(closeables);
 }
 } finally {
 // The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
 // No-op if already complete
 completed.complete(null);
 }
 }

 AutoCloseables.close(closeables);
 // Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
 // No-op if already complete; do this after the check in the AutoCloseable lambda above
 completed.complete(null);
 }

 /**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
 super();
 }

 /** Helper to add an item to the queue under the appropriate lock. */
 private void enqueue(AutoCloseable message) {
 synchronized (completed) {
 if (completed.isDone()) {
 // The stream is already closed (RPC ended), discard the message
 AutoCloseables.closeNoChecked(message);
 } else {
 queue.add(message);
 }
 }
 }

 @Override
 public void onNext(ArrowMessage msg) {
 // Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
 // close().
 requestOutstanding();
 switch (msg.getMessageType()) {
 case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
 }
 break;
 }
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
 try {
 MetadataV4UnionChecker.checkRead(schema, metadataVersion);
 } catch (IOException e) {
 queue.add(DONE_EX);
 ex = e;
 enqueue(DONE_EX);
 break;
 }

 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 synchronized (completed) {
 if (!completed.isDone()) {
 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 root.set(fulfilledRoot);
 }
 }
 root.set(fulfilledRoot);
 break;
 }
 case RECORD_BATCH:
 queue.add(msg);
 break;
 case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
 break;
 case TENSOR:
 default:
 queue.add(DONE_EX);
 ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 enqueue(DONE_EX);
 }
 }

 @Override
 public void onError(Throwable t) {
 ex = StatusUtils.fromThrowable(t);
 queue.add(DONE_EX);
 cancelled.complete(null);
 root.setException(ex);
 }

 @Override
 public void onCompleted() {
 // Depends on gRPC calling onNext and onCompleted non-concurrently
 cancelled.complete(null);
 queue.add(DONE);
 }
 }

 /**
 * Cancels sending the stream to a client.
 *
 * @throws UnsupportedOperationException on a stream being uploaded from the client.
 * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
 * received and to wait for the underlying transport to acknowledge cancellation.
 */
 public void cancel(String message, Throwable exception) {
 completed.completeExceptionally(
 CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
 if (cancellable != null) {
 cancellable.cancel(message, exception);
 } else {
 if (cancellable == null) {
 throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
 "Instead, server should reject incoming messages.");
 }
 cancellable.cancel(message, exception);
 // Do not mark the stream as completed, as gRPC may still be delivering messages.
 }

 StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

lidavidm
wants to merge
3
commits into

from

 Closed


ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut
#8010
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


3 commits


Select commit
Hold shift + click to select a range
c3a66cc

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…


lidavidm Aug 19, 2020
1e535b7

ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation


lidavidm Aug 31, 2020
80abecb

Fix TestLeak taking the full 60 second timeout


lidavidm Aug 31, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
 final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
 ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
 final ClientStreamListener writer = new PutObserver(
 descriptor, observer, stream.completed::isDone,
 descriptor, observer, stream.cancelled::isDone,
 () -> {
 try {
 stream.completed.get();
Expand Down

2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
 @Override
 public String toString() {
 String s = getClass().getName();
 return String.format("%s: %s", s, status);
 return String.format("%s: %s: %s", s, status.code(), status.description());
 }
}

53 changes: 24 additions & 29 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);

 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When the ackStream is completed, the FlightStream will be closed with it
 ackStream.setAutoCloseable(fs);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
 executors.submit(() -> {
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 try {
 producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
 } catch (Exception ex) {
 ackStream.onError(ex);
 } finally {
 // Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
 try {
 fs.close();
 } catch (Exception e) {
 handleExceptionWithMiddleware(e);
 }
 // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
 // We don't do this for other streams since the implementation may be asynchronous
 ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
 */
 private void handleExceptionWithMiddleware(Throwable t) {
 final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
 if (middleware == null) {
 if (middleware == null || middleware.isEmpty()) {
 logger.error("Uncaught exception in Flight method body", t);
 return;
 }
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

 /** Ensures that other resources are cleaned up when the service finishes its call. */
 private static class ExchangeListener extends GetListener {
 private final AutoCloseable resource;

 private AutoCloseable resource;
 private boolean closed = false;
 private Runnable onCancelHandler = null;

 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
 AutoCloseable resource) {
 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
 super(responseObserver, errorHandler);
 this.resource = resource;
 this.resource = null;
 super.setOnCancelHandler(() -> {
 try {
 if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
 }
 closed = true;
 try {
this.resource.close();
AutoCloseables.close(resource);
 } catch (Exception e) {
 throw CallStatus.INTERNAL
 .withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
 public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
 final ServerCallStreamObserver<ArrowMessage> responseObserver =
 (ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 final ExchangeListener listener = new ExchangeListener(
 responseObserver,
 this::handleExceptionWithMiddleware,
 () -> {
 // Force the stream to "complete" so it will close without incident. At this point, we don't care since
 // we are about to end the call. (Normally it will raise an error.)
 fs.completed.complete(null);
 fs.close();
 });
 this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 listener.resource = fs;
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down

121 changes: 85 additions & 36 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
 * An adaptor between protobuf streams and flight data streams.
 */
public class FlightStream implements AutoCloseable {

 // Use AutoCloseable sentinel objects to simplify logic in #close
 private final AutoCloseable DONE = () -> {
 };
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
 private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
 private final int pendingTarget;
 private final Requestor requestor;
 // The completion flags.
 // This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
 // data and closed the stream
 final CompletableFuture<Void> completed;
 // This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
 // we don't block forever trying to write to a server that has rejected a call.
 final CompletableFuture<Void> cancelled;

 private volatile int pending = 1;
 private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
 *
 * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
 * @param pendingTarget Target number of messages to receive.
 * @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
 * @param cancellable Used to cancel mid-stream requests.
 * @param requestor A callback to determine how many pending items there are.
 */
 public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
 Objects.requireNonNull(allocator);
 Objects.requireNonNull(requestor);
 this.allocator = allocator;
 this.pendingTarget = pendingTarget;
 this.cancellable = cancellable;
 this.requestor = requestor;
 this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
 this.completed = new CompletableFuture<>();
 this.cancelled = new CompletableFuture<>();
 }

 /**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
 /**
 * Closes the stream (freeing any existing resources).
 *
 * <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
 * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
 */
 public void close() throws Exception {
 final List<AutoCloseable> closeables = new ArrayList<>();
 // cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
 closeables.add(() -> {
 if (!completed.isDone() && cancellable != null) {
 cancel("Stream closed before end.", /* no exception to report */ null);
 Throwable suppressor = null;
 if (cancellable != null) {
 // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
 // On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
 // may never run https://github.com/grpc/grpc-java/issues/5882
 try {
 synchronized (cancellable) {
 if (!cancelled.isDone()) {
 // Only cancel if the call is not done on the gRPC side
 cancellable.cancel("Stream closed before end", /* no exception to report */null);
 }
 }
 // Drain the stream without the lock (as next() implicitly needs the lock)
 while (next()) { }
 } catch (FlightRuntimeException e) {
 suppressor = e;
 }
 });
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 // Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the
@rymurrrymurrSep 5, 2020

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification
@lidavidmlidavidmSep 5, 2020

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.
 // middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
 // by the lambda above.
 synchronized (completed) {
 try {
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 }
 if (suppressor != null) {
 AutoCloseables.close(suppressor, closeables);
 } else {
 AutoCloseables.close(closeables);
 }
 } finally {
 // The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
 // No-op if already complete
 completed.complete(null);
 }
 }

 AutoCloseables.close(closeables);
 // Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
 // No-op if already complete; do this after the check in the AutoCloseable lambda above
 completed.complete(null);
 }

 /**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
 super();
 }

 /** Helper to add an item to the queue under the appropriate lock. */
 private void enqueue(AutoCloseable message) {
 synchronized (completed) {
 if (completed.isDone()) {
 // The stream is already closed (RPC ended), discard the message
 AutoCloseables.closeNoChecked(message);
 } else {
 queue.add(message);
 }
 }
 }

 @Override
 public void onNext(ArrowMessage msg) {
 // Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
 // close().
 requestOutstanding();
 switch (msg.getMessageType()) {
 case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
 }
 break;
 }
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
 try {
 MetadataV4UnionChecker.checkRead(schema, metadataVersion);
 } catch (IOException e) {
 queue.add(DONE_EX);
 ex = e;
 enqueue(DONE_EX);
 break;
 }

 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 synchronized (completed) {
 if (!completed.isDone()) {
 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 root.set(fulfilledRoot);
 }
 }
 root.set(fulfilledRoot);
 break;
 }
 case RECORD_BATCH:
 queue.add(msg);
 break;
 case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
 break;
 case TENSOR:
 default:
 queue.add(DONE_EX);
 ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 enqueue(DONE_EX);
 }
 }

 @Override
 public void onError(Throwable t) {
 ex = StatusUtils.fromThrowable(t);
 queue.add(DONE_EX);
 cancelled.complete(null);
 root.setException(ex);
 }

 @Override
 public void onCompleted() {
 // Depends on gRPC calling onNext and onCompleted non-concurrently
 cancelled.complete(null);
 queue.add(DONE);
 }
 }

 /**
 * Cancels sending the stream to a client.
 *
 * @throws UnsupportedOperationException on a stream being uploaded from the client.
 * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
 * received and to wait for the underlying transport to acknowledge cancellation.
 */
 public void cancel(String message, Throwable exception) {
 completed.completeExceptionally(
 CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
 if (cancellable != null) {
 cancellable.cancel(message, exception);
 } else {
 if (cancellable == null) {
 throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
 "Instead, server should reject incoming messages.");
 }
 cancellable.cancel(message, exception);
 // Do not mark the stream as completed, as gRPC may still be delivering messages.
 }

 StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

lidavidm
wants to merge
3
commits into

from

 Closed


ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut
#8010
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


3 commits


Select commit
Hold shift + click to select a range
c3a66cc

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…


lidavidm Aug 19, 2020
1e535b7

ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation


lidavidm Aug 31, 2020
80abecb

Fix TestLeak taking the full 60 second timeout


lidavidm Aug 31, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
 final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
 ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
 final ClientStreamListener writer = new PutObserver(
 descriptor, observer, stream.completed::isDone,
 descriptor, observer, stream.cancelled::isDone,
 () -> {
 try {
 stream.completed.get();
Expand Down

2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
 @Override
 public String toString() {
 String s = getClass().getName();
 return String.format("%s: %s", s, status);
 return String.format("%s: %s: %s", s, status.code(), status.description());
 }
}

53 changes: 24 additions & 29 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);

 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When the ackStream is completed, the FlightStream will be closed with it
 ackStream.setAutoCloseable(fs);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
 executors.submit(() -> {
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 try {
 producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
 } catch (Exception ex) {
 ackStream.onError(ex);
 } finally {
 // Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
 try {
 fs.close();
 } catch (Exception e) {
 handleExceptionWithMiddleware(e);
 }
 // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
 // We don't do this for other streams since the implementation may be asynchronous
 ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
 */
 private void handleExceptionWithMiddleware(Throwable t) {
 final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
 if (middleware == null) {
 if (middleware == null || middleware.isEmpty()) {
 logger.error("Uncaught exception in Flight method body", t);
 return;
 }
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

 /** Ensures that other resources are cleaned up when the service finishes its call. */
 private static class ExchangeListener extends GetListener {
 private final AutoCloseable resource;

 private AutoCloseable resource;
 private boolean closed = false;
 private Runnable onCancelHandler = null;

 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
 AutoCloseable resource) {
 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
 super(responseObserver, errorHandler);
 this.resource = resource;
 this.resource = null;
 super.setOnCancelHandler(() -> {
 try {
 if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
 }
 closed = true;
 try {
this.resource.close();
AutoCloseables.close(resource);
 } catch (Exception e) {
 throw CallStatus.INTERNAL
 .withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
 public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
 final ServerCallStreamObserver<ArrowMessage> responseObserver =
 (ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 final ExchangeListener listener = new ExchangeListener(
 responseObserver,
 this::handleExceptionWithMiddleware,
 () -> {
 // Force the stream to "complete" so it will close without incident. At this point, we don't care since
 // we are about to end the call. (Normally it will raise an error.)
 fs.completed.complete(null);
 fs.close();
 });
 this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 listener.resource = fs;
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down

121 changes: 85 additions & 36 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
 * An adaptor between protobuf streams and flight data streams.
 */
public class FlightStream implements AutoCloseable {

 // Use AutoCloseable sentinel objects to simplify logic in #close
 private final AutoCloseable DONE = () -> {
 };
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
 private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
 private final int pendingTarget;
 private final Requestor requestor;
 // The completion flags.
 // This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
 // data and closed the stream
 final CompletableFuture<Void> completed;
 // This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
 // we don't block forever trying to write to a server that has rejected a call.
 final CompletableFuture<Void> cancelled;

 private volatile int pending = 1;
 private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
 *
 * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
 * @param pendingTarget Target number of messages to receive.
 * @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
 * @param cancellable Used to cancel mid-stream requests.
 * @param requestor A callback to determine how many pending items there are.
 */
 public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
 Objects.requireNonNull(allocator);
 Objects.requireNonNull(requestor);
 this.allocator = allocator;
 this.pendingTarget = pendingTarget;
 this.cancellable = cancellable;
 this.requestor = requestor;
 this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
 this.completed = new CompletableFuture<>();
 this.cancelled = new CompletableFuture<>();
 }

 /**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
 /**
 * Closes the stream (freeing any existing resources).
 *
 * <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
 * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
 */
 public void close() throws Exception {
 final List<AutoCloseable> closeables = new ArrayList<>();
 // cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
 closeables.add(() -> {
 if (!completed.isDone() && cancellable != null) {
 cancel("Stream closed before end.", /* no exception to report */ null);
 Throwable suppressor = null;
 if (cancellable != null) {
 // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
 // On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
 // may never run https://github.com/grpc/grpc-java/issues/5882
 try {
 synchronized (cancellable) {
 if (!cancelled.isDone()) {
 // Only cancel if the call is not done on the gRPC side
 cancellable.cancel("Stream closed before end", /* no exception to report */null);
 }
 }
 // Drain the stream without the lock (as next() implicitly needs the lock)
 while (next()) { }
 } catch (FlightRuntimeException e) {
 suppressor = e;
 }
 });
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 // Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the
@rymurrrymurrSep 5, 2020

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification
@lidavidmlidavidmSep 5, 2020

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.
 // middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
 // by the lambda above.
 synchronized (completed) {
 try {
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 }
 if (suppressor != null) {
 AutoCloseables.close(suppressor, closeables);
 } else {
 AutoCloseables.close(closeables);
 }
 } finally {
 // The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
 // No-op if already complete
 completed.complete(null);
 }
 }

 AutoCloseables.close(closeables);
 // Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
 // No-op if already complete; do this after the check in the AutoCloseable lambda above
 completed.complete(null);
 }

 /**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
 super();
 }

 /** Helper to add an item to the queue under the appropriate lock. */
 private void enqueue(AutoCloseable message) {
 synchronized (completed) {
 if (completed.isDone()) {
 // The stream is already closed (RPC ended), discard the message
 AutoCloseables.closeNoChecked(message);
 } else {
 queue.add(message);
 }
 }
 }

 @Override
 public void onNext(ArrowMessage msg) {
 // Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
 // close().
 requestOutstanding();
 switch (msg.getMessageType()) {
 case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
 }
 break;
 }
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
 try {
 MetadataV4UnionChecker.checkRead(schema, metadataVersion);
 } catch (IOException e) {
 queue.add(DONE_EX);
 ex = e;
 enqueue(DONE_EX);
 break;
 }

 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 synchronized (completed) {
 if (!completed.isDone()) {
 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 root.set(fulfilledRoot);
 }
 }
 root.set(fulfilledRoot);
 break;
 }
 case RECORD_BATCH:
 queue.add(msg);
 break;
 case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
 break;
 case TENSOR:
 default:
 queue.add(DONE_EX);
 ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 enqueue(DONE_EX);
 }
 }

 @Override
 public void onError(Throwable t) {
 ex = StatusUtils.fromThrowable(t);
 queue.add(DONE_EX);
 cancelled.complete(null);
 root.setException(ex);
 }

 @Override
 public void onCompleted() {
 // Depends on gRPC calling onNext and onCompleted non-concurrently
 cancelled.complete(null);
 queue.add(DONE);
 }
 }

 /**
 * Cancels sending the stream to a client.
 *
 * @throws UnsupportedOperationException on a stream being uploaded from the client.
 * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
 * received and to wait for the underlying transport to acknowledge cancellation.
 */
 public void cancel(String message, Throwable exception) {
 completed.completeExceptionally(
 CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
 if (cancellable != null) {
 cancellable.cancel(message, exception);
 } else {
 if (cancellable == null) {
 throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
 "Instead, server should reject incoming messages.");
 }
 cancellable.cancel(message, exception);
 // Do not mark the stream as completed, as gRPC may still be delivering messages.
 }

 StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut#8010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

lidavidm
wants to merge
3
commits into

from

 Closed


ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut
#8010
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


3 commits


Select commit
Hold shift + click to select a range
c3a66cc

ARROW-9587: [FlightRPC][Java] clean up FlightStream/DoPut memory hand…


lidavidm Aug 19, 2020
1e535b7

ARROW-9587: [FlightRPC][Java] fix FlightStream with cancellation


lidavidm Aug 31, 2020
80abecb

Fix TestLeak taking the full 60 second timeout


lidavidm Aug 31, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -321,7 +321,7 @@ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption..
 final ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
 ClientCalls.asyncBidiStreamingCall(call, stream.asObserver());
 final ClientStreamListener writer = new PutObserver(
 descriptor, observer, stream.completed::isDone,
 descriptor, observer, stream.cancelled::isDone,
 () -> {
 try {
 stream.completed.get();
Expand Down

2 changes: 1 addition & 1 deletion
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightRuntimeException.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,6 +41,6 @@ public CallStatus status() {
 @Override
 public String toString() {
 String s = getClass().getName();
 return String.format("%s: %s", s, status);
 return String.format("%s: %s: %s", s, status.code(), status.description());
 }
}

53 changes: 24 additions & 29 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,11 +34,11 @@
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorUnloader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

Expand DownExpand Up@@ -189,24 +189,22 @@ public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutR
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);

 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When the ackStream is completed, the FlightStream will be closed with it
 ackStream.setAutoCloseable(fs);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
 executors.submit(() -> {
 final StreamPipe<PutResult, Flight.PutResult> ackStream = StreamPipe
 .wrap(responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware);
 try {
 producer.acceptPut(makeContext(responseObserver), fs, ackStream).run();
 } catch (Exception ex) {
 ackStream.onError(ex);
 } finally {
 // Close this stream before telling gRPC that the call is complete. That way we don't race with server shutdown.
 try {
 fs.close();
 } catch (Exception e) {
 handleExceptionWithMiddleware(e);
 }
 // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself
 // We don't do this for other streams since the implementation may be asynchronous
 ackStream.ensureCompleted();
Expand DownExpand Up@@ -236,7 +234,7 @@ public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<Flight
 */
 private void handleExceptionWithMiddleware(Throwable t) {
 final Map<Key<?>, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get();
 if (middleware == null) {
 if (middleware == null || middleware.isEmpty()) {
 logger.error("Uncaught exception in Flight method body", t);
 return;
 }
Expand All@@ -258,14 +256,14 @@ public void getSchema(Flight.FlightDescriptor request, StreamObserver<Flight.Sch

 /** Ensures that other resources are cleaned up when the service finishes its call. */
 private static class ExchangeListener extends GetListener {
 private final AutoCloseable resource;

 private AutoCloseable resource;
 private boolean closed = false;
 private Runnable onCancelHandler = null;

 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler,
 AutoCloseable resource) {
 public ExchangeListener(ServerCallStreamObserver<ArrowMessage> responseObserver, Consumer<Throwable> errorHandler) {
 super(responseObserver, errorHandler);
 this.resource = resource;
 this.resource = null;
 super.setOnCancelHandler(() -> {
 try {
 if (onCancelHandler != null) {
Expand All@@ -285,7 +283,7 @@ private void cleanup() {
 }
 closed = true;
 try {
this.resource.close();
AutoCloseables.close(resource);
 } catch (Exception e) {
 throw CallStatus.INTERNAL
 .withCause(e)
Expand DownExpand Up@@ -321,19 +319,16 @@ public void setOnCancelHandler(Runnable handler) {
 public StreamObserver<ArrowMessage> doExchangeCustom(StreamObserver<ArrowMessage> responseObserverSimple) {
 final ServerCallStreamObserver<ArrowMessage> responseObserver =
 (ServerCallStreamObserver<ArrowMessage>) responseObserverSimple;
 final FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, (String message, Throwable cause) -> {
 responseObserver.onError(Status.CANCELLED.withCause(cause).withDescription(message).asException());
 }, responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 final ExchangeListener listener = new ExchangeListener(
 responseObserver,
 this::handleExceptionWithMiddleware,
 () -> {
 // Force the stream to "complete" so it will close without incident. At this point, we don't care since
 // we are about to end the call. (Normally it will raise an error.)
 fs.completed.complete(null);
 fs.close();
 });
 this::handleExceptionWithMiddleware);
 final FlightStream fs = new FlightStream(
 allocator,
 PENDING_REQUESTS,
 /* server-upload streams are not cancellable */null,
 responseObserver::request);
 // When service completes the call, this cleans up the FlightStream
 listener.resource = fs;
 responseObserver.disableAutoInboundFlowControl();
 responseObserver.request(1);
 final StreamObserver<ArrowMessage> observer = fs.asObserver();
Expand Down

121 changes: 85 additions & 36 deletions
java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
Expand DownExpand Up@@ -54,7 +55,6 @@
 * An adaptor between protobuf streams and flight data streams.
 */
public class FlightStream implements AutoCloseable {

 // Use AutoCloseable sentinel objects to simplify logic in #close
 private final AutoCloseable DONE = () -> {
 };
Expand All@@ -68,7 +68,13 @@ public class FlightStream implements AutoCloseable {
 private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
 private final int pendingTarget;
 private final Requestor requestor;
 // The completion flags.
 // This flag is only updated as the user iterates through the data, i.e. it tracks whether the user has read all the
 // data and closed the stream
 final CompletableFuture<Void> completed;
 // This flag is immediately updated when gRPC signals that the server has ended the call. This is used to make sure
 // we don't block forever trying to write to a server that has rejected a call.
 final CompletableFuture<Void> cancelled;

 private volatile int pending = 1;
 private volatile VectorSchemaRoot fulfilledRoot;
Expand All@@ -84,16 +90,19 @@ public class FlightStream implements AutoCloseable {
 *
 * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
 * @param pendingTarget Target number of messages to receive.
 * @param cancellable Only provided for streams from server to client, used to cancel mid-stream requests.
 * @param cancellable Used to cancel mid-stream requests.
 * @param requestor A callback to determine how many pending items there are.
 */
 public FlightStream(BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
 Objects.requireNonNull(allocator);
 Objects.requireNonNull(requestor);
 this.allocator = allocator;
 this.pendingTarget = pendingTarget;
 this.cancellable = cancellable;
 this.requestor = requestor;
 this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
 this.completed = new CompletableFuture<>();
 this.cancelled = new CompletableFuture<>();
 }

 /**
Expand DownExpand Up@@ -158,29 +167,52 @@ public FlightDescriptor getDescriptor() {
 /**
 * Closes the stream (freeing any existing resources).
 *
 * <p>If the stream isn't complete and is cancellable, this method will cancel the stream first.</p>
 * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the stream first.
 */
 public void close() throws Exception {
 final List<AutoCloseable> closeables = new ArrayList<>();
 // cancellation can throw, but we still want to clean up resources, so make it an AutoCloseable too
 closeables.add(() -> {
 if (!completed.isDone() && cancellable != null) {
 cancel("Stream closed before end.", /* no exception to report */ null);
 Throwable suppressor = null;
 if (cancellable != null) {
 // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after close() ends.
 // On the server side, we can't rely on draining the stream , because this gRPC bug means the completion callback
 // may never run https://github.com/grpc/grpc-java/issues/5882
 try {
 synchronized (cancellable) {
 if (!cancelled.isDone()) {
 // Only cancel if the call is not done on the gRPC side
 cancellable.cancel("Stream closed before end", /* no exception to report */null);
 }
 }
 // Drain the stream without the lock (as next() implicitly needs the lock)
 while (next()) { }
 } catch (FlightRuntimeException e) {
 suppressor = e;
 }
 });
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 // Perform these operations under a lock. This way the observer can't enqueue new messages while we're in the
@rymurrrymurrSep 5, 2020

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe its just the wording of the comments but this seems like its theoretically possible for an observer to put a message between 187 and the lock gets aquired in 195. Is that true? The chance is prob pretty small and not easy to code for. Just wanted a bit of clarification
@lidavidmlidavidmSep 5, 2020

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes - it's possible because the observer is run in a separate thread (by gRPC) than the application, so the observer can trigger when the application is in the middle of close(). On the client side, draining the stream (as done here) prevents this case, but on the server side, we can't rely on this, unfortunately, hence the lock to protect things.
 // middle of cleanup. This should only be a concern for server-side streams since client-side streams are drained
 // by the lambda above.
 synchronized (completed) {
 try {
 if (fulfilledRoot != null) {
 closeables.add(fulfilledRoot);
 }
 closeables.add(applicationMetadata);
 closeables.addAll(queue);
 if (dictionaries != null) {
 dictionaries.getDictionaryIds().forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
 }
 if (suppressor != null) {
 AutoCloseables.close(suppressor, closeables);
 } else {
 AutoCloseables.close(closeables);
 }
 } finally {
 // The value of this CompletableFuture is meaningless, only whether it's completed (or has an exception)
 // No-op if already complete
 completed.complete(null);
 }
 }

 AutoCloseables.close(closeables);
 // Other code ignores the value of this CompletableFuture, only whether it's completed (or has an exception)
 // No-op if already complete; do this after the check in the AutoCloseable lambda above
 completed.complete(null);
 }

 /**
Expand DownExpand Up@@ -337,8 +369,22 @@ private class Observer implements StreamObserver<ArrowMessage> {
 super();
 }

 /** Helper to add an item to the queue under the appropriate lock. */
 private void enqueue(AutoCloseable message) {
 synchronized (completed) {
 if (completed.isDone()) {
 // The stream is already closed (RPC ended), discard the message
 AutoCloseables.closeNoChecked(message);
 } else {
 queue.add(message);
 }
 }
 }

 @Override
 public void onNext(ArrowMessage msg) {
 // Operations here have to be under a lock so that we don't add a message to the queue while in the middle of
 // close().
 requestOutstanding();
 switch (msg.getMessageType()) {
 case NONE: {
Expand All@@ -347,7 +393,7 @@ public void onNext(ArrowMessage msg) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 if (msg.getApplicationMetadata() != null) {
queue.add(msg);
enqueue(msg);
 }
 break;
 }
Expand All@@ -367,60 +413,63 @@ public void onNext(ArrowMessage msg) {
 try {
 MetadataV4UnionChecker.checkRead(schema, metadataVersion);
 } catch (IOException e) {
 queue.add(DONE_EX);
 ex = e;
 enqueue(DONE_EX);
 break;
 }

 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 synchronized (completed) {
 if (!completed.isDone()) {
 fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
 loader = new VectorLoader(fulfilledRoot);
 if (msg.getDescriptor() != null) {
 descriptor.set(new FlightDescriptor(msg.getDescriptor()));
 }
 root.set(fulfilledRoot);
 }
 }
 root.set(fulfilledRoot);
 break;
 }
 case RECORD_BATCH:
 queue.add(msg);
 break;
 case DICTIONARY_BATCH:
queue.add(msg);
enqueue(msg);
 break;
 case TENSOR:
 default:
 queue.add(DONE_EX);
 ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 enqueue(DONE_EX);
 }
 }

 @Override
 public void onError(Throwable t) {
 ex = StatusUtils.fromThrowable(t);
 queue.add(DONE_EX);
 cancelled.complete(null);
 root.setException(ex);
 }

 @Override
 public void onCompleted() {
 // Depends on gRPC calling onNext and onCompleted non-concurrently
 cancelled.complete(null);
 queue.add(DONE);
 }
 }

 /**
 * Cancels sending the stream to a client.
 *
 * @throws UnsupportedOperationException on a stream being uploaded from the client.
 * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before cancellation are
 * received and to wait for the underlying transport to acknowledge cancellation.
 */
 public void cancel(String message, Throwable exception) {
 completed.completeExceptionally(
 CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException());
 if (cancellable != null) {
 cancellable.cancel(message, exception);
 } else {
 if (cancellable == null) {
 throw new UnsupportedOperationException("Streams cannot be cancelled that are produced by client. " +
 "Instead, server should reject incoming messages.");
 }
 cancellable.cancel(message, exception);
 // Do not mark the stream as completed, as gRPC may still be delivering messages.
 }

 StreamObserver<ArrowMessage> asObserver() {
Expand Down
Loading