raghavyadav01 commented on code in PR #19021: URL: https://github.com/apache/pinot/pull/19021#discussion_r4328290761
########## pinot-foreign-buffers/src/main/java/org/apache/pinot/segment/memory/foreign/ForeignPinotBuffer.java: ########## @@ -0,0 +1,247 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.memory.foreign; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.lang.ref.Cleaner; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import javax.annotation.Nullable; +import org.apache.pinot.segment.spi.memory.NonNativePinotDataBuffer; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// A [PinotDataBuffer] backed by the Foreign Function & Memory API ([MemorySegment]). +/// +/// This is the Java 21+ replacement for the `sun.misc.Unsafe` based buffer. Like the Unsafe implementation it only +/// reads and writes in [native byte order][ByteOrder#nativeOrder()]; non-native order is provided by wrapping it in a +/// [NonNativePinotDataBuffer]. Because [MemorySegment] natively supports `long` indexing, a single implementation +/// covers buffers larger than [Integer#MAX_VALUE] without any of the manual pointer arithmetic, page alignment or +/// reflective `DirectByteBuffer` construction the Unsafe path required. +/// +/// All accessors use the `*_UNALIGNED` [ValueLayout] constants because callers read and write at arbitrary byte +/// offsets. Bounds checking is performed by [MemorySegment] itself (throwing [IndexOutOfBoundsException]), so this +/// class does not add its own checks. Access is memory-safe: reads and writes after the backing [Arena] is closed fail +/// deterministically instead of corrupting memory. +/// +/// Thread-safety: instances created through [ForeignPinotBufferFactory] use a [shared arena][Arena#ofShared()], so the +/// buffer may be read, written and closed from any thread. As with every [PinotDataBuffer], closing a buffer while +/// another thread accesses it is a caller error. Note that closing a shared arena does not fail fast: [#release()] +/// blocks until all in-flight accesses on this buffer and its views have completed (and the inherited synchronized +/// [#close()] serialises concurrent close calls on the same instance). +public class ForeignPinotBuffer extends PinotDataBuffer { + + private static final Logger LOGGER = LoggerFactory.getLogger(ForeignPinotBuffer.class); + /// Leak safety net for owning buffers. A shared [Arena] is not GC-managed (unlike [Arena#ofAuto()]), so if a buffer + /// is never closed the native memory would leak. This cleaner closes the arena on collection, mirroring the + /// `finalize()` safety net the `Unsafe` implementation used (but without the deprecated finalizer). + private static final Cleaner CLEANER = Cleaner.create(); + + // Native-order, byte-aligned (alignment == 1) layouts. Single-byte access is always aligned. + private static final ValueLayout.OfByte BYTE = ValueLayout.JAVA_BYTE; + private static final ValueLayout.OfChar CHAR = ValueLayout.JAVA_CHAR_UNALIGNED; + private static final ValueLayout.OfShort SHORT = ValueLayout.JAVA_SHORT_UNALIGNED; + private static final ValueLayout.OfInt INT = ValueLayout.JAVA_INT_UNALIGNED; + private static final ValueLayout.OfLong LONG = ValueLayout.JAVA_LONG_UNALIGNED; + private static final ValueLayout.OfFloat FLOAT = ValueLayout.JAVA_FLOAT_UNALIGNED; + private static final ValueLayout.OfDouble DOUBLE = ValueLayout.JAVA_DOUBLE_UNALIGNED; + + private final MemorySegment _segment; + /// The cleanup action registered with the [Cleaner] for owning buffers, or {@code null} for a non-owning view. + @Nullable + private final ArenaCleanup _cleanup; + /// The handle to the registered cleaning action, used to free eagerly (and de-register) on explicit close. + @Nullable + private final Cleaner.Cleanable _cleanable; + + /// Creates a buffer over the given segment. + /// + /// @param segment the memory segment this buffer reads and writes. Its [MemorySegment#byteSize()] is the buffer size. + /// @param arena the arena to close when this buffer is released, or {@code null} for a non-owning view. + public ForeignPinotBuffer(MemorySegment segment, @Nullable Arena arena) { + super(arena != null); + _segment = segment; + if (arena != null) { + _cleanup = new ArenaCleanup(arena, segment.byteSize()); + // The registered action (ArenaCleanup) must not reference this buffer, otherwise it would never be collected. + _cleanable = CLEANER.register(this, _cleanup); + } else { + _cleanup = null; + _cleanable = null; + } + } + + @Override + public byte getByte(long offset) { + return _segment.get(BYTE, offset); + } + + @Override + public void putByte(long offset, byte value) { + _segment.set(BYTE, offset, value); + } + + @Override + public char getChar(long offset) { + return _segment.get(CHAR, offset); + } + + @Override + public void putChar(long offset, char value) { + _segment.set(CHAR, offset, value); + } + + @Override + public short getShort(long offset) { + return _segment.get(SHORT, offset); + } + + @Override + public void putShort(long offset, short value) { + _segment.set(SHORT, offset, value); + } + + @Override + public int getInt(long offset) { + return _segment.get(INT, offset); + } + + @Override + public void putInt(long offset, int value) { + _segment.set(INT, offset, value); + } + + @Override + public long getLong(long offset) { + return _segment.get(LONG, offset); + } + + @Override + public void putLong(long offset, long value) { + _segment.set(LONG, offset, value); + } + + @Override + public float getFloat(long offset) { + return _segment.get(FLOAT, offset); + } + + @Override + public void putFloat(long offset, float value) { + _segment.set(FLOAT, offset, value); + } + + @Override + public double getDouble(long offset) { + return _segment.get(DOUBLE, offset); + } + + @Override + public void putDouble(long offset, double value) { + _segment.set(DOUBLE, offset, value); + } + + @Override + public void copyTo(long offset, byte[] buffer, int destOffset, int size) { + // Direct segment-to-array copy: a bounds-checked, JIT-intrinsified memcpy with no intermediate ByteBuffer. + MemorySegment.copy(_segment, BYTE, offset, buffer, destOffset, size); + } + + @Override + public void readFrom(long offset, byte[] buffer, int srcOffset, int size) { + // Direct array-to-segment copy: a bounds-checked, JIT-intrinsified memcpy with no intermediate ByteBuffer. + MemorySegment.copy(buffer, srcOffset, _segment, BYTE, offset, size); + } + + @Override + public long size() { + return _segment.byteSize(); + } + + @Override + public ByteOrder order() { + return ByteOrder.nativeOrder(); + } + + @Override + public PinotDataBuffer view(long start, long end, ByteOrder byteOrder) { + // asSlice validates start/size against the segment bounds. + ForeignPinotBuffer nativeView = new ForeignPinotBuffer(_segment.asSlice(start, end - start), null); Review Comment: The view doesn't keep the owning buffer reachable, so if a caller holds only the view, the owner's Cleaner can close the arena under it and the next read throws `Already closed`. The Unsafe view keeps the parent's memory alive. Could the view hold a reference to the owner? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
