[20240814] XTDB Internals 1 - MemoryBufferPool

• [공식문서]​
• [GITHUB SOURCE CODE]​ : 별기능은 없고 BaseAllocator 를 확장함.

• [BaseAllocator SOURCE CODE]​

  BaseAllocator의 wrapForeignAllocation(ForeignAllocation allocation) 에 대해 조금만 보면 ByteBuffer를 ArrowBuf 로 변환하고 있다. 자세한 내용은 아래설명에서 찾아볼 것.

  final ArrowBuf buf =
      new ArrowBuf(ledger, /*bufferManager=*/ null, size, allocation.memoryAddress());
  return buf;
  

locking 은 기본 clojure api 이다. (https://clojuredocs.org/clojure.core/locking)

이건 동기화를 위해 사용하는 기능이고, 내부구현은 special form인 monitor-enter 를 사용한다. (https://clojure.org/reference/special_forms#monitor-enter)

java.nio.ByteBuffer 를 ArrowBuf 로 변환하기 위한 함수.

중요한 것은 .wrapForeignallocation 이다. 외부 할당된 녀석(ByteBuffer) 를 ArrowBuf 구현으로 가져오는 것이다.

BufferAllocator 는 Arrow 에서 사용하는 메모리를 할당하여 버퍼를 생성하는 녀석이 아닌가 싶다.

https://github.com/xtdb/xtdb/blob/main/core/src/main/clojure/xtdb/util.clj#L434

(defn ->arrow-buf-view
  (^org.apache.arrow.memory.ArrowBuf [^BufferAllocator allocator ^ByteBuffer nio-buffer]
   (->arrow-buf-view allocator nio-buffer nil))

  (^org.apache.arrow.memory.ArrowBuf [^BufferAllocator allocator ^ByteBuffer nio-buffer release-fn]
 (let [nio-buffer (if (and (.isDirect nio-buffer) (zero? (.position nio-buffer)))
		    ;; 빈 nio-buffer 다이렉트버퍼면 그대로 리턴.
		    nio-buffer
		    ;; 비어있지 않거나, 다이렉트 버퍼가 아니면 다이렉트버퍼로 다시 만들어서 리턴.
		    (-> (ByteBuffer/allocateDirect (.remaining nio-buffer))
			(.put (.duplicate nio-buffer))
			(.clear)))]
   ;; BufferAllocator를 이용하여 nio-buffer(ArrowBuf) 의 주소와 사이즈를 리턴하면, 그것으로 ArrowBuf로 할당하는 듯.
   (.wrapForeignAllocation allocator
			   ;; ForeignAllocation 클래스는 거의 데이터 클래스와 같음. 별거 없음.
			   ;; 주소와 사이즈를 머금고 있고 release0 을 이용해서 free 방법만 위임해준다.
			   (proxy [ForeignAllocation] [(.remaining nio-buffer) (MemoryUtil/getByteBufferAddress nio-buffer)]
			     (release0 []
			      (try-free-direct-buffer nio-buffer)
			       (when release-fn
				 (release-fn))))))))

1. 입력으로 BufferAllocator와 ByteBuffer를 받는다.

2. 입력 ByteBuffer가 다이렉트 버퍼가 아니거나 position이 0이 아닌 경우, 새로운 다이렉트 버퍼를 생성하고 데이터를 복사함.

   (let [nio-buffer (if (and (.isDirect nio-buffer) (zero? (.position nio-buffer)))
   	      nio-buffer
   	      (-> (ByteBuffer/allocateDirect (.remaining nio-buffer))
   		  (.put (.duplicate nio-buffer))
   		  (.clear)))]
   

3. allocator.wrapForeignAllocation을 사용하여 ArrowBuf를 생성합니다. 이는 외부 메모리를 Arrow 버퍼로 래핑.
4. ForeignAllocation 프록시를 사용하여 버퍼 해제 로직을 구현합니다.
5. 선택적으로 release-fn을 받아 버퍼 해제 시 추가 정리 작업을 수행할 수 있게 합니다.

6. 네티(Netty)의 PlatformDependent.freeDirectBuffer를 사용하여 다이렉트 버퍼를 해제합니다.

   네티 기능은 try-free-direct-buffer 함수에서 사용함. Netty로 이용하는 함수는 (io.netty.util.internal.PlatformDependent/freeDirectBuffer nio-buffer) 이다.

   (.wrapForeignAllocation allocator
   		   (proxy [ForeignAllocation] [(.remaining nio-buffer) (MemoryUtil/getByteBufferAddress nio-buffer)]  ;;[size memroy-addres]
   		     (release0 []
   		      (try-free-direct-buffer nio-buffer)
   		       (when release-fn
   			 (release-fn)))))
   

Arrow의 메모미 관리 핵심 부분.

1. 구조
   • ArrowBuf는 기본적으로 off-heap 메모리를 래핑함.
   • 내부적으로 nio.ByteBuffer를 사용하여 실제 메모리를 관리(?)
2. 메모리 관리
   • 참조 카운팅 방식 이용.
   • BufferAllocator를 통해 생성, 관리

https://docs.oracle.com/javase/8/docs/api/java/nio/ByteBuffer.html

java.nio 에 있는 ByteBuffer 또한 핵심 개념임.

특이 Direct 여부가 중요하다.

1. 다이렉트 버퍼 (Direct Buffer):
   • `isDirect()` 메소드가 true를 반환합니다.
   • JVM은 이 버퍼에 대해 네이티브 I/O 작업을 직접 수행하려고 노력합니다.
   • 중간 버퍼로의 복사 과정을 피하여 I/O 성능을 향상시킵니다.
   • `ByteBuffer.allocateDirect()` 메소드로 생성할 수 있습니다.
   • 일반적으로 할당과 해제 비용이 더 높습니다.
   • 내용이 일반적인 가비지 컬렉션 힙 외부에 있을 수 있어, 메모리 사용량 파악이 어려울 수 있습니다.
   • 파일의 특정 영역을 메모리에 직접 매핑하여 생성할 수도 있습니다.
   • 주로 크고 수명이 긴 버퍼에 사용되며, 특히 네이티브 I/O 작업에 적합합니다.
2. 비다이렉트 버퍼 (Non-direct Buffer):
   • `isDirect()` 메소드가 false를 반환합니다.
   • 일반적인 JVM 힙 메모리에 할당됩니다.
   • 다이렉트 버퍼에 비해 할당과 해제 비용이 낮습니다.
3. 사용 지침:
   • 다이렉트 버퍼는 주로 크고 수명이 긴 버퍼에 사용하는 것이 좋습니다.
   • 프로그램 성능에 측정 가능한 이득이 있을 때만 다이렉트 버퍼를 할당하는 것이 좋습니다.
   • 성능이 중요한 코드에서는 `isDirect()` 메소드를 사용하여 명시적인 버퍼 관리를 할 수 있습니다.
4. 주의사항:
   • 다이렉트 버퍼가 접근 불가능한 메모리 영역을 참조하는 경우, 해당 영역에 접근하려고 하면 버퍼의 내용은 변경되지 않고 특정되지 않은 예외가 발생할 수 있습니다.
5. 추가 생성 방법:
   • JNI를 통해 네이티브 코드에서 다이렉트 바이트 버퍼를 생성할 수도 있습니다(Java 플랫폼 구현에 따라 선택적으로 지원).

이러한 특성들로 인해, 다이렉트 버퍼는 특히 대용량 데이터 처리나 빈번한 네이티브 I/O 작업이 필요한 경우에 성능상의 이점을 제공할 수 있습니다.

우리가 보는 코드에서는 Direct Buffer를 만든다.

그리고 allocator.wrapForeignAllocation 으로 아까 만든 DirectByteBuffer 를 Apache Arrow 메모리 관리 시스템으로 통합한다.

메모리 관리 통합:

Arrow의 BufferAllocator가 외부 메모리를 관리할 수 있게 합니다. 이를 통해 Arrow의 메모리 추적 및 참조 카운팅 메커니즘을 활용할 수 있습니다.

1. 성능 최적화: 다이렉트 버퍼의 네이티브 메모리를 직접 사용하여 추가적인 복사를 방지합니다. I/O 작업이나 네이티브 코드와의 상호 작용 시 효율성을 높입니다.

1. 리소스 해제 관리: ForeignAllocation 프록시를 통해 버퍼가 더 이상 필요하지 않을 때 적절히 해제될 수 있도록 합니다. try-free-direct-buffer 함수를 통해 다이렉트 버퍼의 메모리를 명시적으로 해제합니다.

1. Arrow 생태계와의 호환성: 외부 메모리를 Arrow의 ArrowBuf로 감싸 Arrow의 다른 컴포넌트들과 원활하게 사용할 수 있게 합니다.

1. 유연한 메모리 소스 지원: NIO 버퍼뿐만 아니라 다양한 외부 메모리 소스를 Arrow 시스템에 통합할 수 있는 방법을 제공합니다.

5.제로-카피 연산 지원: 가능한 경우 메모리 복사를 피하고 직접 참조를 사용하여 성능을 최적화합니다.

이 접근 방식은 Java NIO의 다이렉트 버퍼의 이점(네이티브 I/O 성능)과 Apache Arrow의 메모리 관리 기능을 결합하여, 고성능 데이터 처리 파이프라인에서 효율적인 메모리 사용을 가능하게 합니다.

allocator.wrapForeignallocation 구현에 대해 알아보자.

• ForeignAllocation 프록시 객체를 생성한다. (https://arrow.apache.org/docs/dev/java/reference/org/apache/arrow/memory/ForeignAllocation.html)
  • 생성자 : protected ForeignALlocation(long size, long memoryAddress)
  • 메소드 :
    • public long getSize() : get the size of this allocation.
    • protected long memoryAddress : get the address of this allocation
    • protected abstract void release0() : free this allocation. will only called once.
  • 코드상 생성자에 버퍼의 크기와 주소를 전달한다. (proxy [ForeignAllocation] [(.remaining nio-buffer) (MemoryUtil/getByteBufferAddress nio-buffer)]

  • [소스코드 링크]​

     /*
     * 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.arrow.memory;
    
    /**
     * EXPERIMENTAL: a memory allocation that does not come from a BufferAllocator, but rather an
     * outside source (like JNI).
     *
     * <p>To use this, subclass this class and implement {@link #release0()} to free the allocation.
     */
    public abstract class ForeignAllocation {
      private final long memoryAddress;
      private final long size;
    
      /**
       * Create a new AllocationManager representing an imported buffer.
       *
       * @param size The buffer size.
       * @param memoryAddress The buffer address.
       */
      protected ForeignAllocation(long size, long memoryAddress) {
        this.memoryAddress = memoryAddress;
        this.size = size;
      }
    
      /** Get the size of this allocation. */
      public long getSize() {
        return size;
      }
    
      /** Get the address of this allocation. */
      protected long memoryAddress() {
        return memoryAddress;
      }
    
      /** Free this allocation. Will only be called once. */
      protected abstract void release0();
    }
    

• 메모리 해제 로직 구현
  • try-free-direct-buffer 함수를 호출하여 다이렉트 버퍼의 메모리를 해제
  • 추가적인 release-fn 이 제공되는 경우 이를 호출.

이 코드는 Java NIO의 ByteBuffer를 Apache Arrow의 ArrowBuf로 효율적으로 변환하며, 메모리 관리를 Arrow 시스템에 위임합니다. 이를 통해 네이티브 메모리 사용의 이점을 유지하면서 Arrow의 메모리 관리 기능을 활용할 수 있습니다.

org.apache.arrow.memory.util.MemoryUtil/getByteBufferAddress 는 nio의 ByteBuffer 에서 주소를 가져오는 함수.

메모리 주소 관리: MemoryUtil은 Java의 직접 메모리(Direct Memory)와 네이티브 메모리 사이의 상호 작용을 용이하게 합니다.

안전한 메모리 접근: 이 클래스는 JVM의 경계를 넘어 직접 메모리에 접근할 때 필요한 안전 장치를 제공합니다.

플랫폼 독립성: 다양한 플랫폼에서 일관된 방식으로 메모리 주소를 다룰 수 있게 해줍니다.

API 문서의 설명을 보자.

public static long getByteBufferAddress(ByteBuffer buf)

Given a ByteBuffer, gets the address the underlying memory space.

Parameters:
  buf - the byte buffer.

Returns:
  address of the underlying memory.

소스코드를 보자. GITHUB SOURCE

/** The offset of the address field with the {@link java.nio.ByteBuffer} object. */
private static final long BYTE_BUFFER_ADDRESS_OFFSET;

static {
  // ...
  // get the offset of the address field in a java.nio.Buffer object
  Field addressField = java.nio.Buffer.class.getDeclaredField("address");
  addressField.setAccessible(true);
  // ...
}
/**
 * Given a {@link ByteBuffer}, gets the address the underlying memory space.
 *
 * @param buf the byte buffer.
 * @return address of the underlying memory.
 */
 public static long getByteBufferAddress(ByteBuffer buf) {
   return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET);
 }

UNSAFE 라는 기능을 쓴다. 오랫만에 본다. UNSAFE 는 예전에 블로그를 쓴 적이 있다. 다른 플랫폼에 적었던 것인데 찾아서 이곳에도 옮겨놓아야겠다.

Unsafe 는 Java 버전이 올라가면서 쓰지 말라고 권고하는 클래스 중에 하나이다. [GITHUB SOURCE CODE LINK]​

/**
  * @deprecated Use {@link java.lang.foreign} to access off-heap memory.
  *
  * @see #getByte(long)
  */
 @Deprecated(since="23", forRemoval=true)
 @ForceInline
 public long getLong(long address) {
   beforeMemoryAccess();
   return theInternalUnsafe.getLong(address);
}

// Reports the location of a given field in the storage allocation of its class.
// Do not expect to perform any sort of arithmetic on this offset;
// it is just a cookie which is passed to the unsafe heap memory accessors.
public long objectFieldOffset(Field f) {
  return objectFieldOffset0(f);
}

private native long objectFieldOffset0(Field f);

theInternalUnsafe.getLong 은 native 기능이다.

/** @see #getInt(Object, long) */
@IntrinsicCandidate
public native long    getLong(Object o, long offset);

이 메소드는 주어진 java 객체의 offset 에서 값을 가져온다.

putObject와는 다른 것인가?

이는 뭔가 Writer를 이용해서 저장해야 하는 값이 있는 것 같다. 하지만 이는 여러개를 한번에 저장하는 기능은 아닌 것 같다.

시작은 아래 이슈때문인 것 같다.

ArrowWriter 는 xtdb 내부에서 쓰는 BufferPool 구현체들이 사용하는 ArrowWriter 를 일관성있게 사용하려고 만든 인터페이스 같음.

우리가 이전에 IBufferPool.kt 의 인터페이스 중에 openArrowWriter 메소드의 리턴값이다.

package xtdb

interface ArrowWriter : AutoCloseable {
  fun writeBatch()
  fun end()
}

[stdb/core/src/main/kotlin/xtdb/ArrowWriter.kt]​

그런데 이 구현을 MemoryBefferPool 구현체에서 보면 이상한 것이 있다.

(openArrowWriter [this k rel]
  (let [baos (ByteArrayOutputStream.)]
    (util/with-close-on-catch [write-ch (Channels/newChannel baos)
			       unl (.startUnload rel write-ch)]
      (reify ArrowWriter
	(writeBatch [_] (.writeBatch unl))

	(end [_]
	  (.endFile unl)
	  (.close write-ch)
	  (.putObject this k (ByteBuffer/wrap (.toByteArray baos))))

	(close [_]
	  (util/close unl)
	  (when (.isOpen write-ch)
	    (.close write-ch)))))))

writeBatch 를 구현하라고 했더니, 그대로 (.writeBatch unl) 에 위임한다. 이미 writeBatch 를 구현해놓아져있다. 돌아가보면 rel 이 대체뭔지를 알아봐야겠다고 생각이 들었다.

[xtdb.arrow.Relation]​

Relation 은 RelationReader 인터페이스를 구현한다.

그냥 봐선 별기능이 있는 것 같지 않다. 그저 데이터를 Arrow 와 연결이 쉽도록 하기 위한 코드같아보인다.

궁금했던 writeBatch() 에 대해 알아보자.

// Apache Arrow 파일 포맷에서 사용하는 매직 바이트 인듯?
private val MAGIC = "ARROW1".toByteArray()

inner class RelationUnloader(private val ch: WriteChannel) : AutoCloseable {

	private val vectors = this@Relation.vectors.values
	private val schema = Schema(vectors.map { it.field })
	private val recordBlocks = mutableListOf<ArrowBlock>()

	init {
	    // 파일 시작에 매직 바이트 쓰기.
	    ch.write(MAGIC)
	    ch.align()
	    MessageSerializer.serialize(ch, schema)
	}

	fun writeBatch() {
	    val nodes = mutableListOf<ArrowFieldNode>()
	    val buffers = mutableListOf<ArrowBuf>()

	    vectors.forEach { it.unloadBatch(nodes, buffers) }

	    ArrowRecordBatch(rowCount, nodes, buffers).use { recordBatch ->
		MessageSerializer.serialize(ch, recordBatch)
		    .also { recordBlocks.add(it) }
	    }
	}

	fun endStream() {
	    ch.writeIntLittleEndian(MessageSerializer.IPC_CONTINUATION_TOKEN)
	    ch.writeIntLittleEndian(0)
	}

	fun endFile() {
	    endStream()

	    val footerStart = ch.currentPosition
	    ch.write(ArrowFooter(schema, emptyList(), recordBlocks), false)

	    val footerLength = ch.currentPosition - footerStart
	    check(footerLength > 0) { "Footer length must be positive" }
	    ch.writeIntLittleEndian(footerLength.toInt())
	    // 파일 마지막에 다시 매직 바이트 쓰기 
	    ch.write(MAGIC)
	}

	override fun close() {
	    ch.close()
	}
    }

단순히 arrow 파일로 쓰기 위한 것인듯. 매직바이트를 쓰는 것이 특이하다. 아마 읽을 때도 이 매직바이트를 읽을 것이다.

[xtdb.arrow.RelationReader.kt]​

주요구현을 보자.

• get(colName: String) : 컬럼 이름을 통한 데이터 접근과 반복자를 제공
• select 메서드를 통해 데이터 부분 집합을 쉽게 생성할 수 있음.
• from 메서드로 OldRelationReader 를 새로운 RelationReader 로 변환한다.

package xtdb.arrow

import org.apache.arrow.vector.types.pojp.Schema
import java.util.*
import xtdb.vector.RelationReader as OldRelationReader

interface RelationReader : Interable<VectorReader>, AutoCloseable {
  val schema: Schema
  val rowCount: Int

  operator fun get(colName: String): VectorReader?

  private class IndirectRelation(
    private val vecs: SequencedMap<String, VectorReader>,
    override val rowCount: Int
  ) : RelationReader {
      override val schema get() = Schema(vecs.map { it.value field })
      override fun get(colName: String) = vecs[colName]
      override fun iterator() = vecs.values.iterator()
      override fun close() = vecs.forEach { it.value.close() }
  }

  fun select(idxs: IntArray): RelationReader =
      IndirectRelation(associateTo(linkedMapOf()) { it.name to it.select(idxs) }, idxs.size)

  override fun close() = forEach { it.close() }

  companion object {
    @JvmStatic
    fun from(oldReader: OldRelationReader) = object : RelationReader {
	override val schema = Schema(oldReader.map { it.field })
	override val rowCount: Int get() = oldReader.rowCount()

	override operator fun get(colName: String) =
	    oldReader.readerForName(colName)?.let { VectorReader.from(it) }

	override fun iterator() = oldReader.asSequence().map { VectorReader.from(it) }.iterator()
    }
  }
}