Github user sachouche commented on a diff in the pull request:
https://github.com/apache/drill/pull/1237#discussion_r184138876
--- Diff:
exec/java-exec/src/main/java/org/apache/drill/exec/record/RawFragmentBatch.java
---
@@ -77,4 +83,46 @@ public long getByteCount() {
public boolean isAckSent() {
return ackSent.get();
}
+
+ /**
+ * Transfer ownership of this DrillBuf to the target allocator. This is
done for better memory
+ * accounting (that is, the operator should be charged with the body's
Drillbuf memory).
+ *
+ * <p><b>NOTES -</b>
+ * <ul>
+ * <li>This operation is a NOOP when a) the current allocator
(associated with the DrillBuf) is not the
+ * owning allocator or b) the target allocator is already the owner
+ * <li>When transfer happens, a new RawFragmentBatch instance is
allocated; this is done for proper
+ * DrillBuf reference count accounting
+ * <li>The RPC handling code caches a reference to this RawFragmentBatch
object instance; release()
+ * calls should be routed to the previous DrillBuf
+ * </ul>
+ *
+ * @param targetAllocator target allocator
+ * @return a new {@link RawFragmentBatch} object instance on success
(where the buffer ownership has
+ * been switched to the target allocator); otherwise this
operation is a NOOP (current instance
+ * returned)
+ */
+ public RawFragmentBatch transferBodyOwnership(BufferAllocator
targetAllocator) {
+ if (body == null) {
+ return this; // NOOP
+ }
+
+ if (!body.getLedger().isOwningLedger()
+ || body.getLedger().isOwner(targetAllocator)) {
+
+ return this;
+ }
+
+ int writerIndex = body.writerIndex();
+ TransferResult transferResult =
body.transferOwnership(targetAllocator);
+
+ // Set the index and increment reference count
+ transferResult.buffer.writerIndex(writerIndex);
+
+ // Clear the current Drillbuffer since caller will perform release()
on the new one
+ body.release();
+
+ return new RawFragmentBatch(getHeader(), transferResult.buffer,
getSender(), false);
--- End diff --
That was my original code which failed miserably:
- The RPC code has references on the DrillBuf and the RawFragmentBatch
- This means, we need to ensure that a release call is routed to the right
DrillBuf (otherwise, the reference count logic stops working)
- Creating a new RawFragmentBatch instance essentially provided just that
(proper reference count accounting)
---
