@Override
    public BufferWithSubpartition getNextBuffer(@Nullable MemorySegment transitBuffer) {
        checkState(isFinished, "Sort buffer is not ready to be read.");
        checkState(!isReleased, "Sort buffer is already released.");

        if (!hasRemaining()) {
            freeSegments.add(transitBuffer);
            return null;
        }

        int numBytesRead = 0;
        Buffer.DataType bufferDataType = Buffer.DataType.DATA_BUFFER;
        int currentReadingSubpartitionId = subpartitionReadOrder[readOrderIndex];

        do {
            // Get the buffer index and offset from the index entry
            int toReadBufferIndex = getSegmentIndexFromPointer(readIndexEntryAddress);
            int toReadOffsetInBuffer = getSegmentOffsetFromPointer(readIndexEntryAddress);

            // Get the lengthAndDataType buffer according the buffer index
            MemorySegment toReadBuffer = segments.get(toReadBufferIndex);

            // From the lengthAndDataType buffer, read and get the length and the data type
            long lengthAndDataType = toReadBuffer.getLong(toReadOffsetInBuffer);
            int recordLength = getSegmentIndexFromPointer(lengthAndDataType);
            Buffer.DataType dataType =
                    Buffer.DataType.values()[getSegmentOffsetFromPointer(lengthAndDataType)];

            // If the buffer is an event and some data has been read, return it directly to ensure
            // that the event will occupy one buffer independently
            if (dataType.isEvent() && numBytesRead > 0) {
                break;
            }
            bufferDataType = dataType;

            // Get the next index entry address and move the read position forward
            long nextReadIndexEntryAddress = toReadBuffer.getLong(toReadOffsetInBuffer + 8);
            toReadOffsetInBuffer += INDEX_ENTRY_SIZE;

            // Allocate a temp buffer for the event, recycle the original buffer
            if (bufferDataType.isEvent()) {
                freeSegments.add(transitBuffer);
                transitBuffer = MemorySegmentFactory.allocateUnpooledSegment(recordLength);
            }

            if (!isPartialRecordAllowed
                    && !isLastBufferPartialRecord
                    && numBytesRead > 0
                    && numBytesRead + recordLength > transitBuffer.size()) {
                break;
            }

            // Start reading data from the data buffer
            numBytesRead +=
                    copyRecordOrEvent(
                            transitBuffer,
                            numBytesRead,
                            toReadBufferIndex,
                            toReadOffsetInBuffer,
                            recordLength);

            if (recordRemainingBytes == 0) {
                // move to next subpartition if the current subpartition has been finished
                if (readIndexEntryAddress
                        == lastIndexEntryAddresses[currentReadingSubpartitionId]) {
                    isLastBufferPartialRecord = false;
                    updateReadSubpartitionAndIndexEntryAddress();
                    break;
                }
                readIndexEntryAddress = nextReadIndexEntryAddress;
                if (isLastBufferPartialRecord) {
                    isLastBufferPartialRecord = false;
                    break;
                }
            } else {
                isLastBufferPartialRecord = true;
            }
        } while (numBytesRead < transitBuffer.size() && bufferDataType.isBuffer());

        if (!isPartialRecordAllowed
                && !isLastBufferPartialRecord
                && bufferDataType == Buffer.DataType.DATA_BUFFER) {
            bufferDataType = Buffer.DataType.DATA_BUFFER_WITH_CLEAR_END;
        }

        numTotalBytesRead += numBytesRead;
        return new BufferWithSubpartition(
                new NetworkBuffer(
                        transitBuffer,
                        bufferDataType.isBuffer() ? bufferRecycler : FreeingBufferRecycler.INSTANCE,
                        bufferDataType,
                        numBytesRead),
                currentReadingSubpartitionId);
    }
@Test
    void testBufferIsRecycledWhenGetEvent() throws Exception {
        int numSubpartitions = 10;
        int bufferPoolSize = 512;
        int bufferSizeBytes = 1024;
        int numBuffersForSort = 20;
        int subpartitionId = 0;
        Random random = new Random(1234);

        NetworkBufferPool globalPool = new NetworkBufferPool(bufferPoolSize, bufferSizeBytes);
        BufferPool bufferPool = globalPool.createBufferPool(bufferPoolSize, bufferPoolSize);

        LinkedList<MemorySegment> segments = new LinkedList<>();
        for (int i = 0; i < numBuffersForSort; ++i) {
            segments.add(bufferPool.requestMemorySegmentBlocking());
        }
        TieredStorageSortBuffer sortBuffer =
                new TieredStorageSortBuffer(
                        segments,
                        bufferPool,
                        numSubpartitions,
                        bufferSizeBytes,
                        numBuffersForSort,
                        true);

        byte[] bytes = new byte[1];
        random.nextBytes(bytes);
        ByteBuffer dataRecord = ByteBuffer.wrap(bytes);
        sortBuffer.append(dataRecord, subpartitionId, Buffer.DataType.DATA_BUFFER);
        ByteBuffer eventRecord = ByteBuffer.wrap(bytes);
        sortBuffer.append(eventRecord, subpartitionId, Buffer.DataType.EVENT_BUFFER);
        sortBuffer.finish();

        MemorySegment memorySegment = bufferPool.requestMemorySegmentBlocking();
        BufferWithSubpartition bufferWithSubpartition = sortBuffer.getNextBuffer(memorySegment);
        assertThat(bufferWithSubpartition.getBuffer().isBuffer()).isTrue();
        assertThat(bufferWithSubpartition.getSubpartitionIndex()).isEqualTo(subpartitionId);
        bufferWithSubpartition.getBuffer().recycleBuffer();
        assertThat(bufferPool.bestEffortGetNumOfUsedBuffers()).isEqualTo(numBuffersForSort);

        bufferWithSubpartition = sortBuffer.getNextBuffer(memorySegment);
        assertThat(bufferWithSubpartition.getBuffer().isBuffer()).isFalse();
        assertThat(bufferWithSubpartition.getSubpartitionIndex()).isEqualTo(subpartitionId);
        assertThat(bufferPool.bestEffortGetNumOfUsedBuffers()).isEqualTo(numBuffersForSort);
    }